Upload gradio_demo.py with huggingface_hub

gradio_demo.py

@@ -0,0 +1,138 @@
+import os
+import cv2
+import time
+import torch
+import random
+import gradio as gr
+import numpy as np
+from loguru import logger
+from transformers import VJEPA2ForVideoClassification, AutoVideoProcessor
+
+# Config
+CHECKPOINT = "HaithemH/vjepa2-vitl-fpc16-256-ssv2-66K-220cat" 
+TORCH_DTYPE = torch.float16
+TORCH_DEVICE = "cuda:4"  # Change if needed
+UPDATE_EVERY_N_FRAMES = 16
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Load model & processor
+model = VJEPA2ForVideoClassification.from_pretrained(CHECKPOINT, torch_dtype=torch.bfloat16)
+model = model.to(TORCH_DEVICE)
+video_processor = AutoVideoProcessor.from_pretrained(CHECKPOINT)
+frames_per_clip = model.config.frames_per_clip
+
+
+def add_text_on_image(image, text):
+    image[:70] = 0
+    line_spacing = 10
+    top_margin = 20
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = 0.5
+    thickness = 1
+    color = (255, 255, 255)
+    words = text.split()
+    lines = []
+    current_line = ""
+    img_width = image.shape[1]
+    for word in words:
+        test_line = current_line + (" " if current_line else "") + word
+        (test_width, _), _ = cv2.getTextSize(test_line, font, font_scale, thickness)
+        if test_width > img_width - 20:
+            lines.append(current_line)
+            current_line = word
+        else:
+            current_line = test_line
+    if current_line:
+        lines.append(current_line)
+    y = top_margin
+    for line in lines:
+        (line_width, line_height), _ = cv2.getTextSize(line, font, font_scale, thickness)
+        x = (img_width - line_width) // 2
+        cv2.putText(image, line, (x, y + line_height), font, font_scale, color, thickness, cv2.LINE_AA)
+        y += line_height + line_spacing
+    return image
+
+
+class RunningFramesCache:
+    def __init__(self, max_frames=16):
+        self.max_frames = max_frames
+        self._frames = []
+        self.counter = 0
+
+    def add_frame(self, frame):
+        self.counter += 1
+        self._frames.append(frame)
+        if len(self._frames) > self.max_frames:
+            self._frames.pop(0)
+
+    def get_last_n_frames(self, n):
+        return self._frames[-n:]
+
+    def __len__(self):
+        return len(self._frames)
+
+
+class RunningResult:
+    def __init__(self, max_predictions=4):
+        self.predictions = []
+        self.max_predictions = max_predictions
+
+    def add_prediction(self, prediction):
+        current_time = time.strftime("%H:%M:%S", time.gmtime(time.time()))
+        self.predictions.append((current_time, prediction))
+        if len(self.predictions) > self.max_predictions:
+            self.predictions.pop(0)
+
+    def get_formatted(self):
+        if not self.predictions:
+            return "Starting..."
+        current, *past = self.predictions[::-1]
+        text = f">>> {current[1]}\n\n" + "\n".join(
+            f"[{time_str}] {pred}" for time_str, pred in past
+        )
+        return text
+
+    def get_last(self):
+        return self.predictions[-1][1] if self.predictions else "Starting..."
+
+
+# Shared state
+frames_cache = RunningFramesCache(max_frames=frames_per_clip)
+results_cache = RunningResult()
+
+
+def classify_frame(image):
+    image = cv2.flip(image, 1)  # mirror webcam
+    frames_cache.add_frame(image)
+
+    if frames_cache.counter % UPDATE_EVERY_N_FRAMES == 0 and len(frames_cache) >= frames_per_clip:
+        frames = frames_cache.get_last_n_frames(frames_per_clip)
+        frames = np.array(frames)
+        inputs = video_processor(frames, device=TORCH_DEVICE, return_tensors="pt")
+        inputs = inputs.to(dtype=TORCH_DTYPE)
+        with torch.no_grad():
+            logits = model(**inputs).logits
+        top_idx = logits.argmax(dim=-1).item()
+        class_name = model.config.id2label[top_idx]
+        logger.info(f"Predicted: {class_name}")
+        results_cache.add_prediction(class_name)
+
+    annotated_image = add_text_on_image(image.copy(), results_cache.get_last())
+    return annotated_image, results_cache.get_formatted()
+
+
+# Gradio UI
+demo = gr.Interface(
+    fn=classify_frame,
+    inputs=gr.Image(sources=["webcam"], streaming=True),
+    outputs=[
+        gr.Image(label="Live Prediction", type="numpy"),
+        gr.TextArea(label="Recent Predictions", lines=10),
+    ],
+    live=True,
+    title="V-JEPA 2: Streaming Video Action Recognition - SSV2",
+    description="This demo showcases a specialized version of V-JEPA 2, fine-tuned for real-time video action recognition!",
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)