Fix tracking

data_pipeline/extract_tracks.py

@@ -10,12 +10,17 @@ references ``action_path``).
 Tracks are stored in ORIGINAL video pixel coordinates. The full 2500-point grid + visibility
 are kept; the trainer samples 1-200 points per step.
 
+If ``--detect-entries`` is set, FastSAM detects objects entering after frame 0. For each
+entry event at frame T_entry a second CoTracker pass runs the full grid from T_entry. Grid
+points landing on the new object replace dead background slots (those permanently occluded
+from T_entry onwards), keeping N = grid_size^2 throughout.
+
 Run on a GPU node (never the login node), e.g.:
 
-    srun --jobid=<shao_wm jobid> --overlap --ntasks=1 \
-        .venv/bin/python data_pipeline/extract_tracks.py \
-        --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p \
-        --grid-size 50
+    srun --jobid=<shao_wm jobid> --overlap --ntasks=1 \\
+        .venv/bin/python data_pipeline/extract_tracks.py \\
+        --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p \\
+        --grid-size 50 --detect-entries
 
 torch.hub note: prefetch once on the login node (internet) so the shared cache is warm:
     .venv/bin/python -c "import torch; torch.hub.load('facebookresearch/co-tracker','cotracker3_offline')"
@@ -48,6 +53,24 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--limit", type=int, default=None, help="Only process the first N videos (for smoke tests).")
     p.add_argument("--index", type=int, nargs="+", default=None, metavar="IDX",
                    help="Only process videos with these indices (e.g. --index 4 7 12).")
+    p.add_argument("--rank", type=int, default=0, help="GPU rank for sharding (0-indexed).")
+    p.add_argument("--world-size", type=int, default=1, help="Total number of parallel processes.")
+    p.add_argument("--force", action="store_true", help="Re-extract even if .npz already exists.")
+    p.add_argument("--verbose", action="store_true", help="Print per-mask debug info for entry detection.")
+    # Entry-frame detection
+    p.add_argument("--detect-entries", action="store_true",
+                   help="Detect objects entering after frame 0 and replace dead background slots.")
+    p.add_argument("--sam-model", type=str, default="FastSAM-s.pt")
+    p.add_argument("--sam-conf", type=float, default=0.75)
+    p.add_argument("--sam-iou", type=float, default=0.9)
+    p.add_argument("--sam-imgsz", type=int, default=1024)
+    p.add_argument("--entry-sample-every", type=int, default=5,
+                   help="Check for new objects every N frames.")
+    p.add_argument("--entry-new-area", type=float, default=0.3,
+                   help="A mask triggers entry detection if at least this fraction of its area "
+                        "is not covered by any frame-0 mask (catches partially-entering objects).")
+    p.add_argument("--entry-min-area", type=float, default=0.005,
+                   help="Min area fraction for a new-object mask to be considered.")
     return p.parse_args()
 
 
@@ -70,16 +93,12 @@ def load_cotracker(model_name: str, device: str):
 
 
 def read_video(path: Path) -> tuple[torch.Tensor, int, int]:
-    """Return (video[1,T,C,H,W] float 0-255, H, W).
-
-    Recent torchvision dropped ``torchvision.io.read_video``; use decord (fast,
-    reliable), falling back to imageio's ffmpeg reader.
-    """
+    """Return (video[1,T,C,H,W] float 0-255, H, W)."""
     try:
         from decord import VideoReader, cpu
         vr = VideoReader(str(path), ctx=cpu(0))
         frames = vr.get_batch(list(range(len(vr)))).asnumpy()  # (T, H, W, C) uint8
-    except Exception:  # noqa: BLE001 - fall back to ffmpeg
+    except Exception:  # noqa: BLE001
         import imageio
         reader = imageio.get_reader(str(path), format="ffmpeg")
         frames = np.stack([np.asarray(f) for f in reader], axis=0)
