#!/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)