run_vos_demo.py

gistfile1.txt

"""Training-free video object segmentation with frozen LingBot-Vision features."""

from __future__ import annotations

import argparse
import math
import time
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

torch.set_grad_enabled(False)

from lingbot_vision import load_pretrained_backbone
from lingbot_vision.loader import extract_patch_tokens
from lingbot_vision.preprocess import IMAGENET_MEAN, IMAGENET_STD

MAX_CONTEXT_LENGTH = 7

QUALITY_PRESETS = {
    "fast": {
        "short_side": 480,
        "topk": 5,
        "neighborhood_size": 12,
        "temperature": 0.2,
        "upsample": "nearest",
        "refine_seed": False,
        "refine_edges": False,
        "feature_blend": 0.0,
    },
    "balanced": {
        "short_side": 640,
        "topk": 5,
        "neighborhood_size": 12,
        "temperature": 0.2,
        "upsample": "bilinear",
        "refine_seed": True,
        "refine_edges": True,
        "feature_blend": 0.15,
    },
    "high": {
        "short_side": 720,
        "topk": 8,
        "neighborhood_size": 14,
        "temperature": 0.15,
        "upsample": "bilinear",
        "refine_seed": True,
        "refine_edges": True,
        "feature_blend": 0.25,
    },
}


def _repo_root() -> Path:
    return Path(__file__).resolve().parent


def _round_up(side: float, multiple: int) -> int:
    return math.ceil(side / multiple) * multiple


def resize_shortest_side(img_rgb: np.ndarray, short_side: int, patch_size: int) -> np.ndarray:
    """Resize so the shorter side equals short_side, rounded up to patch_size multiples."""
    height, width = img_rgb.shape[:2]
    if width > height:
        new_height = _round_up(short_side, patch_size)
        new_width = _round_up(width * new_height / height, patch_size)
    else:
        new_width = _round_up(short_side, patch_size)
        new_height = _round_up(height * new_width / width, patch_size)
    return cv2.resize(img_rgb, (new_width, new_height), interpolation=cv2.INTER_LINEAR)


def normalize_rgb(img_rgb: np.ndarray) -> torch.Tensor:
    img_t = torch.from_numpy(img_rgb.astype(np.float32) / 255.0).permute(2, 0, 1)
    return (img_t.unsqueeze(0) - IMAGENET_MEAN) / IMAGENET_STD


def extract_features(backbone, img_norm: torch.Tensor, device: str, dtype: torch.dtype) -> torch.Tensor:
    tokens, (h, w) = extract_patch_tokens(backbone, img_norm, device, dtype)
    feats = tokens.reshape(h, w, -1).float()
    return F.normalize(feats, dim=-1, p=2)


