code

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+video_example.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__/

README.md

@@ -5,8 +5,8 @@ emoji: 🔥
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio
-sdk_version: 5.0.2
-app_file: run.py
+sdk_version: 5.1.0
+app_file: app.py
 pinned: false
 hf_oauth: true
 ---

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 import torch
 import time
 import numpy as np
-import uuid
+from gradio_webrtc import WebRTC
 
 from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
 
@@ -21,21 +21,6 @@ SUBSAMPLE = 2
 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)
@@ -62,15 +47,11 @@ def stream_object_detection(video, conf_threshold):
                 frame = np.array(pil_image)
                 # Convert RGB to BGR
                 frame = frame[:, :, ::-1].copy()
-                segment_file.write(frame)
+                yield 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
@@ -80,7 +61,7 @@ with gr.Blocks() as app:
     gr.HTML(
         """
     <h1 style='text-align: center'>
-    Video Object Detection with RT-DETR
+    Video Object Detection with RT-DETR (Powered by WebRTC ⚡️)
     </h1>
     """)
     gr.HTML(
@@ -100,13 +81,21 @@ with gr.Blocks() as app:
                 value=0.30,
             )
         with gr.Column():
-            output_video = gr.Video(label="Processed Video", streaming=True, autoplay=True)
+            output = WebRTC(label="WebRTC Stream",
+                                   rtc_configuration=None,
+                                   mode="receive",
+                                   modality="video")
+            detect = gr.Button("Detect", variant="primary")
 
-    video.upload(
+    output.stream(
         fn=stream_object_detection,
         inputs=[video, conf_threshold],
-        outputs=[output_video],
+        outputs=[output],
+        trigger=detect.click
     )
 
+    gr.Examples(examples=["video_example.mp4"],
+                inputs=[video])
+
 if __name__ == '__main__':
     app.launch()

requirements.txt

@@ -3,3 +3,4 @@ opencv-python
 torch
 transformers>=4.43.0
 Pillow
+gradio-webrtc

run.ipynb

@@ -1 +0,0 @@
-{"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

@@ -1,115 +0,0 @@
-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()

video_example.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58f4485e53ab6877244b23c73c74151475a7f3814b8832e32ca02562e37ea0a5
+size 5373394