| |
| """Probe which SAM3 text prompts detect the Stack-Can objects at 384x384. |
| Loads SAM3 once, squashes ep000 to 384, tries a list of prompts on a short clip, |
| prints #objects + #mask px, and dumps a frame-0 mask overlay png per prompt.""" |
| 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 |
|
|
| VID = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can/tableHF_ep000.mp4" |
| OUT = "/tmp/stackcan_probe"; os.makedirs(OUT, exist_ok=True) |
| PROMPTS = ["Green cup", "green cup", "cup", "mug", "White plate on the table", |
| "white plate", "plate", "red can", "orange can", "can", "soda can", |
| "tin can", "red object", "the object on the table"] |
|
|
| frames = [cv2.resize(f, (384, 384), interpolation=cv2.INTER_AREA) for f in load_video(VID, None)] |
| clip = frames[:16] |
| tmp = os.path.join(OUT, "clip.mp4"); imageio.mimwrite(tmp, clip, fps=15, macro_block_size=1) |
| imageio.imwrite(os.path.join(OUT, "frame0.png"), frames[0]) |
|
|
| 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) |
|
|
| for p in PROMPTS: |
| masks, n = sam3_propagate(pred, tmp, p) |
| if masks is None: |
| print(f" {p!r:40s} -> NONE", flush=True); continue |
| m0 = masks[0] |
| px = int(np.asarray(m0).sum()) |
| print(f" {p!r:40s} -> n_obj={n} frame0_px={px}", flush=True) |
| ov = frames[0].copy() |
| mm = np.asarray(m0).astype(bool) |
| if mm.ndim == 2 and mm.any(): |
| ov[mm] = (0.5 * ov[mm] + np.array([0, 0, 255]) * 0.5).astype(np.uint8) |
| imageio.imwrite(os.path.join(OUT, f"mask_{p.replace(' ','_')}.png"), ov) |
| print("done ->", OUT, flush=True) |
|
|