def upsample_features(feats: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
    """Bilinearly upsample patch features to pixel resolution and re-normalize."""
    feat_map = feats.permute(2, 0, 1).unsqueeze(0)
    feat_up = F.interpolate(feat_map, size=size, mode="bilinear", align_corners=False)
    feat_up = feat_up.squeeze(0).permute(1, 2, 0)
    return F.normalize(feat_up, dim=-1, p=2)


def refine_seed_mask_with_features(
    frame_rgb: np.ndarray,
    hard_mask: np.ndarray,
    backbone,
    device: str,
    dtype: torch.dtype,
    short_side: int,
    patch_size: int,
) -> np.ndarray:
    """Expand a coarse bbox seed into a feature-similarity mask on the first frame."""
    resized = resize_shortest_side(frame_rgb, short_side, patch_size)
    feats = extract_features(backbone, normalize_rgb(resized).to(device), device, dtype)
    h, w = feats.shape[:2]

    refined_small = np.zeros((h, w), dtype=np.uint8)
    for label in range(1, int(hard_mask.max()) + 1):
        label_small = cv2.resize((hard_mask == label).astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST) > 0
        if not label_small.any():
            continue
        label_t = torch.from_numpy(label_small).to(feats.device)
        proto = F.normalize(feats[label_t].mean(dim=0, keepdim=True), dim=-1)
        sim = torch.einsum("hwc,c->hw", feats, proto.squeeze())
        thresh = torch.quantile(sim[label_t], 0.20).item()
        refined_small[sim.cpu().numpy() >= thresh] = label

    if refined_small.max() == 0:
        return hard_mask
    return cv2.resize(refined_small, (hard_mask.shape[1], hard_mask.shape[0]), interpolation=cv2.INTER_NEAREST)


def upsample_probs(
    probs: torch.Tensor,
    size: tuple[int, int],
    mode: str,
) -> torch.Tensor:
    align = False if mode == "bilinear" else None
    upsampled = F.interpolate(probs.movedim(-1, -3)[None], size=size, mode=mode, align_corners=align)
    return postprocess_probs(upsampled).squeeze(0)


def refine_prob_edges(probs: torch.Tensor) -> torch.Tensor:
    """Edge-aware smoothing on class probability maps."""
    refined = []
    for idx in range(probs.shape[0]):
        channel = probs[idx].detach().float().cpu().numpy()
        smooth = cv2.bilateralFilter(channel, d=7, sigmaColor=0.08, sigmaSpace=7)
        refined.append(smooth)
    out = torch.from_numpy(np.stack(refined)).to(device=probs.device, dtype=torch.float32)
    out = out / out.sum(dim=0, keepdim=True).clamp_min(1e-6)
    return out


def blend_with_feature_similarity(
    probs: torch.Tensor,
    current_feats: torch.Tensor,
    first_feats: torch.Tensor,
    first_probs: torch.Tensor,
    proc_size: tuple[int, int],
    blend_weight: float,
) -> torch.Tensor:
    """Sharpen boundaries by blending propagated probs with first-frame feature affinity."""
    if blend_weight <= 0:
        return probs

    current_up = upsample_features(current_feats, proc_size)
    num_masks = probs.shape[0]
    affinity = torch.zeros_like(probs)
    for label in range(num_masks):
        weights = first_probs[..., label]
        if float(weights.sum()) < 1e-6:
            continue
        proto = F.normalize((first_feats * weights.unsqueeze(-1)).sum(dim=(0, 1)) / weights.sum(), dim=0)
        sim = torch.einsum("hwc,c->hw", current_up, proto)
        affinity[label] = sim.to(probs.dtype)

    affinity = postprocess_probs(affinity.unsqueeze(0)).squeeze(0)
    blended = (1.0 - blend_weight) * probs + blend_weight * affinity
    return blended / blended.sum(dim=0, keepdim=True).clamp_min(1e-6)


def make_neighborhood_mask(h: int, w: int, size: float, device: torch.device) -> torch.Tensor:
    ij = torch.stack(
        torch.meshgrid(
            torch.arange(h, dtype=torch.float32, device=device),
            torch.arange(w, dtype=torch.float32, device=device),
            indexing="ij",
        ),
        dim=-1,
    )
    norm = torch.linalg.vector_norm(
        ij[:, :, None, None, :] - ij[None, None, :, :, :],
        ord=2,
        dim=-1,
    )
    return norm <= size


def propagate(
    current_features: torch.Tensor,
    context_features: torch.Tensor,
    context_probs: torch.Tensor,
    neighborhood_mask: torch.Tensor,
    topk: int,
    temperature: float,
) -> torch.Tensor:
    _, h, w, _ = context_features.shape
    dot = torch.einsum("ijd, tuvd -> ijtuv", current_features, context_features)
    dot = torch.where(neighborhood_mask[:, :, None, :, :], dot, -torch.inf)
    dot = dot.flatten(2, -1).flatten(0, 1)
    k_th_largest = torch.topk(dot, dim=1, k=topk).values
    dot = torch.where(dot >= k_th_largest[:, -1:], dot, -torch.inf)
    weights = F.softmax(dot / temperature, dim=1)
    current_probs = torch.mm(weights, context_probs.flatten(0, 2))
    current_probs = current_probs / current_probs.sum(dim=1, keepdim=True).clamp_min(1e-6)
    return current_probs.unflatten(0, (h, w))


def postprocess_probs(probs: torch.Tensor) -> torch.Tensor:
    vmin = probs.flatten(2, 3).min(dim=2).values
    vmax = probs.flatten(2, 3).max(dim=2).values
    denom = (vmax[:, :, None, None] - vmin[:, :, None, None]).clamp_min(1e-6)
    probs = (probs - vmin[:, :, None, None]) / denom
    return torch.nan_to_num(probs, nan=0.0)


def mask_to_rgb(mask: np.ndarray, num_masks: int) -> np.ndarray:
    background = mask == 0
    indexed = mask.copy()
    rgb = np.zeros((*mask.shape, 3), dtype=np.uint8)
    if num_masks <= 1:
        return rgb
    palette = [
        (255, 0, 0),
        (0, 255, 0),
        (0, 128, 255),
        (255, 165, 0),
        (255, 0, 255),
        (0, 255, 255),
        (255, 255, 0),
        (128, 0, 255),
    ]
    for label in range(1, num_masks):
        color = palette[(label - 1) % len(palette)]
        rgb[mask == label] = color
    rgb[background] = 0
    return rgb


def overlay_mask(frame_rgb: np.ndarray, mask_rgb: np.ndarray, alpha: float = 0.45) -> np.ndarray:
    out = frame_rgb.copy()
    fg = mask_rgb.sum(axis=-1) > 0
    out[fg] = (
        (1.0 - alpha) * frame_rgb[fg].astype(np.float32) + alpha * mask_rgb[fg].astype(np.float32)
    ).astype(np.uint8)
    return out


def read_video(path: Path) -> tuple[list[np.ndarray], float]:
    cap = cv2.VideoCapture(str(path))
    if not cap.isOpened():
        raise FileNotFoundError(f"cannot open video: {path}")
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    frames: list[np.ndarray] = []
    while True:
        ok, frame_bgr = cap.read()
        if not ok:
            break
        frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
    cap.release()
    if not frames:
        raise RuntimeError(f"video has no frames: {path}")
    return frames, fps


def paint_mask_tk(frame_rgb: np.ndarray, brush_size: int = 12) -> np.ndarray | None:
    """Paint the first-frame mask: left=draw, right=erase, Enter=confirm, Esc=cancel."""
    import tkinter as tk
    from PIL import ImageTk

    height, width = frame_rgb.shape[:2]
    mask = np.zeros((height, width), dtype=np.uint8)
    overlay_color = np.array([255, 64, 64], dtype=np.float32)

    root = tk.Tk()
    root.title("Paint mask: LMB=draw, RMB=erase, [/]=brush, Enter=confirm, Esc=cancel")

    header = tk.Label(
        root,
        text=f"Brush size: {brush_size}px  |  Left drag=paint  Right drag=erase  |  Enter=confirm",
        anchor="w",
    )
    header.pack(fill="x", padx=8, pady=6)

    canvas = tk.Canvas(root, width=width, height=height, cursor="dotbox", highlightthickness=0)
    canvas.pack()

    state = {
        "brush": max(1, brush_size),
        "mode": None,
        "last_x": None,
        "last_y": None,
        "photo": None,
        "result": None,
    }

    def render_display() -> ImageTk.PhotoImage:
        display = frame_rgb.copy()
        painted = mask > 0
        display[painted] = (
            0.55 * frame_rgb[painted].astype(np.float32) + 0.45 * overlay_color
        ).astype(np.uint8)
        photo = ImageTk.PhotoImage(Image.fromarray(display))
        state["photo"] = photo
        return photo

    def refresh_canvas() -> None:
        canvas.delete("all")
        canvas.create_image(0, 0, anchor=tk.NW, image=render_display())

    def clamp_xy(x: int, y: int) -> tuple[int, int]:
        return max(0, min(width - 1, x)), max(0, min(height - 1, y))

    def apply_stroke(x0: int, y0: int, x1: int, y1: int, erase: bool) -> None:
        value = 0 if erase else 1
        thickness = max(1, state["brush"] * 2)
        cv2.line(mask, (x0, y0), (x1, y1), value, thickness=thickness, lineType=cv2.LINE_AA)
        cv2.circle(mask, (x1, y1), state["brush"], value, thickness=-1, lineType=cv2.LINE_AA)

    def begin_draw(event: tk.Event, erase: bool) -> None:
        state["mode"] = "erase" if erase else "draw"
        x, y = clamp_xy(event.x, event.y)
        state["last_x"], state["last_y"] = x, y
        value = 0 if erase else 1
        cv2.circle(mask, (x, y), state["brush"], value, thickness=-1, lineType=cv2.LINE_AA)
        refresh_canvas()

    def continue_draw(event: tk.Event) -> None:
        if state["mode"] is None or state["last_x"] is None or state["last_y"] is None:
            return
        x, y = clamp_xy(event.x, event.y)
        erase = state["mode"] == "erase"
        apply_stroke(state["last_x"], state["last_y"], x, y, erase=erase)
        state["last_x"], state["last_y"] = x, y
        refresh_canvas()

    def end_draw(_event: tk.Event) -> None:
        state["mode"] = None
        state["last_x"] = None
        state["last_y"] = None

    def set_brush_size(size: int) -> None:
        state["brush"] = max(1, min(96, size))
        header.config(
            text=(
                f"Brush size: {state['brush']}px  |  "
                "Left drag=paint  Right drag=erase  |  Enter=confirm"
            )
        )

    def on_brush_smaller(_event: tk.Event | None = None) -> None:
        set_brush_size(state["brush"] - 2)

    def on_brush_larger(_event: tk.Event | None = None) -> None:
        set_brush_size(state["brush"] + 2)

    def on_confirm(_event: tk.Event | None = None) -> None:
        if mask.any():
            state["result"] = mask.copy()
        root.destroy()

    def on_cancel(_event: tk.Event | None = None) -> None:
        state["result"] = None
        root.destroy()

    canvas.bind("<ButtonPress-1>", lambda e: begin_draw(e, erase=False))
    canvas.bind("<B1-Motion>", continue_draw)
    canvas.bind("<ButtonRelease-1>", end_draw)
    canvas.bind("<ButtonPress-3>", lambda e: begin_draw(e, erase=True))
    canvas.bind("<B3-Motion>", continue_draw)
    canvas.bind("<ButtonRelease-3>", end_draw)
    root.bind("[", on_brush_smaller)
    root.bind("]", on_brush_larger)
    root.bind("-", on_brush_smaller)
    root.bind("=", on_brush_larger)
    root.bind("<Return>", on_confirm)
    root.bind("<Escape>", on_cancel)
    root.protocol("WM_DELETE_WINDOW", on_cancel)

    refresh_canvas()
    root.mainloop()
    return state["result"]


def load_first_mask(
    mask_path: Path | None,
    frame_rgb: np.ndarray,
    pick_mask: bool,
    seed_bbox: tuple[int, int, int, int] | None,
    brush_size: int = 12,
) -> tuple[np.ndarray, str]:
    height, width = frame_rgb.shape[:2]
    if mask_path is not None:
        mask = np.array(Image.open(mask_path))
        if mask.ndim == 3:
            mask = mask[..., 0]
        if mask.shape[:2] != (height, width):
            mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST)
        return mask.astype(np.uint8), "file"

    if pick_mask:
        painted = paint_mask_tk(frame_rgb, brush_size=brush_size)
        if painted is not None and painted.any():
            return painted.astype(np.uint8), "paint"
        raise RuntimeError("no mask painted; use --first-mask or --seed-bbox")

    mask = np.zeros((height, width), dtype=np.uint8)
    if seed_bbox is not None:
        x, y, w, h = seed_bbox
        mask[y : y + h, x : x + w] = 1
        return mask, "bbox"

    raise RuntimeError("provide --first-mask, --pick-mask, or --seed-bbox")


