Add video input

README.md

@@ -61,13 +61,3 @@ The following variants are available out of the box:
 - **app.py**: Main Gradio application script.
 - **requirements.txt**: Python dependencies, including MMCV and MMPose.
 - **README.md**: This documentation file.
-
-## Development
-
-To update dependencies, edit `requirements.txt`. To extend functionality or add new variants, modify `app.py` accordingly.
-
-## Future Plans
-
-1. Support video input streams.
-2. Enable ONNX model inference via `rtmlib`.
-

app.py

@@ -1,10 +1,14 @@
 #!/usr/bin/env python3
 import spaces
-import os, sys, importlib.util, re
+import os
+import sys
+import importlib.util
+import re
 import gradio as gr
 from PIL import Image
 import torch
 import requests  # for downloading remote checkpoints
+import shutil
 
 # CUDA info
 try:
@@ -92,22 +96,45 @@ def load_inferencer(checkpoint_path=None, device=None):
 # —─── Prediction function ────
 @spaces.GPU()
 def predict(image: Image.Image,
+            video,                         # new video input
             remote_ckpt: str,
             upload_ckpt,
             bbox_thr: float,
             nms_thr: float):
-    inp_path = "/tmp/upload.jpg"
-    image.save(inp_path)
+    # 1) Write image or pick up video file
+    if video:
+        # Gradio Video can come in as a filepath string or dict
+        if isinstance(video, dict) and 'name' in video:
+            inp_path = video['name']
+        elif hasattr(video, "name"):
+            inp_path = video.name
+        else:
+            inp_path = video
+    else:
+        inp_path = "/tmp/upload.jpg"
+        image.save(inp_path)
+
+    # 2) Determine by extension if this is video
+    ext = os.path.splitext(inp_path)[1].lower()
+    is_video = ext in (".mp4", ".mov", ".avi", ".mkv", ".webm")
+
+    # checkpoint selection
     if upload_ckpt:
         ckpt_path = upload_ckpt.name
         active = os.path.basename(ckpt_path)
     else:
         ckpt_path = get_checkpoint(remote_ckpt)
         active = remote_ckpt
+
+    # prepare (and clear) output dir
     vis_dir = "/tmp/vis"
+    if os.path.exists(vis_dir):
+        shutil.rmtree(vis_dir)
     os.makedirs(vis_dir, exist_ok=True)
+
+    # run inferencer (handles both image & video)
     inferencer = load_inferencer(checkpoint_path=ckpt_path, device=None)
