File size: 2,038 Bytes
23e1f91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | #!/usr/bin/env python3
"""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] # start frame of the 16-frame grounding clip
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)
|