File size: 3,057 Bytes
7eee2d7 | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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()
|