def parse_seed_bbox(value: str | None) -> tuple[int, int, int, int] | None:
    if value is None:
        return None
    parts = [int(x.strip()) for x in value.split(",")]
    if len(parts) != 4:
        raise ValueError("--seed-bbox expects x,y,w,h")
    return parts[0], parts[1], parts[2], parts[3]


def resolve_settings(args: argparse.Namespace) -> dict:
    settings = dict(QUALITY_PRESETS[args.quality])
    if args.short_side is not None:
        settings["short_side"] = args.short_side
    if args.topk is not None:
        settings["topk"] = args.topk
    if args.neighborhood_size is not None:
        settings["neighborhood_size"] = args.neighborhood_size
    if args.no_refine_seed:
        settings["refine_seed"] = False
    if args.no_refine_edges:
        settings["refine_edges"] = False
    return settings


def build_parser() -> argparse.ArgumentParser:
    root = _repo_root()
    parser = argparse.ArgumentParser(description="LingBot-Vision training-free video segmentation demo.")
    parser.add_argument("--video", type=Path, default=root / "examples" / "tom-and-jerry-segment-p1.mp4")
    parser.add_argument("--model-dir", type=Path, default=root / "lingbot-vision-vit-giant")
    parser.add_argument("--out", type=Path, default=root / "outputs" / "vos_demo")
    parser.add_argument("--first-mask", type=Path, default=None, help="First-frame mask PNG (0=bg, 1..N=objects)")
    parser.add_argument("--pick-mask", action="store_true", help="Paint the first-frame mask with mouse brush")
    parser.add_argument("--brush-size", type=int, default=12, help="Brush radius in pixels for --pick-mask")
    parser.add_argument("--seed-bbox", type=str, default=None, help="Fallback bbox x,y,w,h on the first frame")
    parser.add_argument(
        "--quality",
        choices=sorted(QUALITY_PRESETS),
        default="balanced",
        help="fast=paper-like speed; balanced=default; high=finer pixel boundaries",
    )
    parser.add_argument("--short-side", type=int, default=None, help="Override inference resize (multiple of 16)")
    parser.add_argument("--topk", type=int, default=None, help="Top-k context patches for label propagation")
    parser.add_argument("--neighborhood-size", type=int, default=None, help="Local search radius in patch units")
    parser.add_argument("--no-refine-seed", action="store_true", help="Disable first-frame feature mask refinement")
    parser.add_argument("--no-refine-edges", action="store_true", help="Disable bilateral edge smoothing")
    parser.add_argument("--max-frames", type=int, default=0, help="Process only the first N frames (0 = all)")
    parser.add_argument("--overlay-alpha", type=float, default=0.45)
    return parser


