"""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 = 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()