@@ -103,15 +122,141 @@ def track_one(model, video: torch.Tensor, grid_size: int, downscale: float, devi
         sh, sw = h, w
 
     pred_tracks, pred_vis = model(track_video.to(device), grid_size=grid_size)
-    tracks = pred_tracks[0].float().cpu().numpy()      # (T, N, 2) in tracking-res px
-    vis = pred_vis[0].cpu().numpy()                    # (T, N)
+    tracks = pred_tracks[0].float().cpu().numpy()
+    vis = pred_vis[0].cpu().numpy()
 
-    if downscale != 1.0:  # rescale coords back to original resolution
+    if downscale != 1.0:
         tracks[..., 0] *= w / float(sw)
         tracks[..., 1] *= h / float(sh)
     return tracks.astype(np.float32), vis
 
 
+@torch.no_grad()
+def track_with_queries(model, video: torch.Tensor, queries_txy: np.ndarray, downscale: float,
+                       device: str, H: int, W: int) -> tuple[np.ndarray, np.ndarray]:
+    """Track explicit query points. queries_txy: [K,3] as (t, x, y) in original pixel coords."""
+    track_video = video
+    sh, sw = H, W
+    q = queries_txy.astype(np.float32).copy()
+    if downscale != 1.0:
+        sh = max(1, int(round(H * downscale)))
+        sw = max(1, int(round(W * downscale)))
+        track_video = F.interpolate(video[0], size=(sh, sw), mode="bilinear", align_corners=False).unsqueeze(0)
+        q[:, 1] *= sw / float(W)
+        q[:, 2] *= sh / float(H)
+
+    q_tensor = torch.from_numpy(q).unsqueeze(0).to(device)  # [1, K, 3]
+    pred_tracks, pred_vis = model(track_video.to(device), queries=q_tensor)
+    tracks = pred_tracks[0].float().cpu().numpy()
+    vis = pred_vis[0].cpu().numpy()
+
+    if downscale != 1.0:
+        tracks[..., 0] *= W / float(sw)
+        tracks[..., 1] *= H / float(sh)
+    return tracks.astype(np.float32), vis
+
+
+def make_grid_queries(grid_size: int, H: int, W: int, frame_t: int) -> np.ndarray:
+    """Generate a grid_size x grid_size uniform query grid at frame_t. Returns [N,3] (t,x,y)."""
+    ys = np.linspace(0, H - 1, grid_size)
+    xs = np.linspace(0, W - 1, grid_size)
+    xx, yy = np.meshgrid(xs, ys)
+    pts_xy = np.column_stack([xx.ravel(), yy.ravel()])  # [N, 2]
+    return np.column_stack([np.full(len(pts_xy), float(frame_t)), pts_xy]).astype(np.float32)
+
+
+
+def fastsam_masks(sam_model, frame_rgb: np.ndarray, conf: float, iou: float, imgsz: int,
+                  H: int, W: int) -> np.ndarray:
+    """Run FastSAM on a uint8 HxWx3 RGB frame, return bool masks [M,H,W]."""
+    res = sam_model(frame_rgb, device="cuda", retina_masks=True,
+                    imgsz=imgsz, conf=conf, iou=iou, verbose=False)
+    if not res or res[0].masks is None:
+        return np.zeros((0, H, W), bool)
+    masks = res[0].masks.data.cpu().numpy().astype(bool)
+    if masks.shape[0] and masks.shape[1:] != (H, W):
+        import cv2
+        masks = np.stack([
+            cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST).astype(bool)
+            for m in masks
+        ])
+    return masks
+
+
+def max_iou_with_set(mask: np.ndarray, others: np.ndarray) -> float:
+    """Max IOU of a single [H,W] bool mask against a set [M,H,W]."""
+    if others.shape[0] == 0:
+        return 0.0
+    inter = (mask & others).reshape(others.shape[0], -1).sum(1)
+    union = (mask | others).reshape(others.shape[0], -1).sum(1)
+    return float(np.where(union > 0, inter / union, 0.0).max())
+
+
+def detect_entry_events(video: torch.Tensor, sam_model, H: int, W: int, conf: float, iou: float,
+                        imgsz: int, sample_every: int, new_area_thresh: float,
+                        min_area_frac: float) -> tuple[dict[int, np.ndarray], np.ndarray]:
+    """Return ({frame_t: new_masks [M,H,W]}, masks0_union [H,W]) for frames where new
+    regions appear that weren't covered by any frame-0 mask.
+
+    Uses new-area fraction rather than IOU so partially-entering objects (partially visible
+    at frame 0, more visible later) are correctly detected: only the newly visible region
+    triggers detection and receives replacement tracks.
+    """
+    T = video.shape[1]
+    min_area_px = int(min_area_frac * H * W)
+
+    frame0 = video[0, 0].permute(1, 2, 0).numpy().astype(np.uint8)
+    masks0 = fastsam_masks(sam_model, frame0, conf, iou, imgsz, H, W)
+    # Exclude large background masks (table surface, floor, walls) from masks0_union.
+    # Only object-sized masks define "known territory" — background covers the whole frame
+    # and would suppress detection of the ball moving to a new position on that surface.
+    if masks0.shape[0] > 0:
+        areas = masks0.reshape(masks0.shape[0], -1).sum(1)
+        object_masks0 = masks0[areas < 0.4 * H * W]
+        masks0_union = object_masks0.any(axis=0) if object_masks0.shape[0] > 0 else np.zeros((H, W), bool)
+    else:
+        masks0_union = np.zeros((H, W), bool)
+
+    # Tracks newly-covered regions across entry events to avoid re-detecting the same
+    # entering object at multiple sample frames.
+    claimed_new = np.zeros((H, W), bool)
+
+    entry_events: dict[int, np.ndarray] = {}
+    for frame_t in range(sample_every, T, sample_every):
+        frame = video[0, frame_t].permute(1, 2, 0).numpy().astype(np.uint8)
+        masks_t = fastsam_masks(sam_model, frame, conf, iou, imgsz, H, W)
+        if masks_t.shape[0] == 0:
+            continue
+
+        new_masks = []
+        for m in masks_t:
+            area = int(m.sum())
+            if area < min_area_px:
+                continue
+            # New area: part of this mask not present in ANY frame-0 mask.
+            new_region = m & ~masks0_union
+            new_area_frac = float(new_region.sum()) / max(area, 1)
+            if new_area_frac < new_area_thresh:
+                continue
+            # Skip if this new region was already claimed by a prior entry event.
+            if float((new_region & claimed_new).sum()) / max(int(new_region.sum()), 1) > 0.5:
+                continue
+            new_masks.append(m)
+            # Dilate claimed region by ~2x the object radius so a fast-moving
+            # object doesn't re-trigger entry detection on subsequent sample frames.
+            radius = int(np.sqrt(int(new_region.sum()) / np.pi))
+            dil = max(1, radius * 2)
+            ys, xs = np.where(new_region)
+            y0, y1 = max(0, ys.min() - dil), min(H, ys.max() + dil + 1)
+            x0, x1 = max(0, xs.min() - dil), min(W, xs.max() + dil + 1)
+            claimed_new[y0:y1, x0:x1] = True
+
+        if new_masks:
+            entry_events[frame_t] = np.stack(new_masks)
+
+    return entry_events, masks0_union
+
+
 def patch_manifest(manifest_path: Path, stem_to_points: dict[str, Path]) -> int:
     if not manifest_path.exists():
         return 0
