File size: 1,916 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 | #!/usr/bin/env python3
"""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)
|