"""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)) @spaces.GPU 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 [])) @spaces.GPU(duration=video_gpu_duration) 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"))