File size: 4,428 Bytes
feb730b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 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) |