-    for result in inferencer(
+    for _ in inferencer(
         inputs=inp_path,
         bbox_thr=bbox_thr,
         nms_thr=nms_thr,
@@ -116,9 +143,18 @@ def predict(image: Image.Image,
         vis_out_dir=vis_dir,
     ):
         pass
+
+    # collect and return results
     out_files = sorted(os.listdir(vis_dir))
-    vis_img = Image.open(os.path.join(vis_dir, out_files[0])) if out_files else None
-    return vis_img, active
+    if is_video:
+        # return only the annotated video path
+        out_vid = next((f for f in out_files if f.lower().endswith((".mp4", ".mov", ".avi"))), None)
+        return None, os.path.join(vis_dir, out_vid) if out_vid else None, active
+    else:
+        # return only the annotated image
+        img_f = out_files[0] if out_files else None
+        vis_img = Image.open(os.path.join(vis_dir, img_f)) if img_f and not img_f.lower().endswith((".mp4", ".mov", ".avi")) else None
+        return vis_img, None, active
 
 # —─── Gradio UI ────
 def main():
@@ -126,43 +162,44 @@ def main():
         gr.Markdown("## RTMO Pose Demo")
         with gr.Row():
             with gr.Column(scale=1, min_width=300):
-                img_input = gr.Image(type="pil", label="Upload Image")
-                remote_dd = gr.Dropdown(label="Select Remote Checkpoint",
-                                      choices=list(REMOTE_CHECKPOINTS.keys()),
-                                      value=list(REMOTE_CHECKPOINTS.keys())[0])
+                img_input   = gr.Image(type="pil", label="Upload Image")
+                video_input = gr.Video(label="Upload Video")
+                remote_dd   = gr.Dropdown(
+                    label="Select Remote Checkpoint",
+                    choices=list(REMOTE_CHECKPOINTS.keys()),
+                    value=list(REMOTE_CHECKPOINTS.keys())[0]
+                )
                 upload_ckpt = gr.File(file_types=['.pth'], label="Or Upload Your Own Checkpoint (optional)")
-                bbox_thr = gr.Slider(minimum=0.0, maximum=1.0, step=0.01,
-                                     value=0.1, label="Bounding Box Threshold")
-                nms_thr = gr.Slider(minimum=0.0, maximum=1.0, step=0.01,
-                                    value=0.65, label="NMS Threshold")
-                run_btn = gr.Button("Run Inference")
+                bbox_thr    = gr.Slider(0.0, 1.0, value=0.1, step=0.01, label="Bounding Box Threshold")
+                nms_thr     = gr.Slider(0.0, 1.0, value=0.65, step=0.01, label="NMS Threshold")
+                run_btn     = gr.Button("Run Inference")
             with gr.Column(scale=2):
-                output_img = gr.Image(type="pil", label="Annotated Image",
-                                      elem_id="output_image", interactive=False)
-                active_tb = gr.Textbox(label="Active Checkpoint", interactive=False)
-        
+                output_img   = gr.Image(type="pil", label="Annotated Image", elem_id="output_image", interactive=False)
+                output_video = gr.Video(label="Annotated Video", interactive=False)
+                active_tb    = gr.Textbox(label="Active Checkpoint", interactive=False)
+
         # Examples for quick testing
         gr.Examples(
             examples=[
-                ["https://images.pexels.com/photos/1858175/pexels-photo-1858175.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=614",
-                 "rtmo-s_coco_retrainable", None, 0.1, 0.65],
-                ["https://images.pexels.com/photos/3779706/pexels-photo-3779706.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=614",
-                 "rtmo-t_8xb32-600e_body7", None, 0.1, 0.65],
-                ["https://images.pexels.com/photos/220453/pexels-photo-220453.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=614",
-                 "rtmo-s_8xb32-600e_coco", None, 0.1, 0.65],
+                ["https://images.pexels.com/photos/1858175/pexels-photo-1858175.jpeg?auto=compress&cs=tinysrgb&h=614&w=614", None, "rtmo-s_coco_retrainable", None, 0.1, 0.65],
+                ["https://images.pexels.com/photos/3779706/pexels-photo-3779706.jpeg?auto=compress&cs=tinysrgb&h=614&w=614", None, "rtmo-t_8xb32-600e_body7", None, 0.1, 0.65],
+                ["https://images.pexels.com/photos/220453/pexels-photo-220453.jpeg?auto=compress&cs=tinysrgb&h=614&w=614", None, "rtmo-s_8xb32-600e_coco", None, 0.1, 0.65],
             ],
-            inputs=[img_input, remote_dd, upload_ckpt, bbox_thr, nms_thr],
-            outputs=[output_img, active_tb],
+            inputs=[img_input, video_input, remote_dd, upload_ckpt, bbox_thr, nms_thr],
+            outputs=[output_img, output_video, active_tb],
             fn=predict,
             cache_examples=False,
             label="Examples",
             examples_per_page=3
         )
 
-        run_btn.click(predict,
-                      inputs=[img_input, remote_dd, upload_ckpt, bbox_thr, nms_thr],
-                      outputs=[output_img, active_tb])
+        run_btn.click(
+            predict,
+            inputs=[img_input, video_input, remote_dd, upload_ckpt, bbox_thr, nms_thr],
+            outputs=[output_img, output_video, active_tb]
+        )
+
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()