@@ -141,23 +286,91 @@ def main() -> None:
         videos = [v for v in videos if v.name in wanted]
     if args.limit is not None:
         videos = videos[:args.limit]
+    if args.world_size > 1:
+        videos = videos[args.rank::args.world_size]
     if not videos:
         print(f"[track] no videos found in {videos_dir}", flush=True)
         return
 
-    model = load_cotracker(args.model, args.device)
+    cotracker = load_cotracker(args.model, args.device)
     print(f"[track] loaded {args.model}; {len(videos)} videos, grid={args.grid_size}x{args.grid_size}", flush=True)
 
+    sam_model = None
+    if args.detect_entries:
+        from ultralytics import FastSAM
+        sam_model = FastSAM(args.sam_model)
+        print(f"[track] entry detection enabled: sam={args.sam_model}, "
+              f"conf={args.sam_conf}, every={args.entry_sample_every}f", flush=True)
+
     stem_to_points: dict[str, Path] = {}
     for k, vpath in enumerate(videos, 1):
         out_path = out_dir / f"{vpath.stem}.npz"
         stem_to_points[vpath.stem] = out_path
-        if out_path.exists():
+        if out_path.exists() and not args.force:
             continue
         video, h, w = read_video(vpath)
-        tracks, vis = track_one(model, video, args.grid_size, args.downscale, args.device)
-        # np.savez appends ".npz" unless the name already ends in it, so keep the
-        # tmp name ".npz"-terminated to make the atomic replace below match.
+        tracks, vis = track_one(cotracker, video, args.grid_size, args.downscale, args.device)
+
+        n_replaced = 0
+        if sam_model is not None:
+            entry_events, masks0_union = detect_entry_events(
+                video, sam_model, h, w,
+                conf=args.sam_conf, iou=args.sam_iou, imgsz=args.sam_imgsz,
+                sample_every=args.entry_sample_every,
+                new_area_thresh=args.entry_new_area,
+                min_area_frac=args.entry_min_area,
+            )
+            for frame_t, new_masks in sorted(entry_events.items()):
+                for mi, new_mask in enumerate(new_masks):
+                    new_region = new_mask & ~masks0_union
+
+                    # Covered slots: original frame-0 grid points whose position at
+                    # T_entry falls inside the new object's region.
+                    orig_xi = np.clip(tracks[frame_t, :, 0].round().astype(int), 0, w - 1)
+                    orig_yi = np.clip(tracks[frame_t, :, 1].round().astype(int), 0, h - 1)
+                    covered_dst = np.where(new_region[orig_yi, orig_xi])[0]
+                    dead_dst = covered_dst if covered_dst.size > 0 else \
+                        np.where((~vis[frame_t:].astype(bool)).all(axis=0))[0]
+                    if dead_dst.size == 0:
+                        continue
+
+                    # Run full grid from T_entry — same density as frame-0 grid.
+                    queries_txy = make_grid_queries(args.grid_size, h, w, frame_t)
+                    e_tracks, e_vis = track_with_queries(
+                        cotracker, video, queries_txy, args.downscale, args.device, h, w)
+                    e_vis[:frame_t] = False  # object didn't exist before entry frame
+
+                    pts = e_tracks[frame_t]
+                    xi = np.clip(pts[:, 0].round().astype(int), 0, w - 1)
+                    yi = np.clip(pts[:, 1].round().astype(int), 0, h - 1)
+                    object_src = np.where(new_region[yi, xi])[0]
+
+                    n = min(len(object_src), len(dead_dst))
+                    if n > 0:
+                        dst = dead_dst[:n]
+                        src = object_src[:n]
+                        tracks[:, dst] = e_tracks[:, src]
+                        vis[:, dst] = e_vis[:, src]
+                        # CoTracker can mark points invisible even at their query frame
+                        # when the object is entering from the edge. Force vis=True at
+                        # T_entry so segment_tracks sees these as first-visible there.
+                        vis[frame_t, dst] = True
+                        n_replaced += n
+                    if args.verbose:
+                        mask_area = int(new_mask.sum())
+                        new_region_area = int(new_region.sum())
+                        vis_at_entry = int(vis[frame_t, dst].sum()) if n > 0 else 0
+                        print(f"  [entry] t={frame_t} mask#{mi}: area={mask_area}px "
+                              f"new_region={new_region_area}px "
+                              f"object_src={len(object_src)} covered_dst={len(covered_dst)} "
+                              f"dead_dst={len(dead_dst)} -> replacing {n}, "
+                              f"vis[{frame_t}, dst].sum()={vis_at_entry}", flush=True)
+
+            frames_str = ",".join(str(t) for t in sorted(entry_events)) if entry_events else "none"
+            print(f"[track] [{k}/{len(videos)}] {vpath.name}: "
+                  f"{len(entry_events)} entry event(s) at frames [{frames_str}], "
+                  f"{n_replaced} slots replaced", flush=True)
+
         tmp = out_path.with_name(out_path.stem + ".tmp.npz")
         np.savez(
             tmp, tracks=tracks, visibility=vis,

data_pipeline/run_extract_tracks.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# Run extract_tracks.py in parallel across 4 GPUs.
+# Usage: bash data_pipeline/run_extract_tracks.sh [extra args]
+
+DATA_DIR=/home/hal-kevin/data/motion-stream-test
+WORLD_SIZE=4
+LOG_FILE=data_pipeline/extract_tracks.log
+
+> $LOG_FILE  # truncate on each run
+
+echo "[track] launching $WORLD_SIZE workers... logging to $LOG_FILE"
+
+for RANK in $(seq 0 $((WORLD_SIZE - 1))); do
+    CUDA_VISIBLE_DEVICES=$RANK python -u data_pipeline/extract_tracks.py \
+        --data-dir $DATA_DIR \
+        --grid-size 50 \
+        --device cuda \
+        --detect-entries \
+        --sam-conf 0.75 \
+        --sam-iou 0.9 \
+        --sam-imgsz 1024 \
+        --entry-sample-every 2 \
+        --entry-min-area 0.001 \
+        --entry-new-area 0.5 \
+        --rank $RANK --world-size $WORLD_SIZE \
+        "$@" \
+        >> $LOG_FILE 2>&1 &
+done
+
+wait
+echo "[track] all done. log at $LOG_FILE"

data_pipeline/run_segment_tracks.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Run segment_tracks.py in parallel across 4 GPUs.
+# Usage: bash data_pipeline/run_segment_tracks.sh [extra args]
+# Example: bash data_pipeline/run_segment_tracks.sh --limit 20
+
+DATA_DIR=/home/hal-kevin/data/motion-stream-test
+WORLD_SIZE=4
+LOG_FILE=data_pipeline/segment_tracks.log
+
+> $LOG_FILE  # truncate on each run
+
+echo "[seg] launching $WORLD_SIZE workers... logging to $LOG_FILE"
+
+for RANK in $(seq 0 $((WORLD_SIZE - 1))); do
+    CUDA_VISIBLE_DEVICES=$RANK python -u data_pipeline/segment_tracks.py \
+        --data-dir $DATA_DIR \
+        --conf 0.75 --iou 0.9 --imgsz 1024 \
+        --vis-override-every 1 --force --viz \
+        --rank $RANK --world-size $WORLD_SIZE \
+        "$@" \
+        >> $LOG_FILE 2>&1 &
+done
+
+wait
+echo "[seg] all done. log at $LOG_FILE"

data_pipeline/segment_tracks.py

@@ -23,6 +23,63 @@ from pathlib import Path
 import numpy as np
 
 
+def read_all_frames(path: str) -> np.ndarray:
+    """Read all frames from a video, return [T,H,W,3] uint8."""
+    try:
+        from decord import VideoReader, cpu
+        vr = VideoReader(path, ctx=cpu(0))
+        return vr.get_batch(list(range(len(vr)))).asnumpy()
+    except Exception:  # noqa: BLE001
+        import av
+        c = av.open(path)
+        return np.stack([f.to_ndarray(format="rgb24") for f in c.decode(video=0)])
+
+
+def _colors(n: int) -> np.ndarray:
+    import colorsys
+    return np.array(
+        [[int(255 * c) for c in colorsys.hsv_to_rgb((i * 0.61803) % 1.0, 0.65, 1.0)]
+         for i in range(max(1, n))],
+        np.uint8,
+    )
+
+
+def render_viz(frames: np.ndarray, tracks: np.ndarray, vis: np.ndarray,
+               object_ids: np.ndarray, out_path: Path, fps: int = 24) -> None:
+    import imageio.v2 as imageio
+    from fastvideo.train.callbacks.track_validation import _draw_overlay
+
+    objs = sorted(int(o) for o in np.unique(object_ids) if int(o) >= 0)
+    ocols = _colors(len(objs) + 1)
+    N = tracks.shape[1]
+    pcols = np.tile(np.array([[110, 110, 110]], np.uint8), (N, 1))
+    for oi, o in enumerate(objs):
+        pcols[object_ids == o] = ocols[oi % len(ocols)]
+
+    G = int(round(N**0.5))
+    if G * G == N:
+        k = max(1, G // 50)
+        sel = np.arange(N).reshape(G, G)[::k, ::k].reshape(-1)
+    else:
+        st = max(1, N // 1500)
+        sel = np.arange(0, N, st)
+
+    ov = _draw_overlay(frames, tracks[:, sel].copy(), vis[:, sel], pcols[sel], 12, 2, 0.5)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    imageio.mimsave(str(out_path), ov, fps=fps, macro_block_size=1)
+
+
+def render_seg_viz(frame: np.ndarray, masks: np.ndarray, out_path: Path) -> None:
+    import imageio.v2 as imageio
+    img = frame.copy()
+    if masks.shape[0] > 0:
+        colors = _colors(masks.shape[0])
+        for i, mask in enumerate(masks):
+            img[mask] = (img[mask] * 0.5 + colors[i % len(colors)] * 0.5).astype(np.uint8)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    imageio.imwrite(str(out_path), img)
+
+
 def read_frame(path: str, idx: int) -> np.ndarray:
     """Read a single frame by index, return HxWx3 uint8."""
     try:
@@ -138,16 +195,33 @@ def main() -> None:
                    help="drop masks smaller than this fraction of the frame")
     p.add_argument("--max-masks", type=int, default=0, help="keep only the N largest masks (0 = keep all)")
     p.add_argument("--limit", type=int, default=None)
+    p.add_argument("--index", type=int, nargs="+", default=None, metavar="IDX",
+                   help="Only process videos at these manifest indices (e.g. --index 4 7 12).")
+    p.add_argument("--rank", type=int, default=0, help="GPU rank for sharding (0-indexed).")
+    p.add_argument("--world-size", type=int, default=1, help="Total number of parallel processes.")
     p.add_argument("--force", action="store_true", help="re-run even if object_ids already present")
+    p.add_argument("--vis-override-every", type=int, default=0,
+                   help="Run FastSAM every N frames and set vis=True for object points inside masks. "
+                        "Fixes CoTracker vis=0 on edge-of-frame objects. 0 = disabled.")
+    p.add_argument("--viz", action="store_true", help="Render a track-overlay mp4 after each video.")
+    p.add_argument("--viz-dir", type=str, default=None,
+                   help="Output directory for viz mp4s (default: <data-dir>/viz).")
+    p.add_argument("--verbose", action="store_true", help="Print per-frame debug info.")
     args = p.parse_args()
+    if args.viz and args.viz_dir is None:
+        args.viz_dir = str(args.data_dir / "viz")
 
     from ultralytics import FastSAM
     model = FastSAM(args.model)
 
     manifest_path = args.data_dir / args.manifest
     items = json.loads(manifest_path.read_text()) if manifest_path.exists() else []
+    if args.index is not None:
+        items = [items[i] for i in args.index if i < len(items)]
     if args.limit:
         items = items[:args.limit]
+    if args.world_size > 1:
+        items = items[args.rank::args.world_size]
     n_ok = 0
     for k, item in enumerate(items, 1):
         vpath = args.data_dir / args.videos_subdir / item["path"]
@@ -169,17 +243,65 @@ def main() -> None:
         first_visible = np.where(ever_visible, np.argmax(vis, axis=0), -1)     # [N]
         unique_frames = sorted(set(first_visible[ever_visible].tolist()))
 
+        if args.verbose:
+            fv_counts = {f: int((first_visible == f).sum()) for f in unique_frames}
+            print(f"  [seg] first_visible frames: {fv_counts}", flush=True)
+
         # Run FastSAM on each unique first-visible frame
         frame_masks: dict[int, np.ndarray] = {}
         for frame_t in unique_frames:
             frame = read_frame(str(vpath), frame_t)
             res = model(frame, device=args.device, retina_masks=True,
                         imgsz=args.imgsz, conf=args.conf, iou=args.iou, verbose=False)
-            frame_masks[frame_t] = extract_masks(res, H, W, args.min_area_frac, args.max_masks)
+            masks = extract_masks(res, H, W, args.min_area_frac, args.max_masks)
+            frame_masks[frame_t] = masks
+            if args.verbose:
+                areas = masks.reshape(masks.shape[0], -1).sum(1).tolist() if masks.shape[0] else []
+                print(f"  [seg] frame {frame_t}: {masks.shape[0]} masks, areas={[int(a) for a in areas]}", flush=True)
 
         oid = assign_object_ids_multiframe(frame_masks, first_visible, tracks, H, W)
         n_objects = int(np.unique(oid[oid >= 0]).shape[0]) if (oid >= 0).any() else 0
 
+        if args.verbose:
+            for frame_t in unique_frames:
+                pt_sel = first_visible == frame_t
+                labeled = int((oid[pt_sel] >= 0).sum())
+                print(f"  [seg] frame {frame_t}: {pt_sel.sum()} pts, {labeled} got oid>=0", flush=True)
+
+        # Vis override: run FastSAM every N frames and set vis=True for object points
+        # that fall inside any mask — fixes CoTracker vis=0 on edge-of-frame objects.
+        n_overrides = 0
+        if args.vis_override_every > 0 and (oid >= 0).any():
+            T = tracks.shape[0]
+            object_pts = np.where(oid >= 0)[0]
+            for frame_t in range(0, T, args.vis_override_every):
+                frame = read_frame(str(vpath), frame_t)
+                res = model(frame, device=args.device, retina_masks=True,
+                            imgsz=args.imgsz, conf=args.conf, iou=args.iou, verbose=False)
+                masks = extract_masks(res, H, W, args.min_area_frac, args.max_masks)
+                if masks.shape[0] == 0:
+                    continue
+                pts = tracks[frame_t, object_pts]                              # [K,2]
+                xi = np.clip(pts[:, 0].round().astype(int), 0, W - 1)
+                yi = np.clip(pts[:, 1].round().astype(int), 0, H - 1)
+                in_any_mask = masks[:, yi, xi].any(axis=0)                    # [K] bool
+                newly_visible = in_any_mask & ~vis[frame_t, object_pts]
+                vis[frame_t, object_pts] |= in_any_mask
+                n_overrides += int(newly_visible.sum())
+
+            # Fill gaps between True frames caused by the sampling interval.
+            # If vis is True at frame T and True again at T+k (k <= override_every),
+            # the frames in between should also be True — the object didn't disappear.
+            for pt_idx in object_pts:
+                last_true = -1
+                for t in range(T):
+                    if vis[t, pt_idx]:
+                        if last_true >= 0 and (t - last_true) <= args.vis_override_every:
+                            vis[last_true:t, pt_idx] = True
+                        last_true = t
+
+            d["visibility"] = vis
+
         d["object_ids"] = oid.astype(np.int64)
         d["n_objects"] = np.int64(n_objects)
         d["track_weights"] = lowrank_track_weights(tracks)
@@ -187,10 +309,26 @@ def main() -> None:
         np.savez(tmp, **d)
         tmp.replace(npz_path)
         n_ok += 1
+        if args.viz:
+            frames = read_all_frames(str(vpath))[:tracks.shape[0]]
+            T_v = frames.shape[0]
+            stem_dir = Path(args.viz_dir) / vpath.stem
+            stem_dir.mkdir(parents=True, exist_ok=True)
+            render_viz(frames, tracks, vis, oid, stem_dir / "tracks.mp4",
+                       fps=int(item.get("fps", 24)))
+            for label, fidx in [("000", 0), ("mid", T_v // 2), ("last", T_v - 1)]:
+                frame = frames[fidx]
+                res = model(frame, device=args.device, retina_masks=True,
+                            imgsz=args.imgsz, conf=args.conf, iou=args.iou, verbose=False)
+                fmasks = extract_masks(res, H, W, args.min_area_frac, args.max_masks)
+                render_seg_viz(frame, fmasks, stem_dir / f"seg_frame{label}.jpg")
+            print(f"  viz -> {stem_dir}/", flush=True)
+
         cov = int((oid >= 0).sum())
+        override_str = f", {n_overrides} vis overrides" if args.vis_override_every > 0 else ""
         print(
             f"[seg] [{k}/{len(items)}] {vpath.name}: {n_objects} objs across {len(unique_frames)} frames, "
-            f"{cov}/{oid.shape[0]} grid pts labeled, "
+            f"{cov}/{oid.shape[0]} grid pts labeled{override_str}, "
             f"w[mean={d['track_weights'].mean():.3f} >0.5={(d['track_weights'] > 0.5).mean():.2f}]",
             flush=True)
     print(f"[seg] done; {n_ok}/{len(items)} npz have object_ids", flush=True)

data_pipeline/segment_viz.py

@@ -156,7 +156,7 @@ def cmd_render(args: argparse.Namespace) -> None:
     from PIL import Image, ImageDraw
     from ultralytics import FastSAM
     from fastvideo.train.callbacks.track_validation import _draw_overlay
-    from segment_tracks import object_ids_for_points
+    from segment_tracks import assign_object_ids_multiframe, extract_masks
     import imageio.v2 as imageio
 
     configs = _load_configs(args)
@@ -181,16 +181,25 @@ def cmd_render(args: argparse.Namespace) -> None:
         vis = d["visibility"].astype(np.float32)[:frames.shape[0]]  # [T,N]
         disp = np.sqrt(((tracks - tracks[0:1])**2).sum(-1)).max(0)  # [N] max displacement
 
+        ever_visible = vis.astype(bool).any(axis=0)                                      # [N]
+        first_visible = np.where(ever_visible, np.argmax(vis.astype(bool), axis=0), -1)  # [N]
+        unique_frames = sorted(set(first_visible[ever_visible].tolist()))
+
         for cfg in configs:
-            res = model(frames[0],
-                        device=args.device,
-                        retina_masks=True,
-                        imgsz=cfg["imgsz"],
-                        conf=cfg["conf"],
-                        iou=cfg["iou"],
-                        verbose=False)
-            masks = filter_masks(_extract_masks(res, H, W), cfg["min_area_frac"], cfg["max_masks"])
-            oid = object_ids_for_points(masks, tracks[0], H, W)  # [N] mask idx per point (-1 none)
+            # Run FastSAM on each unique first-visible frame for this config.
+            frame_masks: dict[int, np.ndarray] = {}
+            for frame_t in unique_frames:
+                res = model(frames[frame_t],
+                            device=args.device,
+                            retina_masks=True,
+                            imgsz=cfg["imgsz"],
+                            conf=cfg["conf"],
+                            iou=cfg["iou"],
+                            verbose=False)
+                frame_masks[frame_t] = extract_masks(res, H, W, cfg["min_area_frac"], cfg["max_masks"])
+
+            oid = assign_object_ids_multiframe(frame_masks, first_visible, tracks, H, W)
+            masks = frame_masks.get(0, np.zeros((0, H, W), bool))  # frame-0 masks for SAM panel
             objs = sorted(int(o) for o in np.unique(oid) if int(o) >= 0)
 
             # (1) SAM panel: colored masks + chosen point (highest-motion track) per object