| |
| """For the flaky demos, try each prompt-variant grounded at several anchor frames, |
| report which (prompt, anchor) succeed. Goal: a robust grounding recipe.""" |
| import sys, os |
| from pathlib import Path |
| import numpy as np, cv2, 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 track_objects_in_video import load_video, sam3_propagate |
|
|
| DEMOS = ["tableHF_ep000", "tableHF_ep001", "tableHF_ep002", "tableHF_ep003", "tableHF_ep004"] |
| ROOT = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can" |
| CAN_P = ["red can", "can", "orange can"] |
| PLATE_P = ["White plate on the table", "white plate", "plate", "white round plate"] |
| ANCHORS = [0, 8, 16] |
| OUT = "/tmp/stackcan_robust"; os.makedirs(OUT, exist_ok=True) |
|
|
| bpe = str(Path(_mb.__file__).parent / "assets" / "bpe_simple_vocab_16e6.txt.gz") |
| pred = build_sam3_video_predictor(bpe_path=bpe, gpus_to_use=[0], apply_temporal_disambiguation=False) |
|
|
| def ground(frames, anchor, prompt): |
| clip = frames[anchor:anchor+16] |
| if len(clip) < 4: |
| return -1 |
| tmp = os.path.join(OUT, "clip.mp4"); imageio.mimwrite(tmp, clip, fps=15, macro_block_size=1) |
| m, n = sam3_propagate(pred, tmp, prompt) |
| if m is None: |
| return 0 |
| return int(np.asarray(m[0]).sum()) |
|
|
| for d in DEMOS: |
| vid = f"{ROOT}/{d}.mp4" |
| frames = [cv2.resize(f, (384, 384), interpolation=cv2.INTER_AREA) for f in load_video(vid, None)] |
| print(f"\n=== {d} ({len(frames)} frames) ===", flush=True) |
| for name, plist in [("CAN", CAN_P), ("PLATE", PLATE_P)]: |
| row = [] |
| for p in plist: |
| cells = [] |
| for a in ANCHORS: |
| px = ground(frames, a, p) |
| cells.append(f"a{a}={px}") |
| row.append(f"{p!r}: " + " ".join(cells)) |
| print(f" {name}:", flush=True) |
| for r in row: |
| print(" ", r, flush=True) |
| print("\ndone", flush=True) |
|
|
