Add video detection

app.py

@@ -1,6 +1,8 @@
+import os.path as osp
 from typing import List
 
 import gradio as gr
+import numpy as np
 import supervision as sv
 import torch
 from PIL import Image
@@ -60,7 +62,7 @@ def annotate(
     return Image.fromarray(annotated_bgr_image[:, :, ::-1])
 
 
-def inference(image_rgb_pil: Image.Image, confidence: float) -> List[Image.Image]:
+def inference_image(image_rgb_pil: Image.Image, confidence: float) -> List[Image.Image]:
     output = YOLO_MODEL(image_rgb_pil, imgsz=640, verbose=False)[0]
     detections = sv.Detections.from_ultralytics(output)
 
@@ -79,30 +81,86 @@ def inference(image_rgb_pil: Image.Image, confidence: float) -> List[Image.Image
     )
 
 
-custom_theme = gr.themes.Soft(primary_hue="blue").set(
-    button_secondary_background_fill="*neutral_100",
-    button_secondary_background_fill_hover="*neutral_200",
-)
+def process_frame(frame: np.ndarray, _) -> np.ndarray:
+    output = YOLO_MODEL(frame, imgsz=640, verbose=False)[0]
 
+    detections = sv.Detections.from_ultralytics(output)
+
+    labels = [
+        f"{output.names[class_id]} {confidence:0.2f}"
+        for _, _, confidence, class_id, _ in detections
+    ]
+
+    thickness = 2
+    text_thickness = 1
+    text_scale = 1.0
+
+    height, width, _ = output.orig_img.shape
+
+    thickness_ratio = ((width + height) / 2) / 400
+    text_scale_ratio = ((width + height) / 2) / 600
+    text_thickness_ratio = ((width + height) / 2) / 400
+
+    BOX_ANNOTATOR.thickness = int(thickness * thickness_ratio)
+    BOX_ANNOTATOR.text_scale = float(text_scale * text_scale_ratio)
+    BOX_ANNOTATOR.text_thickness = int(text_thickness * text_thickness_ratio)
+
+    annotated_frame = BOX_ANNOTATOR.annotate(
+        scene=output.orig_img.copy(), detections=detections, labels=labels
+    )
+    return annotated_frame
+
+
+def inference_video(path_video):
+    path_output_video = "temp.mp4"
+    sv.process_video(
+        source_path=path_video,
+        target_path=path_output_video,
+        callback=process_frame,
+    )
+    return path_output_video
+
+
+custom_theme = gr.themes.Soft(primary_hue="green")
 with gr.Blocks(theme=custom_theme, css="style.css") as demo:
     gr.Markdown(MARKDOWN)
 
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(image_mode="RGB", type="pil", height=500)
-            confidence_slider = gr.Slider(
-                label="Confidence", minimum=0.1, maximum=1.0, step=0.05, value=0.6
-            )
-            submit_button = gr.Button("Let's find orang outans 🦧 !")
-        output_image = gr.Image(label="Results", type="pil")
-
-    submit_button.click(
-        inference,
-        inputs=[input_image, confidence_slider],
+    with gr.Tab("Detect on an image 🖼️"):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    image_mode="RGB",
+                    sources=["upload", "clipboard"],
+                    type="pil",
+                )
+                confidence_image_slider = gr.Slider(
+                    label="Confidence", minimum=0.1, maximum=1.0, step=0.05, value=0.6
+                )
+                submit_button_image = gr.Button("Let's find orang outans 🦧 !")
+            output_image = gr.Image(label="Results", type="pil")
+    with gr.Tab("Detect on a video 📹"):
+        with gr.Row():
+            with gr.Column():
+                input_video = gr.Video(sources=["upload"])
+                # confidence_video_slider = gr.Slider(
+                #     label="Confidence", minimum=0.1, maximum=1.0, step=0.05, value=0.6
+                # )
+                submit_button_video = gr.Button("Let's find orang outans 🦧 !")
+            output_video = gr.Video(label="Results")
+
+    submit_button_image.click(
+        inference_image,
+        inputs=[input_image, confidence_image_slider],
         outputs=output_image,
         queue=True,
     )
 
+    submit_button_video.click(
+        inference_video,
+        inputs=[input_video],
+        outputs=output_video,
+        queue=True,
+    )
 
 if __name__ == "__main__":
     demo.queue(max_size=20, api_open=False).launch()