Revise to support video processing with Supervision

This update entirely overhauls the application to replace the existing image processing functionality with a custom video processing implementation.

README.md

@@ -4,7 +4,7 @@ emoji: 🎨
 colorFrom: pink
 colorTo: purple
 sdk: gradio
-sdk_version: 4.0.2
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,33 +1,53 @@
-from typing import Dict
+import time
+import uuid
+from typing import Tuple
 
 import gradio as gr
-import torch
-from PIL import Image
-from transformers import SamModel, SamProcessor
+import supervision as sv
+from tqdm import tqdm
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL = SamModel.from_pretrained("facebook/sam-vit-large").to(DEVICE)
-PROCESSOR = SamProcessor.from_pretrained("facebook/sam-vit-large")
+START_FRAME = 0
+END_FRAME = 10
+TOTAL = END_FRAME - START_FRAME
 
 
-def inference(masked_image: Dict[str, Image.Image]) -> Image.Image:
-    image = masked_image['image']
-    mask = masked_image['mask'].resize((256, 256), Image.Resampling.LANCZOS)
-    return image
+def process(
+    source_video: str,
+    prompt: str,
+    confidence: float,
+    progress=gr.Progress(track_tqdm=True)
+) -> Tuple[str, str]:
+    name = str(uuid.uuid4())
+    video_info = sv.VideoInfo.from_video_path(source_video)
+    frame_iterator = iter(sv.get_video_frames_generator(
+        source_path=source_video, start=START_FRAME, end=END_FRAME))
+
+    with sv.VideoSink(f"{name}.mp4", video_info=video_info) as sink:
+        for _ in tqdm(range(TOTAL), desc="Masking frames"):
+            frame = next(frame_iterator)
+            sink.write_frame(frame)
+            time.sleep(0.1)
+
+    return f"{name}.mp4", f"{name}.mp4"
 
 
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(
-                image_mode='RGB', type='pil', tool="sketch", interactive=True,
-                brush_radius=20.0, brush_color="#FFFFFF", height=500)
+            source_video_player = gr.Video(
+                label="Source video", source="upload", format="mp4")
+            prompt_text = gr.Textbox(
+                label="Prompt", value="person")
+            confidence_slider = gr.Slider(
+                label="Confidence", minimum=0.5, maximum=1.0, step=0.05, value=0.6)
             submit_button = gr.Button("Submit")
-        output_image = gr.Image(image_mode='RGB', type='pil')
+        with gr.Column():
+            masked_video_player = gr.Video(label="Masked video")
+            painted_video_player = gr.Video(label="Painted video")
 
     submit_button.click(
-        inference,
-        inputs=[input_image],
-        outputs=output_image)
+        process,
+        inputs=[source_video_player, prompt_text, confidence_slider],
+        outputs=[masked_video_player, painted_video_player])
 
-demo.launch(debug=False, show_error=True)
+demo.queue().launch(debug=False, show_error=True)

requirements.txt

@@ -3,6 +3,7 @@ torch
 torchvision
 
 numpy
+opencv-python
 pillow
 gradio==3.50.2
 transformers