Upload folder using huggingface_hub

README.md

@@ -1,44 +1,12 @@
+
 ---
-license: mit
-tags:
-- object-detection
-- computer-vision
-- yolov10
-datasets:
-- detection-datasets/coco
+title: rt-detr-object-detection 
+emoji: 🔥
+colorFrom: indigo
+colorTo: indigo
 sdk: gradio
-sdk_version: 4.42.0
+sdk_version: 5.0.0
+app_file: run.py
+pinned: false
+hf_oauth: true
 ---
-
-### Model Description
-[YOLOv10: Real-Time End-to-End Object Detection](https://arxiv.org/abs/2405.14458v1)
-
-- arXiv: https://arxiv.org/abs/2405.14458v1
-- github: https://github.com/THU-MIG/yolov10
-
-### Installation
-```
-pip install supervision git+https://github.com/THU-MIG/yolov10.git
-```
-
-### Yolov10 Inference
-```python
-from ultralytics import YOLOv10
-import supervision as sv
-import cv2
-
-IMAGE_PATH = 'dog.jpeg'
-
-model = YOLOv10.from_pretrained('jameslahm/yolov10{n/s/m/b/l/x}')
-model.predict(IMAGE_PATH, show=True)
-```
-
-### BibTeX Entry and Citation Info
- ```
-@article{wang2024yolov10,
-  title={YOLOv10: Real-Time End-to-End Object Detection},
-  author={Wang, Ao and Chen, Hui and Liu, Lihao and Chen, Kai and Lin, Zijia and Han, Jungong and Ding, Guiguang},
-  journal={arXiv preprint arXiv:2405.14458},
-  year={2024}
-}
-```

draw_boxes.py

@@ -1,7 +1,7 @@
-from PIL import Image, ImageDraw, ImageFont
-import numpy as np
+from PIL import ImageDraw, ImageFont  # type: ignore
 import colorsys
 
+
 def get_color(label):
     # Simple hash function to generate consistent colors for each label
     hash_value = hash(label)
@@ -11,31 +11,35 @@ def get_color(label):
     rgb = colorsys.hsv_to_rgb(hue, saturation, value)
     return tuple(int(x * 255) for x in rgb)
 
-def draw_bounding_boxes(image: Image, results: dict, model, threshold=0.3):
+
+def draw_bounding_boxes(image, results: dict, model, threshold=0.3):
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
 
-    for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
+    for score, label_id, box in zip(
+        results["scores"], results["labels"], results["boxes"]
+    ):
         if score > threshold:
             label = model.config.id2label[label_id.item()]
             box = [round(i, 2) for i in box.tolist()]
             color = get_color(label)
-            
+
             # Draw bounding box
-            draw.rectangle(box, outline=color, width=3)
-            
+            draw.rectangle(box, outline=color, width=3) # type: ignore
+
             # Prepare text
             text = f"{label}: {score:.2f}"
             text_bbox = draw.textbbox((0, 0), text, font=font)
             text_width = text_bbox[2] - text_bbox[0]
             text_height = text_bbox[3] - text_bbox[1]
-            
+
             # Draw text background
-            draw.rectangle([box[0], box[1] - text_height - 4, box[0] + text_width, box[1]], fill=color)
-            
+            draw.rectangle(
+                [box[0], box[1] - text_height - 4, box[0] + text_width, box[1]], # type: ignore
+                fill=color, # type: ignore
+            )
+
             # Draw text
             draw.text((box[0], box[1] - text_height - 4), text, fill="white", font=font)
-    
+
     return image
-  
-import numpy as np

requirements.txt

@@ -1,7 +1,5 @@
---extra-index-url https://download.pytorch.org/whl/cu113
 safetensors==0.4.3
 opencv-python
 torch
-transformers
-gradio-client @ git+https://github.com/gradio-app/gradio@66349fe26827e3a3c15b738a1177e95fec7f5554#subdirectory=client/python
-https://gradio-pypi-previews.s3.amazonaws.com/66349fe26827e3a3c15b738a1177e95fec7f5554/gradio-4.42.0-py3-none-any.whl
+transformers>=4.43.0
+Pillow

run.ipynb

@@ -0,0 +1 @@
+{"cells": [{"cell_type": "markdown", "id": "302934307671667531413257853548643485645", "metadata": {}, "source": ["# Gradio Demo: rt-detr-object-detection"]}, {"cell_type": "code", "execution_count": null, "id": "272996653310673477252411125948039410165", "metadata": {}, "outputs": [], "source": ["!pip install -q gradio safetensors==0.4.3 opencv-python torch transformers>=4.43.0 Pillow "]}, {"cell_type": "code", "execution_count": null, "id": "288918539441861185822528903084949547379", "metadata": {}, "outputs": [], "source": ["# Downloading files from the demo repo\n", "import os\n", "!wget -q https://github.com/gradio-app/gradio/raw/main/demo/rt-detr-object-detection/draw_boxes.py"]}, {"cell_type": "code", "execution_count": null, "id": "44380577570523278879349135829904343037", "metadata": {}, "outputs": [], "source": ["import spaces\n", "import gradio as gr\n", "import cv2\n", "from PIL import Image\n", "import torch\n", "import time\n", "import numpy as np\n", "import uuid\n", "\n", "from transformers import RTDetrForObjectDetection, RTDetrImageProcessor  # type: ignore\n", "\n", "from draw_boxes import draw_bounding_boxes\n", "\n", "image_processor = RTDetrImageProcessor.from_pretrained(\"PekingU/rtdetr_r50vd\")\n", "model = RTDetrForObjectDetection.from_pretrained(\"PekingU/rtdetr_r50vd\").to(\"cuda\")\n", "\n", "\n", "SUBSAMPLE = 2\n", "\n", "\n", "@spaces.GPU\n", "def stream_object_detection(video, conf_threshold):\n", "    cap = cv2.VideoCapture(video)\n", "\n", "    video_codec = cv2.VideoWriter_fourcc(*\"mp4v\")  # type: ignore\n", "    fps = int(cap.get(cv2.CAP_PROP_FPS))\n", "\n", "    desired_fps = fps // SUBSAMPLE\n", "    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2\n", "    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2\n", "\n", "    iterating, frame = cap.read()\n", "\n", "    n_frames = 0\n", "\n", "    name = f\"output_{uuid.uuid4()}.mp4\"\n", "    segment_file = cv2.VideoWriter(name, video_codec, desired_fps, (width, height))  # type: ignore\n", "    batch = []\n", "\n", "    while iterating:\n", "        frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)\n", "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n", "        if n_frames % SUBSAMPLE == 0:\n", "            batch.append(frame)\n", "        if len(batch) == 2 * desired_fps:\n", "            inputs = image_processor(images=batch, return_tensors=\"pt\").to(\"cuda\")\n", "\n", "            print(f\"starting batch of size {len(batch)}\")\n", "            start = time.time()\n", "            with torch.no_grad():\n", "                outputs = model(**inputs)\n", "            end = time.time()\n", "            print(\"time taken for inference\", end - start)\n", "\n", "            start = time.time()\n", "            boxes = image_processor.post_process_object_detection(\n", "                outputs,\n", "                target_sizes=torch.tensor([(height, width)] * len(batch)),\n", "                threshold=conf_threshold,\n", "            )\n", "\n", "            for _, (array, box) in enumerate(zip(batch, boxes)):\n", "                pil_image = draw_bounding_boxes(\n", "                    Image.fromarray(array), box, model, conf_threshold\n", "                )\n", "                frame = np.array(pil_image)\n", "                # Convert RGB to BGR\n", "                frame = frame[:, :, ::-1].copy()\n", "                segment_file.write(frame)\n", "\n", "            batch = []\n", "            segment_file.release()\n", "            yield name\n", "            end = time.time()\n", "            print(\"time taken for processing boxes\", end - start)\n", "            name = f\"output_{uuid.uuid4()}.mp4\"\n", "            segment_file = cv2.VideoWriter(\n", "                name, video_codec, desired_fps, (width, height)\n", "            )  # type: ignore\n", "\n", "        iterating, frame = cap.read()\n", "        n_frames += 1\n", "\n", "\n", "with gr.Blocks() as demo:\n", "    gr.HTML(\n", "        \"\"\"\n", "    <h1 style='text-align: center'>\n", "    Video Object Detection with <a href='https://huggingface.co/PekingU/rtdetr_r101vd_coco_o365' target='_blank'>RT-DETR</a>\n", "    </h1>\n", "    \"\"\"\n", "    )\n", "    with gr.Row():\n", "        with gr.Column():\n", "            video = gr.Video(label=\"Video Source\")\n", "            conf_threshold = gr.Slider(\n", "                label=\"Confidence Threshold\",\n", "                minimum=0.0,\n", "                maximum=1.0,\n", "                step=0.05,\n", "                value=0.30,\n", "            )\n", "        with gr.Column():\n", "            output_video = gr.Video(\n", "                label=\"Processed Video\", streaming=True, autoplay=True\n", "            )\n", "\n", "    video.upload(\n", "        fn=stream_object_detection,\n", "        inputs=[video, conf_threshold],\n", "        outputs=[output_video],\n", "    )\n", "\n", "if __name__ == \"__main__\":\n", "    demo.launch()\n"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}

run.py

@@ -0,0 +1,115 @@
+import spaces
+import gradio as gr
+import cv2
+from PIL import Image
+import torch
+import time
+import numpy as np
+import uuid
+
+from transformers import RTDetrForObjectDetection, RTDetrImageProcessor  # type: ignore
+
+from draw_boxes import draw_bounding_boxes
+
+image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd")
+model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd").to("cuda")
+
+
+SUBSAMPLE = 2
+
+
+@spaces.GPU
+def stream_object_detection(video, conf_threshold):
+    cap = cv2.VideoCapture(video)
+
+    video_codec = cv2.VideoWriter_fourcc(*"mp4v")  # type: ignore
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+
+    desired_fps = fps // SUBSAMPLE
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2
+
+    iterating, frame = cap.read()
+
+    n_frames = 0
+
+    name = f"output_{uuid.uuid4()}.mp4"
+    segment_file = cv2.VideoWriter(name, video_codec, desired_fps, (width, height))  # type: ignore
+    batch = []
+
+    while iterating:
+        frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        if n_frames % SUBSAMPLE == 0:
+            batch.append(frame)
+        if len(batch) == 2 * desired_fps:
+            inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
+
+            print(f"starting batch of size {len(batch)}")
+            start = time.time()
+            with torch.no_grad():
+                outputs = model(**inputs)
+            end = time.time()
+            print("time taken for inference", end - start)
+
+            start = time.time()
+            boxes = image_processor.post_process_object_detection(
+                outputs,
+                target_sizes=torch.tensor([(height, width)] * len(batch)),
+                threshold=conf_threshold,
+            )
+
+            for _, (array, box) in enumerate(zip(batch, boxes)):
+                pil_image = draw_bounding_boxes(
+                    Image.fromarray(array), box, model, conf_threshold
+                )
+                frame = np.array(pil_image)
+                # Convert RGB to BGR
+                frame = frame[:, :, ::-1].copy()
+                segment_file.write(frame)
+
+            batch = []
+            segment_file.release()
+            yield name
+            end = time.time()
+            print("time taken for processing boxes", end - start)
+            name = f"output_{uuid.uuid4()}.mp4"
+            segment_file = cv2.VideoWriter(
+                name, video_codec, desired_fps, (width, height)
+            )  # type: ignore
+
+        iterating, frame = cap.read()
+        n_frames += 1
+
+
+with gr.Blocks() as demo:
+    gr.HTML(
+        """
+    <h1 style='text-align: center'>
+    Video Object Detection with <a href='https://huggingface.co/PekingU/rtdetr_r101vd_coco_o365' target='_blank'>RT-DETR</a>
+    </h1>
+    """
+    )
+    with gr.Row():
+        with gr.Column():
+            video = gr.Video(label="Video Source")
+            conf_threshold = gr.Slider(
+                label="Confidence Threshold",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.05,
+                value=0.30,
+            )
+        with gr.Column():
+            output_video = gr.Video(
+                label="Processed Video", streaming=True, autoplay=True
+            )
+
+    video.upload(
+        fn=stream_object_detection,
+        inputs=[video, conf_threshold],
+        outputs=[output_video],
+    )
+
+if __name__ == "__main__":
+    demo.launch()