From d4eab03809ce2468491529fc8b3433b94efe9d95 Mon Sep 17 00:00:00 2001 From: kevin314 Date: Fri, 10 Jul 2026 07:13:01 +0000 Subject: [PATCH] Fix tracking --- data_pipeline/extract_tracks.py | 249 ++++++++++++++++++++++++++-- data_pipeline/run_extract_tracks.sh | 31 ++++ data_pipeline/run_segment_tracks.sh | 25 +++ data_pipeline/segment_tracks.py | 142 +++++++++++++++- data_pipeline/segment_viz.py | 29 ++-- 5 files changed, 446 insertions(+), 30 deletions(-) create mode 100755 data_pipeline/run_extract_tracks.sh create mode 100755 data_pipeline/run_segment_tracks.sh diff --git a/data_pipeline/extract_tracks.py b/data_pipeline/extract_tracks.py index e2f27dddd..2c4ba2d9a 100644 --- a/data_pipeline/extract_tracks.py +++ b/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= --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= --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, diff --git a/data_pipeline/run_extract_tracks.sh b/data_pipeline/run_extract_tracks.sh new file mode 100755 index 000000000..71a636cae --- /dev/null +++ b/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" diff --git a/data_pipeline/run_segment_tracks.sh b/data_pipeline/run_segment_tracks.sh new file mode 100755 index 000000000..808dffbd6 --- /dev/null +++ b/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" diff --git a/data_pipeline/segment_tracks.py b/data_pipeline/segment_tracks.py index f814b41ef..f20ba4295 100644 --- a/data_pipeline/segment_tracks.py +++ b/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: /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) diff --git a/data_pipeline/segment_viz.py b/data_pipeline/segment_viz.py index 16b886aac..92d5314b5 100644 --- a/data_pipeline/segment_viz.py +++ b/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