| |
| """Track cheez-it box + red-cross mark on all 100 success push-box demos (384x384), |
| saving combined point tracks WITH per-point object labels (0=box, 1=mark). |
| |
| Source: rise_data_384_pushbox_dualmask videos_small (384 RGB, WM-domain-matched). |
| Success set = episodes whose meta _src starts with 'tableHF' (the 100 teleop demos). |
| Loads SAM3 + CoTracker ONCE and loops (no per-video model reload). |
| |
| Per episode -> <OUT>/<src>/tracks_labeled.npz with: |
| tracks (T,N,2) float32, vis (T,N) bool, labels (N,) int [0=box,1=mark], |
| n_box, n_mark, episode_index, src, H, W |
| Also writes a combined overlay mp4 for the first --n_overlay episodes (spot-check). |
| """ |
| import sys, os, json |
| from pathlib import Path |
| import numpy as np |
| import cv2 |
| import torch |
| import imageio |
|
|
| sys.path.insert(0, "/home/hongyi/PWM/external") |
| import sam3.model_builder as _mb |
| from sam3.model_builder import build_sam3_video_predictor |
| from cotracker.predictor import CoTrackerPredictor |
| from track_objects_in_video import load_video, sam3_propagate, run_cotracker, COTRACKER_CKPT |
|
|
| EPISODES_JSONL = "/data0/hongyi/rise_data_384_pushbox_dualmask/Push-Box/meta/episodes.jsonl" |
| VS_ROOT = "/data0/hongyi/rise_data_384_pushbox_dualmask/Push-Box/videos_small" |
| CAM = "observation.images.top_head" |
| CHUNK_SIZE = 100 |
| OUT_ROOT = "/data0/hongyi/reward_model_data/3pointr_data_push_box_2d_boxmark" |
| PROMPTS = [("box", "cheez-it box"), ("mark", "red cross mark")] |
| STRIDE = 3 |
| N_OVERLAY = 5 |
|
|
|
|
| def find_video(ep_idx): |
| chunk = ep_idx // CHUNK_SIZE |
| return f"{VS_ROOT}/chunk-{chunk:03d}/{CAM}/episode_{ep_idx:06d}.mp4" |
|
|
|
|
| def track_one(predictor, cot, tmp_mp4, frames, prompt): |
| """Returns (tracks (T,N,2), vis (T,N)) or None if SAM3 finds nothing.""" |
| masks, n_obj = sam3_propagate(predictor, tmp_mp4, prompt) |
| if masks is None: |
| return None |
| masks_arr = np.stack([masks[i] for i in range(len(frames))], axis=0) |
| if masks_arr[0].any(): |
| init_idx = 0 |
| else: |
| nz = [i for i in range(len(frames)) if masks_arr[i].any()] |
| if not nz: |
| return None |
| init_idx = nz[0] |
| tr, vis = run_cotracker(cot, frames, masks_arr[init_idx], STRIDE, init_frame_idx=init_idx) |
| if tr is None: |
| return None |
| return tr, vis |
|
|
|
|
| def main(): |
| os.makedirs(OUT_ROOT, exist_ok=True) |
| |
| succ = [] |
| for line in open(EPISODES_JSONL): |
| d = json.loads(line) |
| if str(d.get("_src", "")).startswith("tableHF"): |
| succ.append((int(d["episode_index"]), d["_src"], int(d.get("length", 0)))) |
| succ.sort() |
| print(f"[build] {len(succ)} success episodes (tableHF)", flush=True) |
|
|
| bpe = str(Path(_mb.__file__).parent / "assets" / "bpe_simple_vocab_16e6.txt.gz") |
| print("[build] loading SAM3 ...", flush=True) |
| predictor = build_sam3_video_predictor(bpe_path=bpe, gpus_to_use=[0], |
| apply_temporal_disambiguation=False) |
| print("[build] loading CoTracker3 ...", flush=True) |
| cot = CoTrackerPredictor(checkpoint=COTRACKER_CKPT).cuda() |
|
|
| manifest = [] |
| for k, (ep_idx, src, length) in enumerate(succ): |
| vid = find_video(ep_idx) |
| if not os.path.exists(vid): |
| print(f"[skip] {src}: video missing {vid}", flush=True) |
| continue |
| outdir = os.path.join(OUT_ROOT, src) |
| os.makedirs(outdir, exist_ok=True) |
| frames = load_video(vid, None) |
| H, W = frames[0].shape[:2] |
| tmp_mp4 = os.path.join(outdir, "input.mp4") |
| imageio.mimwrite(tmp_mp4, frames, fps=15, macro_block_size=1) |
|
|
| res = {tag: track_one(predictor, cot, tmp_mp4, frames, prompt) for tag, prompt in PROMPTS} |
| if res["box"] is None: |
| print(f"[warn] {src}: BOX not detected — skipping", flush=True) |
| manifest.append({"src": src, "episode_index": ep_idx, "status": "no_box"}) |
| continue |
| tb, vb = res["box"] |
| T, nb = tb.shape[0], tb.shape[1] |
| if res["mark"] is not None: |
| tm, vm = res["mark"] |
| nm = tm.shape[1] |
| tracks = np.concatenate([tb, tm], axis=1).astype(np.float32) |
| vis = np.concatenate([vb, vm], axis=1) |
| labels = np.concatenate([np.zeros(nb, np.int64), np.ones(nm, np.int64)]) |
| mark_ok = True |
| else: |
| tracks, vis, labels, nm, mark_ok = tb.astype(np.float32), vb, np.zeros(nb, np.int64), 0, False |
| print(f"[warn] {src}: MARK not detected — saving box-only", flush=True) |
|
|
| np.savez(os.path.join(outdir, "tracks_labeled.npz"), |
| tracks=tracks, vis=vis, labels=labels, |
| n_box=nb, n_mark=nm, episode_index=ep_idx, src=src, H=H, W=W) |
|
|
| |
| bc0, bcl = tb[0].mean(0), tb[-1].mean(0) |
| if mark_ok: |
| mc = tm[0].mean(0) |
| d0, dl = float(np.linalg.norm(bc0 - mc)), float(np.linalg.norm(bcl - mc)) |
| else: |
| d0 = dl = -1.0 |
| manifest.append({"src": src, "episode_index": ep_idx, "status": "ok", |
| "n_box": nb, "n_mark": nm, "T": T, "mark_ok": mark_ok, |
| "box_mark_d0": round(d0, 1), "box_mark_dlast": round(dl, 1)}) |
| print(f"[{k+1}/{len(succ)}] {src}: box={nb} mark={nm} T={T} gap {d0:.0f}->{dl:.0f}px", flush=True) |
|
|
| if k < N_OVERLAY: |
| ov = os.path.join(outdir, "combined_overlay.mp4") |
| vw = imageio.get_writer(ov, fps=15, macro_block_size=1) |
| for t in range(T): |
| im = frames[t].copy() |
| for pi in range(nb): |
| x, y = tb[t, pi]; cv2.circle(im, (int(x), int(y)), 1, (0, 255, 0), -1) |
| if mark_ok: |
| for pi in range(nm): |
| x, y = tm[t, pi]; cv2.circle(im, (int(x), int(y)), 1, (255, 0, 0), -1) |
| vw.append_data(im) |
| vw.close() |
|
|
| with open(os.path.join(OUT_ROOT, "manifest.json"), "w") as f: |
| json.dump(manifest, f, indent=2) |
| ok = sum(1 for m in manifest if m.get("status") == "ok") |
| mk = sum(1 for m in manifest if m.get("mark_ok")) |
| print(f"[build] DONE: {ok} ok, {mk} with mark, {len(manifest)-ok} failed. -> {OUT_ROOT}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
