| |
| """Probe NEW plate prompts on the can-only Stack-Can demos (plate previously missed). |
| For each (prompt), report how many of the probed demos it detects at ANY anchor.""" |
| 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 |
|
|
| ROOT = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can" |
| |
| DEMOS = ["ep002","ep007","ep010","ep011","ep012","ep018","ep019","ep020", |
| "ep031","ep042","ep061","ep073","ep080","ep091","ep092"] |
| PROMPTS = ["white round plate","round plate","white dish","dish","white disc", |
| "white circle","ceramic plate","dinner plate","white bowl","saucer", |
| "plate on the table","white object on the table","empty plate","white tray"] |
| ANCHORS = [0,16,32,48] |
| OUT="/tmp/stackcan_plate_v2"; 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) |
| return 0 if m is None else int(np.asarray(m[0]).sum()) |
|
|
| |
| frames_cache={} |
| for d in DEMOS: |
| fr=[cv2.resize(f,(384,384),interpolation=cv2.INTER_AREA) for f in load_video(f"{ROOT}/tableHF_{d}.mp4",None)] |
| frames_cache[d]=fr |
|
|
| hits={p:[] for p in PROMPTS} |
| for p in PROMPTS: |
| for d in DEMOS: |
| fr=frames_cache[d] |
| best=max((ground(fr,a,p) for a in ANCHORS if a<len(fr)), default=0) |
| if best>50: |
| hits[p].append(d) |
| print(f"{p!r:32s} -> {len(hits[p])}/{len(DEMOS)} demos {hits[p]}", flush=True) |
|
|
| print("\n=== ranking (most demos recovered) ===", flush=True) |
| for p,h in sorted(hits.items(), key=lambda kv:-len(kv[1])): |
| print(f" {len(h):2d}/{len(DEMOS)} {p!r}", flush=True) |
|
|
