Molmo2Fish interactive sonar fish tracking demo

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/elwha_2018-07-29.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/kenai_leftfar_2018-06-03.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/nushagak_rb_f15-52.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,55 @@
 ---
-title: Molmo2fish Tracking
-emoji: 🦀
-colorFrom: red
-colorTo: gray
+title: Molmo2Fish Tracking
+emoji: 🐟
+colorFrom: purple
+colorTo: blue
 sdk: gradio
 sdk_version: 6.25.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Track fish in sonar video, fix it with plain English
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🐟 Molmo2Fish — interactive fish tracking with natural language guidance
+
+Demo of [**tidalove/Molmo2Fish**](https://huggingface.co/tidalove/Molmo2Fish), the model from
+*"Teach a Molmo2Fish: Towards interactive fish tracking with natural language guidance"*
+([paper](https://huggingface.co/papers/2608.18602) ·
+[code](https://github.com/tidalove/molmo2fish)).
+
+Molmo2Fish is a LoRA-finetuned [Molmo2](https://huggingface.co/allenai) VLM that tracks
+salmon in ARIS **sonar** video and — crucially — accepts a plain-English critique of its
+own output and re-emits corrected tracks. The paper reports tracking accuracy going from
+~5% to ~79% over a handful of conversational correction turns.
+
+## How the demo works
+
+1. **① Track all fish** sends the clip with the prompt `track all fish`. The model replies
+   with the html-v2 pointing format used in training:
+   `<tracks coords="0.0 1 409 852;0.5 1 436 890 2 300 120;…">fish</tracks>`
+   (timestamps in seconds at 2 FPS, ids, and x/y normalised to 0–1000).
+2. **② Apply correction** rebuilds the conversation — video on the first user turn, the
+   model's previous `<tracks …>` answer as the assistant turn, your critique as the new
+   user turn — exactly as in `olmo/eval/vllm_runner.py::build_multi_turn_chat`, and the
+   model regenerates the track set.
+
+Points are parsed with the same regexes as
+`olmo/preprocessing/point_formatter.py` and drawn back onto the 6 FPS source video.
+
+## Example clips
+
+The three sonar clips shipped with this Space are re-encoded (G channel, 6 FPS, libx264,
+matching the repo's `encode_frames_to_video` recipe) from the Caltech Fish Counting
+release [**perona-lab/cfc26**](https://huggingface.co/datasets/perona-lab/cfc26), which is
+distributed under **CC-BY-4.0** — credit to the Caltech Fish Counting / CFC26 authors.
+The correction prompts pre-filled with each example are verbatim from the validation split
+of [tidalove/cfc-track-instruction](https://huggingface.co/datasets/tidalove/cfc-track-instruction).
+
+## Notes
+
+- Runs on ZeroGPU; the model is loaded in bfloat16 (~16 GB VRAM).
+- Video is sampled at 2 FPS, max 128 frames, per the model's
+  `video_preprocessor_config.json`.
+- A truncated answer (no closing `</tracks>`) means you hit the *Max new tokens* cap —
+  raise it in **Advanced**.

app.py

@@ -0,0 +1,401 @@
+"""Molmo2Fish — interactive fish tracking in ARIS sonar video with natural-language guidance.
+
+Paper: "Teach a Molmo2Fish: Towards interactive fish tracking with natural language
+guidance" (arXiv 2608.18602). Model: tidalove/Molmo2Fish.
+
+The demo mirrors the paper's two-stage correction loop:
+  1. an initial pass ("track all fish") produces `<tracks coords="...">fish</tracks>`
+  2. the user types a plain-English critique and the model re-emits corrected tracks,
+     conditioned on the video, its own previous answer, and the critique.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  — must precede torch / CUDA-touching imports
+
+import re  # noqa: E402
+import tempfile  # noqa: E402
+import time  # noqa: E402
+from collections import defaultdict  # noqa: E402
+
+import cv2  # noqa: E402
+import gradio as gr  # noqa: E402
+import imageio.v2 as imageio  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from molmo_utils import process_vision_info  # noqa: E402
+from transformers import AutoModelForImageTextToText, AutoProcessor  # noqa: E402
+
+MODEL_ID = "tidalove/Molmo2Fish"
+
+# Matches the released video_preprocessor_config.json of tidalove/Molmo2Fish.
+NUM_FRAMES = 128
+FRAME_SAMPLE_MODE = "uniform_last_frame"
+MAX_FPS = 2
+SAMPLING_FPS = 2
+
+TRACK_STYLE = "video_point_track_per_frame"
+DEFAULT_PROMPT = "track all fish"
+
+# html-v2 pointing format, exactly as in olmo/preprocessing/point_formatter.py
+COORD_RE = re.compile(r"<(?:points|tracks).*? coords=\"([0-9\t:;, .]+)\"/?>")
+FRAME_RE = re.compile(r"(?:^|\t|:|,|;)([0-9\.]+) ([0-9\. ]+)")
+POINTS_RE = re.compile(r"([0-9]+) ([0-9]{3,4}) ([0-9]{3,4})")
+
+PALETTE = [
+    (240, 82, 156),   # the authors' pink (scripts/unified_demo.py)
+    (66, 214, 255),
+    (124, 252, 118),
+    (255, 196, 61),
+    (186, 132, 255),
+    (255, 122, 92),
+    (0, 255, 214),
+    (255, 255, 120),
+]
+
+print(f"Loading {MODEL_ID} …", flush=True)
+processor = AutoProcessor.from_pretrained(
+    MODEL_ID, trust_remote_code=True, padding_side="left"
+)
+model = AutoModelForImageTextToText.from_pretrained(
+    MODEL_ID, trust_remote_code=True, dtype=torch.bfloat16
+).to("cuda")
+model.eval()
+print("Model ready.", flush=True)
+
+
+# --------------------------------------------------------------------------- #
+# Track parsing / rendering
+# --------------------------------------------------------------------------- #
+def parse_tracks(text: str, width: int, height: int) -> dict:
+    """Parse `<tracks coords="t id x y …">fish</tracks>` into {time: {id: (x, y)}}.
+
+    Coordinates in the model output are normalised to 0-1000; they are scaled
+    back to pixels here.
+    """
+    out: dict = {}
+    for coord in COORD_RE.finditer(text):
+        for frame in FRAME_RE.finditer(coord.group(1)):
+            t = float(frame.group(1))
+            per_frame = out.setdefault(t, {})
+            for pt in POINTS_RE.finditer(frame.group(2)):
+                idx, xs, ys = pt.group(1), pt.group(2), pt.group(3)
+                x = float(xs) / 1000.0 * width
+                y = float(ys) / 1000.0 * height
+                if 0 <= x <= width and 0 <= y <= height:
+                    per_frame.setdefault(idx, (x, y))
+    return out
+
+
+def render_overlay(video_path: str, tracks: dict, out_path: str) -> None:
+    """Draw the parsed tracks (points + fading trails + ids) onto the source video."""
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 6.0
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    times = sorted(tracks)
+    times_arr = np.asarray(times) if times else None
+
+    # id -> ordered list of (time, x, y), used to draw the trail behind each fish
+    history = defaultdict(list)
+    for t in times:
+        for idx, (x, y) in tracks[t].items():
+            history[idx].append((t, x, y))
+
+    ids = sorted(history, key=lambda s: (len(s), s))
+    color_of = {idx: PALETTE[i % len(PALETTE)] for i, idx in enumerate(ids)}
+
+    radius = max(4, int(max(width, height) * 0.008))
+    thickness = max(2, radius // 2)
+    font_scale = max(0.5, max(width, height) / 1400.0)
+
+    writer = imageio.get_writer(
+        out_path, fps=fps, codec="libx264", quality=7,
+        macro_block_size=1, pixelformat="yuv420p", ffmpeg_log_level="error",
+    )
+    try:
+        frame_ix = 0
+        while True:
+            ok, frame = cap.read()
+            if not ok:
+                break
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            if times_arr is not None:
+                t_now = frame_ix / fps
+                k = int(np.argmin(np.abs(times_arr - t_now)))
+                t_key = times[k]
+                for idx, pts in history.items():
+                    trail = [(x, y) for (t, x, y) in pts if t <= t_key]
+                    if len(trail) > 1:
+                        poly = np.asarray(trail[-24:], dtype=np.int32).reshape(-1, 1, 2)
+                        cv2.polylines(rgb, [poly], False, color_of[idx],
+                                      max(1, thickness - 1), cv2.LINE_AA)
+                for idx, (x, y) in tracks[t_key].items():
+                    c = color_of[idx]
+                    cv2.circle(rgb, (int(x), int(y)), radius, c, thickness, cv2.LINE_AA)
+                    cv2.putText(rgb, str(idx), (int(x) + radius + 3, int(y) - radius - 3),
+                                cv2.FONT_HERSHEY_SIMPLEX, font_scale, c,
+                                max(1, thickness - 1), cv2.LINE_AA)
+            writer.append_data(rgb)
+            frame_ix += 1
+    finally:
+        writer.close()
+        cap.release()
+
+
+def summarise(tracks: dict) -> str:
+    if not tracks:
+        return "No fish tracks were returned for this clip."
+    ids = {i for frame in tracks.values() for i in frame}
+    return (f"**{len(ids)} track(s)** across **{len(tracks)}** sampled timesteps "
+            f"(2 FPS). Track ids: {', '.join(sorted(ids, key=int))}.")
+
+
+# --------------------------------------------------------------------------- #
+# Model plumbing
+# --------------------------------------------------------------------------- #
+def build_messages(video_path: str, turns: list) -> list:
+    """Chat list for Molmo2Fish. The video is attached to the *first* user turn only.
+
+    `turns` is a list of (user_text, assistant_text_or_None), matching
+    olmo/eval/vllm_runner.py::build_multi_turn_chat.
+    """
+    messages = []
+    for i, (user_text, assistant_text) in enumerate(turns):
+        content = [dict(type="text", text=user_text, style=TRACK_STYLE)]
+        if i == 0:
+            content.append(dict(
+                type="video",
+                video=video_path,
+                num_frames=NUM_FRAMES,
+                frame_sample_mode=FRAME_SAMPLE_MODE,
+                max_fps=MAX_FPS,
+                sampling_fps=SAMPLING_FPS,
+            ))
+        messages.append({"role": "user", "content": content})
+        if assistant_text is not None:
+            messages.append({"role": "assistant",
+                             "content": [dict(type="text", text=assistant_text)]})
+    return messages
+
+
+def run_model(video_path: str, turns: list, max_new_tokens: int) -> str:
+    messages = build_messages(video_path, turns)
+    _, videos, video_kwargs = process_vision_info(messages)
+    frames, metadatas = zip(*videos)
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = processor(
+        videos=list(frames),
+        video_metadata=list(metadatas),
+        text=text,
+        padding=True,
+        return_tensors="pt",
+        **video_kwargs,
+    )
+    inputs = {k: (v.to(model.device) if hasattr(v, "to") else v)
+              for k, v in inputs.items()}
+    with torch.inference_mode():
+        with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+            generated = model.generate(
+                **inputs, max_new_tokens=max_new_tokens, do_sample=False
+            )
+    prompt_len = inputs["input_ids"].size(1)
+    return processor.post_process_image_text_to_text(
+        generated[:, prompt_len:],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=False,
+    )[0].strip()
+
+
+def _video_size(video_path: str):
+    cap = cv2.VideoCapture(video_path)
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    cap.release()
+    return w, h
+
+
+def _infer(video_path: str, turns: list, max_new_tokens: int):
+    t0 = time.perf_counter()
+    raw = run_model(video_path, turns, max_new_tokens)
+    elapsed = time.perf_counter() - t0
+    width, height = _video_size(video_path)
+    tracks = parse_tracks(raw, width, height)
+    if not tracks:
+        return video_path, raw, f"{summarise(tracks)}  \n_Inference: {elapsed:.1f}s_"
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    render_overlay(video_path, tracks, out_path)
+    return out_path, raw, f"{summarise(tracks)}  \n_Inference: {elapsed:.1f}s_"
+
+
+# --------------------------------------------------------------------------- #
+# Gradio handlers
+# --------------------------------------------------------------------------- #
+@spaces.GPU(duration=140)
+def track_fish(
+    video_path: str,
+    correction_hint: str = "",
+    max_new_tokens: int = 1600,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Run the first tracking pass over a sonar clip ("track all fish").
+
+    Args:
+        video_path: path to an ARIS sonar clip (mp4).
+        correction_hint: ignored here — it only exists so an example row can
+            pre-fill the correction box alongside the video.
+        max_new_tokens: generation budget for the `<tracks …>` string.
+
+    Returns:
+        (overlay video, raw model output, markdown summary)
+    """
+    if not video_path:
+        raise gr.Error("Please provide a sonar video first.")
+    return _infer(video_path, [(DEFAULT_PROMPT, None)], int(max_new_tokens))
+
+
+@spaces.GPU(duration=140)
+def refine_tracks(
+    video_path: str,
+    previous_tracks: str,
+    correction: str,
+    max_new_tokens: int = 1600,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Correct the current tracks using a natural-language instruction.
+
+    The model sees the video, its own previous `<tracks …>` answer, and the
+    critique, then re-emits a corrected track set.
+
+    Args:
+        video_path: the same sonar clip used for the first pass.
+        previous_tracks: the model's previous `<tracks …>` output.
+        correction: plain-English critique, e.g. "Track 1 is shifted downward".
+        max_new_tokens: generation budget for the corrected `<tracks …>` string.
+
+    Returns:
+        (overlay video, raw model output, markdown summary)
+    """
+    if not video_path:
+        raise gr.Error("Please provide a sonar video first.")
+    if not previous_tracks or not previous_tracks.strip():
+        raise gr.Error("Run 'Track all fish' first — there is nothing to correct yet.")
+    if not correction or not correction.strip():
+        raise gr.Error("Type a correction instruction, e.g. 'Track 1 is shifted downward'.")
+    turns = [(DEFAULT_PROMPT, previous_tracks.strip()), (correction.strip(), None)]
+    return _infer(video_path, turns, int(max_new_tokens))
+
+
+# --------------------------------------------------------------------------- #
+# UI
+# --------------------------------------------------------------------------- #
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+EXAMPLES = [
+    [
+        "examples/elwha_2018-07-29.mp4",
+        "Track 1 looks good overall, just slightly shifted downward from the "
+        "actual fish position throughout.",
+    ],
+    [
+        "examples/kenai_leftfar_2018-06-03.mp4",
+        "Track 1 doesn't correspond to any real fish — you've got a false "
+        "detection moving left that should be removed. The actual fish starts "
+        "in the lower left around 8s and swims upward until the end of the clip, "
+        "and you missed it entirely.",
+    ],
+    [
+        "examples/nushagak_rb_f15-52.mp4",
+        "You missed a fish near the top of the frame — please add it.",
+    ],
+]
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# 🐟 Molmo2Fish — interactive fish tracking\n"
+            "Track salmon in ARIS **sonar** video, then fix the mistakes by *talking to the model*.\n\n"
+            "Step 1 runs the model's `track all fish` pass. Step 2 feeds your plain-English "
+            "critique back in — the model re-emits a corrected track set instead of you "
+            "editing keypoints by hand.\n\n"
+            "[Paper](https://huggingface.co/papers/2608.18602) · "
+            "[Model](https://huggingface.co/tidalove/Molmo2Fish) · "
+            "[Code](https://github.com/tidalove/molmo2fish) · "
+            "[Data](https://huggingface.co/datasets/tidalove/cfc-track-instruction)"
+        )
+
+        with gr.Row():
+            with gr.Column():
+                video_in = gr.Video(label="Sonar clip", height=420)
+                track_btn = gr.Button("① Track all fish", variant="primary")
+                correction = gr.Textbox(
+                    label="② Correction instruction",
+                    placeholder="Track 2 drifts off the fish after about 6s — it should keep "
+                                "following the fish swimming up the right side.",
+                    lines=3,
+                )
+                refine_btn = gr.Button("② Apply correction", variant="secondary")
+            with gr.Column():
+                video_out = gr.Video(label="Tracks", height=420, autoplay=True)
+                summary = gr.Markdown()
+                tracks_box = gr.Textbox(
+                    label="Model output (html-v2 tracks) — edited in place by step ②",
+                    lines=6,
+                    max_lines=12,
+                    show_copy_button=True,
+                )
+
+        with gr.Accordion("Advanced", open=False):
+            max_new_tokens = gr.Slider(
+                256, 3072, value=1600, step=64,
+                label="Max new tokens",
+                info="Long clips with many fish need a bigger budget; an unclosed "
+                     "</tracks> means you hit the cap.",
+            )
+
+        gr.Markdown(
+            "### Examples\n"
+            "Clicking a row loads the clip **and** pre-fills a real correction from the "
+            "paper's CFC validation split, and runs step ① for you."
+        )
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[video_in, correction],
+            outputs=[video_out, tracks_box, summary],
+            fn=track_fish,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Sonar clips (CFC26, CC-BY-4.0)",
+        )
+
+        gr.Markdown(
+            "Sonar clips are re-encoded from the "
+            "[perona-lab/cfc26](https://huggingface.co/datasets/perona-lab/cfc26) "
+            "Caltech Fish Counting release (CC-BY-4.0); correction prompts come from "
+            "[tidalove/cfc-track-instruction](https://huggingface.co/datasets/tidalove/cfc-track-instruction). "
+            "Tracks are predicted at 2 FPS and interpolated onto the 6 FPS source for display."
+        )
+
+    track_btn.click(
+        track_fish,
+        inputs=[video_in, correction, max_new_tokens],
+        outputs=[video_out, tracks_box, summary],
+        api_name="track_fish",
+    )
+    refine_btn.click(
+        refine_tracks,
+        inputs=[video_in, tracks_box, correction, max_new_tokens],
+        outputs=[video_out, tracks_box, summary],
+        api_name="refine_tracks",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

examples/elwha_2018-07-29.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86daba3bd947f30d983a7a2c9b2565b69b706bb4e49e861410274940bf94571a
+size 25601886

examples/kenai_leftfar_2018-06-03.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec527b78fbcdf2762e4956cbf028822f5379197c5950ffb5d124709ceb378124
+size 17401090

examples/nushagak_rb_f15-52.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:567a6f771cff7858cdcb9df77884c94d0af429def51b525a62e4edf5b057b303
+size 6778063

requirements.txt

@@ -0,0 +1,12 @@
+transformers==4.57.6
+accelerate
+molmo_utils
+torchvision
+av
+einops
+timm
+opencv-python-headless
+imageio
+imageio-ffmpeg
+numpy
+pillow