def main() -> None:
    args = build_parser().parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required. Activate the project venv with CUDA PyTorch.")

    settings = resolve_settings(args)
    device = "cuda"
    dtype = torch.float16
    seed_bbox = parse_seed_bbox(args.seed_bbox)

    frames_rgb, fps = read_video(args.video)
    if args.max_frames > 0:
        frames_rgb = frames_rgb[: args.max_frames]
    num_frames = len(frames_rgb)
    mask_height, mask_width = frames_rgb[0].shape[:2]

    first_mask_np, mask_source = load_first_mask(
        args.first_mask,
        frames_rgb[0],
        args.pick_mask,
        seed_bbox,
        brush_size=args.brush_size,
    )
    num_masks = int(first_mask_np.max()) + 1
    print(
        f"[vos] video={args.video.name} frames={num_frames} size={mask_width}x{mask_height} "
        f"objects={num_masks - 1} quality={args.quality}"
    )
    print(
        f"[vos] short_side={settings['short_side']} upsample={settings['upsample']} "
        f"topk={settings['topk']} refine_seed={settings['refine_seed']} "
        f"refine_edges={settings['refine_edges']} feature_blend={settings['feature_blend']}"
    )

    args.out.mkdir(parents=True, exist_ok=True)

    with torch.inference_mode():
        backbone, _ = load_pretrained_backbone(
            repo_id_or_path=str(args.model_dir),
            variant="giant",
            device=device,
            dtype=dtype,
        )
        patch_size = backbone.patch_size
        print(f"[vos] backbone patch_size={patch_size} grad_enabled={torch.is_grad_enabled()}")

        if settings["refine_seed"] and mask_source == "bbox":
            refined = refine_seed_mask_with_features(
                frames_rgb[0],
                first_mask_np,
                backbone,
                device,
                dtype,
                settings["short_side"],
                patch_size,
            )
            if not np.array_equal(refined, first_mask_np):
                first_mask_np = refined
                print("[vos] refined first-frame seed mask with feature similarity")

        cv2.imwrite(
            str(args.out / "first_frame_mask.png"),
            cv2.cvtColor(mask_to_rgb(first_mask_np, num_masks), cv2.COLOR_RGB2BGR),
        )

        resized_first = resize_shortest_side(frames_rgb[0], settings["short_side"], patch_size)
        proc_h, proc_w = resized_first.shape[:2]
        feat_h, feat_w = proc_h // patch_size, proc_w // patch_size
        neighborhood_mask = make_neighborhood_mask(
            feat_h,
            feat_w,
            settings["neighborhood_size"],
            torch.device(device),
        )

        first_mask = torch.from_numpy(first_mask_np).to(device=device, dtype=torch.long)
        first_mask = F.interpolate(
            first_mask[None, None].float(),
            size=(feat_h, feat_w),
            mode="nearest",
        ).squeeze().long()
        first_probs = F.one_hot(first_mask, num_masks).float()

        img_norm = normalize_rgb(resized_first).to(device)
        first_feats = extract_features(backbone, img_norm, device, dtype)

        mask_predictions = np.zeros((num_frames, mask_height, mask_width), dtype=np.uint8)
        mask_predictions[0] = first_mask_np

        features_queue: list[torch.Tensor] = []
        probs_queue: list[torch.Tensor] = []

        writer = cv2.VideoWriter(
            str(args.out / f"{args.video.stem}_tracked.mp4"),
            cv2.VideoWriter_fourcc(*"mp4v"),
            fps,
            (mask_width * 2, mask_height),
        )

        preview = np.concatenate(
            [
                cv2.cvtColor(frames_rgb[0], cv2.COLOR_RGB2BGR),
                cv2.cvtColor(overlay_mask(frames_rgb[0], mask_to_rgb(first_mask_np, num_masks)), cv2.COLOR_RGB2BGR),
            ],
            axis=1,
        )
        writer.write(preview)

        start = time.perf_counter()
        for frame_idx in range(1, num_frames):
            resized = resize_shortest_side(frames_rgb[frame_idx], settings["short_side"], patch_size)
            current_feats = extract_features(backbone, normalize_rgb(resized).to(device), device, dtype)

            context_feats = torch.stack([first_feats, *features_queue], dim=0)
            context_probs = torch.stack([first_probs, *probs_queue], dim=0)
            current_probs = propagate(
                current_feats,
                context_feats,
                context_probs,
                neighborhood_mask,
                settings["topk"],
                settings["temperature"],
            )

            features_queue.append(current_feats)
            probs_queue.append(current_probs)
            if len(features_queue) > MAX_CONTEXT_LENGTH:
                features_queue.pop(0)
            if len(probs_queue) > MAX_CONTEXT_LENGTH:
                probs_queue.pop(0)

            probs_up = upsample_probs(
                current_probs,
                (proc_h, proc_w),
                settings["upsample"],
            )
            if settings["feature_blend"] > 0:
                probs_up = blend_with_feature_similarity(
                    probs_up,
                    current_feats,
                    first_feats,
                    first_probs,
                    (proc_h, proc_w),
                    settings["feature_blend"],
                )
            if settings["refine_edges"]:
                probs_up = refine_prob_edges(probs_up)

            if (proc_h, proc_w) != (mask_height, mask_width):
                probs_up = F.interpolate(
                    probs_up.unsqueeze(0),
                    size=(mask_height, mask_width),
                    mode=settings["upsample"],
                    align_corners=False if settings["upsample"] == "bilinear" else None,
                ).squeeze(0)
                probs_up = postprocess_probs(probs_up.unsqueeze(0)).squeeze(0)

            pred = torch.argmax(probs_up, dim=0).to(torch.uint8).cpu().numpy()
            mask_predictions[frame_idx] = pred

            mask_rgb = mask_to_rgb(pred, num_masks)
            panel = np.concatenate(
                [
                    cv2.cvtColor(frames_rgb[frame_idx], cv2.COLOR_RGB2BGR),
                    cv2.cvtColor(overlay_mask(frames_rgb[frame_idx], mask_rgb, args.overlay_alpha), cv2.COLOR_RGB2BGR),
                ],
                axis=1,
            )
            writer.write(panel)

            if frame_idx % 10 == 0 or frame_idx == num_frames - 1:
                elapsed = time.perf_counter() - start
                print(f"[vos] {frame_idx + 1}/{num_frames} elapsed={elapsed:.1f}s")
            torch.cuda.empty_cache()

        writer.release()
        np.save(args.out / f"{args.video.stem}_masks.npy", mask_predictions)

    elapsed = time.perf_counter() - start
    print(f"[vos] done in {elapsed:.1f}s")
    print(f"[vos] video -> {args.out / (args.video.stem + '_tracked.mp4')}")
    print(f"[vos] masks -> {args.out / (args.video.stem + '_masks.npy')}")
    print(f"[vos] peak_vram={torch.cuda.max_memory_allocated() / 1024**3:.2f} GiB")


if __name__ == "__main__":
    main()