validate_sam3_e2e.py

"""End-to-end validation of the full SAM3 ONNX pipeline.

Runs vision_encoder.onnx → text_encoder.onnx → decoder.onnx with a real seed
image and a "seed" text prompt, and prints the predicted boxes/masks/scores.

If this produces sensible output, we have a complete in-browser-ready SAM3.
"""

from pathlib import Path
import sys
import time
import numpy as np
import onnxruntime as ort
from PIL import Image
from transformers import AutoTokenizer
from transformers.models.sam3.image_processing_sam3 import Sam3ImageProcessor

OUTPUT_DIR = Path("sam3-onnx-test")
MODEL_ID = "facebook/sam3"

VISION_ONNX = OUTPUT_DIR / "vision_encoder.onnx"
TEXT_ONNX = OUTPUT_DIR / "text_encoder.onnx"
DECODER_ONNX = OUTPUT_DIR / "decoder.onnx"


def main() -> None:
    # Allow optional CLI args: <image_path> <text_prompt>
    image_path = sys.argv[1] if len(sys.argv) > 1 else None
    text_prompt = sys.argv[2] if len(sys.argv) > 2 else "seed"

    # --- Load preprocessors (from HF, NOT from the ONNX files) ----------------
    print("Loading preprocessors ...")
    image_processor = Sam3ImageProcessor.from_pretrained(MODEL_ID)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

    # --- Prep image -----------------------------------------------------------
    if image_path:
        print(f"Loading image: {image_path}")
        pil_img = Image.open(image_path).convert("RGB")
    else:
        print("No image given — using gray 640x480 placeholder")
        pil_img = Image.new("RGB", (640, 480), color=(128, 128, 128))
    pixel_values = image_processor(images=pil_img, return_tensors="np")["pixel_values"]
    print(f"  pixel_values: shape={pixel_values.shape} dtype={pixel_values.dtype}")

    # --- Prep text ------------------------------------------------------------
    print(f"Tokenizing prompt: {text_prompt!r}")
    # We hardcode max_len=32 since that's SAM3's config — keeps script standalone
    encoded = tokenizer(
        text_prompt,
        return_tensors="np",
        padding="max_length",
        max_length=32,
        truncation=True,
    )
    input_ids = encoded["input_ids"].astype(np.int64)
    attention_mask = encoded["attention_mask"].astype(np.int64)
    print(f"  input_ids: shape={input_ids.shape} dtype={input_ids.dtype}")

    # --- Build ONNX sessions --------------------------------------------------
    providers = ["CPUExecutionProvider"]
    print(f"\nLoading ONNX sessions on {providers} ...")
    t0 = time.time()
    vision_sess = ort.InferenceSession(str(VISION_ONNX), providers=providers)
    text_sess = ort.InferenceSession(str(TEXT_ONNX), providers=providers)
    decoder_sess = ort.InferenceSession(str(DECODER_ONNX), providers=providers)
    print(f"  Loaded in {time.time() - t0:.1f}s")

    # Print actual input/output names so we know what the exporter kept.
    # The legacy tracer drops unused inputs, so the decoder may not have all
    # the names we declared in the export script.
    print("\n  Decoder ONNX inputs:")
    for inp in decoder_sess.get_inputs():
        print(f"    {inp.name}: shape={inp.shape} dtype={inp.type}")
    print("  Decoder ONNX outputs:")
    for outp in decoder_sess.get_outputs():
        print(f"    {outp.name}: shape={outp.shape} dtype={outp.type}")

    # --- 1. Vision encoder ----------------------------------------------------
    print("\n[1/3] Running vision encoder ...")
    t0 = time.time()
    v_out = vision_sess.run(None, {"pixel_values": pixel_values})
    print(f"  Done in {time.time() - t0:.1f}s")
    # v_out is a list of 8 tensors in our defined order:
    # fpn_hidden_state_0..3, fpn_position_encoding_0..3
    fpn_h = v_out[0:4]
    fpn_p = v_out[4:8]
    for i, t in enumerate(fpn_h):
        print(f"  fpn_hidden_state_{i}: shape={t.shape}")

    # --- 2. Text encoder ------------------------------------------------------
    print("\n[2/3] Running text encoder ...")
    t0 = time.time()
    t_out = text_sess.run(
        None,
        {"input_ids": input_ids, "attention_mask": attention_mask},
    )
    text_features = t_out[0]
    print(f"  Done in {time.time() - t0:.1f}s")
    print(f"  text_features: shape={text_features.shape}")

    # --- 3. Decoder pipeline --------------------------------------------------
    # Decoder uses first 3 of 4 FPN levels (forward() does fpn_hidden_states[:-1])
    print("\n[3/3] Running decoder ...")
    t0 = time.time()
    # Build feed dict dynamically — the legacy tracer may have dropped unused inputs
    candidate_inputs = {
        "fpn_hidden_state_0": fpn_h[0],
        "fpn_hidden_state_1": fpn_h[1],
        "fpn_hidden_state_2": fpn_h[2],
        "fpn_position_encoding_0": fpn_p[0],
        "fpn_position_encoding_1": fpn_p[1],
        "fpn_position_encoding_2": fpn_p[2],
        "text_features": text_features,
        "attention_mask": attention_mask,
    }
    expected_input_names = {inp.name for inp in decoder_sess.get_inputs()}
    feed = {k: v for k, v in candidate_inputs.items() if k in expected_input_names}
    missing = expected_input_names - feed.keys()
    if missing:
        raise RuntimeError(f"Decoder expects inputs we didn't provide: {missing}")
    dropped = candidate_inputs.keys() - feed.keys()
    if dropped:
        print(f"  (note: tracer optimized out unused inputs: {sorted(dropped)})")
    d_out = decoder_sess.run(None, feed)
    pred_masks, pred_boxes, pred_logits = d_out
    print(f"  Done in {time.time() - t0:.1f}s")
    print(f"  pred_masks: shape={pred_masks.shape} dtype={pred_masks.dtype}")
    print(f"  pred_boxes: shape={pred_boxes.shape} dtype={pred_boxes.dtype}")
    print(f"  pred_logits: shape={pred_logits.shape} dtype={pred_logits.dtype}")

    # --- Inspect top detections -----------------------------------------------
    # Apply sigmoid to logits to get scores in [0, 1]
    scores = 1.0 / (1.0 + np.exp(-pred_logits))  # sigmoid
    scores_b0 = scores[0]
    top_k = 10
    top_idx = np.argsort(-scores_b0)[:top_k]

    print(f"\nTop {top_k} detections by score:")
    print(f"  {'idx':>4} {'score':>7} {'box (xyxy normalized)':>30}")
    for i in top_idx:
        x1, y1, x2, y2 = pred_boxes[0, i]
        print(f"  {i:>4} {scores_b0[i]:>7.4f}  ({x1:.3f}, {y1:.3f}, {x2:.3f}, {y2:.3f})")

    # How many detections above a reasonable threshold?
    for thresh in (0.5, 0.3, 0.1, 0.05):
        n_above = (scores_b0 > thresh).sum()
        print(f"  detections with score > {thresh}: {n_above}")

    print("\n✅ End-to-end ONNX pipeline ran without error.")
    print("   If scores look reasonable for your test image, we're shipping.")


if __name__ == "__main__":
    main()