{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9904ec5-391d-4967-9357-c8779d677142",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import required libraries\n",
    "from ultralytics import YOLO\n",
    "import gradio as gr\n",
    "import cv2\n",
    "import math\n",
    "from items import classNames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dbb6ae7-c844-4933-9a5c-f778bb1dfa83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# detection function\n",
    "def yolo_detect(feed, vid):\n",
    "    video = vid\n",
    "    # Load a pretrained YOLOv8n model\n",
    "    model = YOLO('yolov8n.pt')\n",
    "    \n",
    "    # Run inference on the source\n",
    "    results = model(video, stream=True, verbose=False) \n",
    "    frames = list()\n",
    "    \n",
    "    # plot annotations\n",
    "    for frame in results:\n",
    "        boxes = frame.boxes\n",
    "        single = frame.orig_img\n",
    "        for box in boxes:\n",
    "            # bounding box\n",
    "            x1, y1, x2, y2 = box.xyxy[0]\n",
    "            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # convert to int values\n",
    "\n",
    "            # put box in cam\n",
    "            cv2.rectangle(single, (x1, y1), (x2, y2), (255, 0, 255), 3)\n",
    "\n",
    "            # object details\n",
    "            cv2.putText(single, classNames[int(box.cls[0])], (x1,y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1)\n",
    "            \n",
    "        frames.append(single)\n",
    "    cv2.destroyAllWindows()\n",
    "    \n",
    "    h, w, c = frames[1].shape\n",
    "    \n",
    "    out_file = \"output.avi\"\n",
    "    fourcc=cv2.VideoWriter_fourcc('X', 'V', 'I', 'D')\n",
    "    writer = out = cv2.VideoWriter(out_file, fourcc, 25.0, (w, h))\n",
    "    for i in range(len(frames)):\n",
    "        writer.write(frames[i])\n",
    "    writer.release()\n",
    "    return out_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692f5c49-67cd-4c11-8ee9-03dc7cb98809",
   "metadata": {},
   "outputs": [],
   "source": [
    "demo = gr.Interface(fn=yolo_detect, \n",
    "                    inputs=[gr.PlayableVideo(source='webcam'), gr.Video(autoplay=True)],\n",
    "                    outputs=[gr.PlayableVideo(autoplay=True, format='avi')],\n",
    "                    cache_examples=True, allow_flagging='never')\n",
    "demo.queue()\n",
    "demo.launch(inline=False, debug=True, show_api=False, quiet=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}