Spaces:
Running on Zero
Running on Zero
| """Video Intelligence: detect, segment, track and caption images and videos. | |
| UI and Gradio handlers only. Detectors live in `detection.py`, captioners in | |
| `captioning.py`, and model folders under `models/` are managed by `hub.py`. | |
| Zero-shot labels come from a JSON vocabulary (`labels/default.json`) or a typed list. | |
| """ | |
| import shutil | |
| import subprocess | |
| import tempfile | |
| from pathlib import Path | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import supervision as sv | |
| import torch | |
| from PIL import Image, ImageDraw, ImageFont | |
| from trackers import ByteTrackTracker | |
| import hub | |
| from captioning import ( | |
| CAPTION_TASKS, | |
| DEFAULT_CAPTION_DETAIL, | |
| DEFAULT_VIDEO_CAPTIONERS, | |
| DEFAULT_VIDEO_PROMPT, | |
| FLORENCE, | |
| MARLIN_PROMPT, | |
| VIDEO_CAPTIONERS, | |
| caption_image, | |
| caption_video, | |
| count_frames, | |
| format_marlin, | |
| load_image_captioner, | |
| load_video_captioner, | |
| needs_prompt, | |
| ) | |
| from detection import DETECTOR_SPECS, TASKS, ZERO_SHOT_TASK, Vocabulary, detect_at_resolution, load_detector | |
| # ----------------------------- | |
| # Config | |
| # ----------------------------- | |
| DEFAULT_TASK = "Detection (COCO)" | |
| DEFAULT_MODEL = {"Detection (COCO)": "RF-DETR Medium", "Segmentation (COCO)": "RF-DETR Seg Medium", ZERO_SHOT_TASK: "OWLv2 Base"} | |
| DEFAULT_VIDEO_MODEL = {"Detection (COCO)": "RF-DETR Small", "Segmentation (COCO)": "RF-DETR Seg Small", ZERO_SHOT_TASK: "OWLv2 Base"} | |
| LABELS_DIR = Path(__file__).resolve().parent / "labels" | |
| LABEL_PRESETS = {p.stem: str(p) for p in sorted(LABELS_DIR.glob("*.json"))} | |
| DEFAULT_LABEL_PRESET = "default" | |
| # ByteTrack's second-stage association needs low-confidence boxes to keep | |
| # tracks alive through momentary dips, so for video we run the model well | |
| # below the user's threshold and let the tracker gate what gets shown. | |
| VIDEO_LOW_CONF = 0.05 | |
| FALLBACK_FPS = 25 | |
| MAX_DETECT_FPS = 60 | |
| SOURCE_RESOLUTION = 0 # longest side fed to the detector; 0 keeps the source size | |
| def model_spec(task: str, model_name: str) -> dict: | |
| return TASKS[task][model_name] | |
| def vocabulary_for(task: str, labels_text: str, labels_preset: str, labels_upload: str | None) -> Vocabulary | None: | |
| """Typed labels win, then an uploaded JSON, then the preset; closed-set tasks need none.""" | |
| if task != ZERO_SHOT_TASK: | |
| return None | |
| if labels_text and labels_text.strip(): | |
| return Vocabulary.from_names(labels_text) | |
| return Vocabulary.from_json(labels_upload or LABEL_PRESETS[labels_preset]) | |
| def video_fps(path: str) -> float: | |
| info = sv.VideoInfo.from_video_path(path) | |
| return info.fps if info.fps and info.fps > 0 else FALLBACK_FPS | |
| def cap_detections(det: sv.Detections, max_det: int) -> sv.Detections: | |
| if len(det) > max_det and det.confidence is not None: | |
| det = det[np.argsort(-det.confidence)[:max_det]] | |
| return det | |
| # ----------------------------- | |
| # Drawing | |
| # ----------------------------- | |
| def build_annotators(resolution_wh, tracking: bool): | |
| thickness = sv.calculate_optimal_line_thickness(resolution_wh) | |
| text_scale = sv.calculate_optimal_text_scale(resolution_wh) | |
| lookup = sv.ColorLookup.TRACK if tracking else sv.ColorLookup.CLASS | |
| return [ | |
| sv.MaskAnnotator(opacity=0.45, color_lookup=lookup), # no-op without masks | |
| sv.BoxAnnotator(thickness=thickness, color_lookup=lookup), | |
| sv.LabelAnnotator( | |
| text_scale=text_scale, | |
| text_thickness=max(1, thickness - 1), | |
| text_padding=max(4, thickness * 2), | |
| text_position=sv.Position.TOP_LEFT, | |
| color_lookup=lookup, | |
| smart_position=True, | |
| ), | |
| ] | |
| def make_labels(det: sv.Detections) -> list[str]: | |
| names = det.data.get("class_name") | |
| if names is None: | |
| names = [str(c) for c in det.class_id] | |
| confs = det.confidence if det.confidence is not None else [None] * len(det) | |
| labels = [] | |
| for i, (name, c) in enumerate(zip(names, confs)): | |
| text = name if c is None else f"{name} {c:.0%}" | |
| if det.tracker_id is not None: | |
| text = f"#{det.tracker_id[i]} {text}" | |
| labels.append(text) | |
| return labels | |
| def annotate(scene_bgr: np.ndarray, det: sv.Detections, annotators) -> np.ndarray: | |
| scene = scene_bgr.copy() | |
| for ann in annotators: | |
| if isinstance(ann, sv.LabelAnnotator): | |
| scene = ann.annotate(scene, det, labels=make_labels(det)) | |
| else: | |
| scene = ann.annotate(scene, det) | |
| return scene | |
| def placeholder(msg: str, size=(640, 480)) -> Image.Image: | |
| w, h = size | |
| img = Image.new("RGB", (w, h), "white") | |
| draw = ImageDraw.Draw(img) | |
| font = ImageFont.load_default() | |
| bbox = draw.textbbox((0, 0), msg, font=font) | |
| draw.text(((w - (bbox[2] - bbox[0])) / 2, (h - (bbox[3] - bbox[1])) / 2), msg, fill="black", font=font) | |
| return img | |
| def to_h264(src: str) -> str: | |
| """Re-encode to browser-friendly H.264 if ffmpeg is available.""" | |
| ffmpeg = shutil.which("ffmpeg") | |
| if ffmpeg is None: | |
| try: | |
| import imageio_ffmpeg | |
| ffmpeg = imageio_ffmpeg.get_ffmpeg_exe() | |
| except Exception: | |
| return src | |
| dst = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name | |
| cmd = [ | |
| ffmpeg, "-y", "-loglevel", "error", "-i", src, | |
| "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", | |
| "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart", dst, | |
| ] | |
| ok = subprocess.run(cmd).returncode == 0 | |
| return dst if ok else src | |
| # ----------------------------- | |
| # Gradio handlers | |
| # ----------------------------- | |
| def on_task_change(task): | |
| """Swap the model list and show the label inputs for zero-shot. Fires for clicks and example loads.""" | |
| return gr.update(choices=list(TASKS[task])), gr.update(visible=task == ZERO_SHOT_TASK) | |
| def on_task_input(task, defaults): | |
| """On a user click only: pick the task's default model and its threshold.""" | |
| model_name = defaults[task] | |
| return gr.update(value=model_name), gr.update(value=model_spec(task, model_name)["conf"]) | |
| def on_model_input(task, model_name): | |
| return gr.update(value=model_spec(task, model_name)["conf"]) | |
| def on_video_change(video): | |
| """Report the source clip and default the detection rate to its frame rate.""" | |
| if video is None: | |
| return gr.update(value=""), gr.update() | |
| info = sv.VideoInfo.from_video_path(video) | |
| fps = video_fps(video) | |
| w, h = info.resolution_wh | |
| seconds = info.total_frames / fps if info.total_frames else 0 | |
| text = f"Detected fps: {fps:g} · {w}×{h} · {info.total_frames} frames ({seconds:.1f} s)" | |
| return gr.update(value=text), gr.update(value=min(round(fps, 2), MAX_DETECT_FPS)) | |
| def on_captioners_change(captioners): | |
| return gr.update(visible=needs_prompt(captioners)) | |
| def inference_image(image, task, model_name, labels_text, labels_preset, labels_upload, conf_threshold, max_detection, with_caption, caption_detail): | |
| if image is None: | |
| return placeholder("No image provided"), "" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| processor, model = load_detector(model_spec(task, model_name)) | |
| model.to(device) | |
| vocabulary = vocabulary_for(task, labels_text, labels_preset, labels_upload) | |
| rgb = np.array(image.convert("RGB")) | |
| det = detect_at_resolution(processor, model, rgb, conf_threshold, device, vocabulary, SOURCE_RESOLUTION) | |
| det = cap_detections(det, int(max_detection)) | |
| annotators = build_annotators((rgb.shape[1], rgb.shape[0]), tracking=False) | |
| out_bgr = annotate(rgb[..., ::-1], det, annotators) | |
| text = "" | |
| if with_caption: | |
| cap_processor, cap_model = load_image_captioner() | |
| cap_model.to(device) | |
| text = caption_image(cap_processor, cap_model, rgb, caption_detail, device) | |
| return Image.fromarray(out_bgr[..., ::-1]), text | |
| def video_gpu_duration(video, task, model_name, labels_text, labels_preset, labels_upload, conf_threshold, max_detection, tracking, smoothing, hold_low_conf, detect_fps, max_side, lost_seconds, min_frames, match_iou, captioners, *_): | |
| """ZeroGPU budget: tracking plus one minute per selected captioner, capped at five minutes.""" | |
| return min(300, 120 + 60 * len(captioners or [])) | |
| def inference_video( | |
| video, task, model_name, labels_text, labels_preset, labels_upload, conf_threshold, max_detection, | |
| tracking, smoothing, hold_low_conf, | |
| detect_fps, max_side, lost_seconds, min_frames, match_iou, | |
| captioners, prompt, sample_fps, max_frames, | |
| ): | |
| if video is None: | |
| return None, "" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| processor, model = load_detector(model_spec(task, model_name)) | |
| model.to(device) | |
| vocabulary = vocabulary_for(task, labels_text, labels_preset, labels_upload) | |
| info = sv.VideoInfo.from_video_path(video) | |
| fps = video_fps(video) | |
| # Detect every `stride`-th frame and carry the last result across the skipped ones. | |
| # The output keeps the source frame rate either way. | |
| stride = max(1, round(fps / detect_fps)) if detect_fps and detect_fps > 0 else 1 | |
| effective_fps = fps / stride | |
| # With ByteTrack, strict mode never shows a box below the slider; only the smoother | |
| # bridges short gaps. Hold mode lets the tracker keep an already-confirmed track | |
| # alive with boxes down to half the slider value (never below VIDEO_LOW_CONF). | |
| # Without ByteTrack every frame is detected on its own at the slider threshold. | |
| model_conf = max(VIDEO_LOW_CONF, conf_threshold * 0.5) if (tracking and hold_low_conf) else conf_threshold | |
| tracker = smoother = None | |
| if tracking: | |
| tracker = ByteTrackTracker( | |
| frame_rate=float(effective_fps), | |
| track_activation_threshold=conf_threshold, | |
| high_conf_det_threshold=conf_threshold, | |
| lost_track_buffer=max(1, round(lost_seconds * effective_fps)), | |
| minimum_consecutive_frames=int(min_frames), | |
| minimum_iou_threshold=float(match_iou), | |
| ) | |
| smoother = sv.DetectionsSmoother(length=int(smoothing)) | |
| annotators = build_annotators(info.resolution_wh, tracking=tracking) | |
| raw_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name | |
| wrote_any = False | |
| det = sv.Detections.empty() | |
| with sv.VideoSink(raw_out, info, codec="mp4v") as sink: | |
| for i, frame_bgr in enumerate(sv.get_video_frames_generator(video)): | |
| if i % stride == 0: | |
| det = detect_at_resolution(processor, model, frame_bgr[..., ::-1], model_conf, device, vocabulary, int(max_side)) | |
| if tracker is not None: | |
| det = tracker.update(det) | |
| if det.tracker_id is not None: | |
| det = det[det.tracker_id != -1] # drop not-yet-confirmed tracks | |
| det = smoother.update_with_detections(det) | |
| det = cap_detections(det, int(max_detection)) | |
| sink.write_frame(annotate(frame_bgr, det, annotators)) | |
| wrote_any = True | |
| if not wrote_any: | |
| return None, "" | |
| # The detector leaves the GPU first, then captioners run one at a time, so the 9B model fits. | |
| model.to("cpu") | |
| sections = [] | |
| if captioners: | |
| num_frames = count_frames(video, sample_fps, max_frames) | |
| for name in captioners or []: | |
| spec = VIDEO_CAPTIONERS[name] | |
| cap_processor, cap_model = load_video_captioner(spec["repo"]) | |
| cap_model.to(device) | |
| try: | |
| text = caption_video( | |
| cap_processor, cap_model, video, | |
| MARLIN_PROMPT if spec["marlin"] else prompt, | |
| num_frames, device, spec["marlin"], | |
| ) | |
| finally: | |
| cap_model.to("cpu") | |
| if device == "cuda": | |
| torch.cuda.empty_cache() | |
| sections.append(f"### {name}\n{format_marlin(text) if spec['marlin'] else text}") | |
| return to_h264(raw_out), "\n\n".join(sections) | |
| # ----------------------------- | |
| # UI | |
| # ----------------------------- | |
| def model_picker(defaults: dict): | |
| """Task radio, model dropdown and the zero-shot label inputs, wired together.""" | |
| task = gr.Radio(list(TASKS), value=DEFAULT_TASK, label="Task") | |
| model = gr.Dropdown(list(TASKS[DEFAULT_TASK]), value=defaults[DEFAULT_TASK], label="Model") | |
| with gr.Column(visible=False) as labels_box: | |
| labels_preset = gr.Dropdown(list(LABEL_PRESETS), value=DEFAULT_LABEL_PRESET, label="Labels preset") | |
| with gr.Accordion("Custom labels", open=False): | |
| labels_upload = gr.File(file_types=[".json"], type="filepath", label="Upload labels JSON (class → prompts), overrides the preset") | |
| labels_text = gr.Textbox(value="", label="Typed labels (comma-separated), override everything", placeholder="head, tie") | |
| return task, model, labels_box, labels_text, labels_preset, labels_upload | |
| def wire_picker(task, model, labels_box, conf, defaults): | |
| task.change(fn=on_task_change, inputs=task, outputs=[model, labels_box]) | |
| task.input(fn=lambda t: on_task_input(t, defaults), inputs=task, outputs=[model, conf]) | |
| model.input(fn=on_model_input, inputs=[task, model], outputs=conf) | |
| with gr.Blocks() as app: | |
| gr.Markdown("# Video Intelligence") | |
| gr.Markdown( | |
| "Detect, segment, track and caption images and videos. Zero-shot models take labels from a preset " | |
| "vocabulary (default, COCO, Objects365, Open Images), an uploaded JSON, or a typed list. " | |
| ) | |
| with gr.Tabs(): | |
| with gr.Tab("Image"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(type="pil", label="Image") | |
| task_img, model_img, labels_box_img, labels_img, preset_img, upload_img = model_picker(DEFAULT_MODEL) | |
| with gr.Accordion("Detection settings", open=False): | |
| conf_img = gr.Slider(0.05, 0.95, value=0.5, step=0.05, label="Confidence threshold") | |
| max_det_img = gr.Slider(1, 300, step=1, value=300, label="Max detections") | |
| with gr.Accordion("Caption", open=False): | |
| caption_img = gr.Checkbox(value=True, label="Caption with Florence-2") | |
| detail_img = gr.Dropdown(list(CAPTION_TASKS), value=DEFAULT_CAPTION_DETAIL, label="Detail") | |
| infer_image_button = gr.Button("Run", variant="primary") | |
| with gr.Column(): | |
| output_image = gr.Image(type="pil", show_label=False) | |
| output_caption = gr.Textbox(label="Caption", lines=3) | |
| gr.DeepLinkButton(variant="primary") | |
| gr.Examples( | |
| examples=[ | |
| ["zidane.jpg", "Detection (COCO)", "RF-DETR Medium", "", "default", None, 0.5, 300, True, "Brief"], | |
| ["zidane.jpg", "Zero-shot detection", "LLMDet Tiny", "", "coco", None, 0.3, 300, False, "Brief"], | |
| ["zidane.jpg", "Zero-shot detection", "OWLv2 Base", "head, tie", "default", None, 0.2, 300, False, "Brief"], | |
| ["bus.jpg", "Zero-shot detection", "OmDet-Turbo Tiny", "", "objects365", None, 0.3, 300, False, "Brief"], | |
| ["bus.jpg", "Detection (COCO)", "RF-DETR Large", "", "default", None, 0.5, 300, True, "Detailed"], | |
| ["San Diego Airport.jpg", "Segmentation (COCO)", "RF-DETR Seg Large", "", "default", None, 0.5, 300, True, "More detailed"], | |
| ], | |
| fn=inference_image, | |
| inputs=[image, task_img, model_img, labels_img, preset_img, upload_img, conf_img, max_det_img, caption_img, detail_img], | |
| outputs=[output_image, output_caption], | |
| label="Examples", | |
| ) | |
| with gr.Tab("Video"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| video = gr.Video(label="Video") | |
| video_info = gr.Markdown() | |
| task_vid, model_vid, labels_box_vid, labels_vid, preset_vid, upload_vid = model_picker(DEFAULT_VIDEO_MODEL) | |
| with gr.Accordion("Detection settings", open=False): | |
| conf_vid = gr.Slider(0.05, 0.95, value=0.5, step=0.05, label="Confidence threshold") | |
| max_det_vid = gr.Slider(1, 300, step=1, value=100, label="Max detections") | |
| with gr.Accordion("Advanced", open=False): | |
| detect_fps_vid = gr.Slider(1, MAX_DETECT_FPS, step=0.5, value=FALLBACK_FPS, label="Detection fps", info="Defaults to the source frame rate on upload") | |
| max_side_vid = gr.Slider(0, 1920, step=64, value=SOURCE_RESOLUTION, label="Inference resolution (longest side, 0 = source)") | |
| with gr.Accordion("Captions", open=False): | |
| captioners_vid = gr.CheckboxGroup(list(VIDEO_CAPTIONERS), value=DEFAULT_VIDEO_CAPTIONERS, label="Models") | |
| prompt_vid = gr.Textbox(value=DEFAULT_VIDEO_PROMPT, label="Prompt", lines=2) | |
| with gr.Row(): | |
| fps_vid = gr.Slider(0.5, 4, step=0.5, value=2, label="Sampled frames per second") | |
| frames_vid = gr.Slider(8, 240, step=8, value=64, label="Max sampled frames") | |
| with gr.Accordion("Tracking", open=False): | |
| track_vid = gr.Checkbox(value=True, label="Enable ByteTrack") | |
| hold_vid = gr.Checkbox(value=True, label="Hold tracks through low-confidence frames") | |
| smooth_vid = gr.Slider(1, 15, step=1, value=5, label="Smoothing window (frames)") | |
| with gr.Accordion("Advanced", open=False): | |
| lost_vid = gr.Slider(0.1, 5, step=0.1, value=1.0, label="Lost track buffer (seconds)") | |
| min_frames_vid = gr.Slider(1, 10, step=1, value=2, label="Frames before a track shows") | |
| iou_vid = gr.Slider(0.05, 0.9, step=0.05, value=0.1, label="Track match IoU") | |
| infer_video_button = gr.Button("Run", variant="primary") | |
| with gr.Column(): | |
| output_video = gr.Video(show_label=False) | |
| output_captions = gr.Markdown(label="Captions", container=True, show_label=True) | |
| gr.DeepLinkButton(variant="primary") | |
| with gr.Accordion("Models and references", open=False): | |
| gr.Markdown( | |
| """ | |
| All detectors load from `models/` in this Space, through 🤗 Transformers. Apache 2.0 unless noted. | |
| - **RF-DETR** (COCO detection and segmentation): [paper](https://huggingface.co/papers/2511.09554) | |
| - **OWLv2** (zero-shot): [paper](https://huggingface.co/papers/2306.09683) | |
| - **Grounding DINO** and **MM Grounding DINO** (zero-shot): [paper](https://huggingface.co/papers/2303.05499), [paper](https://huggingface.co/papers/2401.02361) | |
| - **LLMDet** (zero-shot, MM Grounding DINO with LLM co-training): [paper](https://huggingface.co/papers/2501.18954) | |
| - **OmDet-Turbo** (zero-shot, real-time): [paper](https://huggingface.co/papers/2403.06892) | |
| - **Florence-2** (image captions, MIT): [paper](https://huggingface.co/papers/2311.06242) | |
| - **Qwen3.5-4B / 9B** and **Marlin-2B** (video captions) | |
| """ | |
| ) | |
| wire_picker(task_img, model_img, labels_box_img, conf_img, DEFAULT_MODEL) | |
| wire_picker(task_vid, model_vid, labels_box_vid, conf_vid, DEFAULT_VIDEO_MODEL) | |
| video.change(fn=on_video_change, inputs=video, outputs=[video_info, detect_fps_vid]) | |
| captioners_vid.change(fn=on_captioners_change, inputs=captioners_vid, outputs=prompt_vid) | |
| infer_image_button.click( | |
| fn=inference_image, | |
| inputs=[image, task_img, model_img, labels_img, preset_img, upload_img, conf_img, max_det_img, caption_img, detail_img], | |
| outputs=[output_image, output_caption], | |
| ) | |
| infer_video_button.click( | |
| fn=inference_video, | |
| inputs=[ | |
| video, task_vid, model_vid, labels_vid, preset_vid, upload_vid, conf_vid, max_det_vid, | |
| track_vid, smooth_vid, hold_vid, | |
| detect_fps_vid, max_side_vid, lost_vid, min_frames_vid, iou_vid, | |
| captioners_vid, prompt_vid, fps_vid, frames_vid, | |
| ], | |
| outputs=[output_video, output_captions], | |
| ) | |
| if __name__ == "__main__": | |
| hub.prefetch([FLORENCE] + DETECTOR_SPECS) | |
| app.launch(mcp_server=True, theme=gr.themes.Ocean(primary_hue="indigo", secondary_hue="blue")) | |