SAM 3 text-prompted segmentation for WebGPU: Standard (fp16) and Compact (8-bit weights)
7eee2d7 verified | """Runs SAM 3 in PyTorch on the test images, whole and as the three split graphs, checks they | |
| agree, and saves the inputs and the split graphs' outputs as references for the ONNX checks.""" | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from transformers import Sam3Model, Sam3Processor | |
| from sam3_graphs import Sam3Detector, Sam3Text, Sam3Vision, position_encoding_72 | |
| ROOT = Path(__file__).parent | |
| WEIGHTS = ROOT / "weights" | |
| CASES = { | |
| "car.png": ["car"], | |
| "shivan_dragon.jpg": ["dragon", "fire"], | |
| "serra_angel.jpg": ["angel", "wings"], | |
| "llanowar_elves.jpg": ["elf", "tree"], | |
| "lightning_bolt.jpg": ["lightning"], | |
| "island.jpg": ["sky", "water"], | |
| } | |
| def main() -> None: | |
| torch.set_grad_enabled(False) | |
| model = Sam3Model.from_pretrained(WEIGHTS, attn_implementation="eager").eval() | |
| processor = Sam3Processor.from_pretrained(WEIGHTS) | |
| vision, text = Sam3Vision(model).eval(), Sam3Text(model).eval() | |
| detector = Sam3Detector(model, position_encoding_72(model)).eval() | |
| out_dir = ROOT / "reference" | |
| out_dir.mkdir(exist_ok=True) | |
| for image_name, phrases in CASES.items(): | |
| image = Image.open(ROOT / "images" / image_name).convert("RGB") | |
| pixels = processor(images=image, return_tensors="pt").pixel_values | |
| start = time.time() | |
| levels = vision(pixels) | |
| print(f"{image_name}: vision {time.time() - start:.1f}s, levels {[tuple(l.shape) for l in levels]}") | |
| for phrase in phrases: | |
| tokens = processor(text=phrase, return_tensors="pt") | |
| ids, mask = tokens.input_ids, tokens.attention_mask | |
| features = text(ids, mask) | |
| split = detector(*levels, features, mask) | |
| whole = model(pixel_values=pixels, input_ids=ids, attention_mask=mask) | |
| pairs = { | |
| "pred_logits": (split[0], whole.pred_logits), | |
| "pred_boxes": (split[1], whole.pred_boxes), | |
| "presence_logits": (split[2], whole.presence_logits), | |
| "pred_masks": (split[3], whole.pred_masks), | |
| "semantic_seg": (split[4], whole.semantic_seg), | |
| } | |
| diffs = {k: float((a - b).abs().max()) for k, (a, b) in pairs.items()} | |
| scores = split[0].sigmoid() * split[2].sigmoid() | |
| kept = int((scores > 0.5).sum()) | |
| print(f" '{phrase}': ids {tuple(ids.shape)}, kept {kept}, max |split - whole| {diffs}") | |
| stem = f"{Path(image_name).stem}__{phrase}" | |
| np.savez_compressed( | |
| out_dir / f"{stem}.npz", | |
| pixel_values=pixels.numpy(), | |
| input_ids=ids.numpy(), | |
| attention_mask=mask.numpy(), | |
| features_288=levels[0].numpy(), | |
| features_144=levels[1].numpy(), | |
| features_72=levels[2].numpy(), | |
| text_features=features.numpy(), | |
| **{k: v[0].numpy() for k, v in pairs.items()}, | |
| ) | |
| sys.stdout.flush() | |
| if __name__ == "__main__": | |
| main() | |