# examples/sam_export.py """Export two licensed real-photo cases with the authenticated official SAM ViT-B. Extract sam_native.zip beside sam.py. Download the checkpoint URL reported by sam.CHECKPOINT_URL; its complete SHA256 is authenticated before torch.load. python sam_export.py --checkpoint sam_vit_b_01ec64.pth --output sam-demo By default the two pinned CC0 photographs are fetched and verified. --photo-dir can instead point to a directory with source-cat.png and source-coffee.png. """ import argparse import json import platform import shutil import sys import time from pathlib import Path from urllib.request import urlopen import numpy as np import PIL import sam import torch from PIL import Image, ImageDraw PHOTOS = [ { "id": "cat", "title": "Chelsea the cat", "point": [240.0, 160.0], "author": "Stefan van der Walt", "filename": "chelsea.png", "sha256": "596aa1e7cb875eb79f437e310381d26b338a81c2da23439704a73c4651e8c4bb", }, { "id": "coffee", "title": "Cup of coffee", "point": [285.0, 215.0], "author": "Rachel Michetti", "filename": "coffee.png", "sha256": "cc02f8ca188b167c775a7101b5d767d1e71792cf762c33d6fa15a4599b5a8de7", }, ] PHOTO_ROOT = "https://raw.githubusercontent.com/scikit-image/scikit-image/v0.25.2/skimage/data/" PHOTO_LICENSE = "https://github.com/scikit-image/scikit-image/blob/v0.25.2/skimage/data/_fetchers.py" def export(checkpoint, output, photo_dir=None): """Encode each photograph once, save all masks, and compare with native decode.""" torch.set_num_threads(2) torch.manual_seed(7) output.mkdir(parents=True, exist_ok=True) model = sam.build_model(checkpoint) results = [] for photo in PHOTOS: image_path = output / f"source-{photo['id']}.png" if photo_dir is not None: source_path = photo_dir / image_path.name if source_path.resolve() != image_path.resolve(): shutil.copyfile(source_path, image_path) elif not image_path.exists(): with urlopen( PHOTO_ROOT + photo["filename"], timeout=60 ) as response: image_path.write_bytes(response.read()) if sam.sha256(image_path) != photo["sha256"]: raise ValueError( f"Pinned CC0 photograph hash mismatch: {photo['id']}" ) original = Image.open(image_path).convert("RGB") started = time.perf_counter() cached = sam.encode_image(model, original) encoder_seconds = time.perf_counter() - started points = torch.tensor([[photo["point"]]]) labels = torch.ones(1, 1, dtype=torch.long) started = time.perf_counter() result = sam.predict(model, cached, points, labels) decoder_seconds = time.perf_counter() - started with torch.no_grad(): transformed = sam.transform_coords( points, cached["original_size"], cached["input_size"] ) sparse, dense = model.prompt_encoder( (transformed, labels), None, None ) native_logits, native_scores = model.mask_decoder( cached["embedding"], model.prompt_encoder.get_dense_pe(), sparse, dense, True, ) native_full = model.postprocess_masks( native_logits, cached["input_size"], cached["original_size"] ) differences = [ float((native_logits - result["low_res_logits"]).abs().max()), float((native_scores - result["scores"]).abs().max()), float((native_full - result["full_logits"]).abs().max()), ] if max(differences) > 1e-5: raise ValueError(f"Native/wrapper parity failed: {differences}") # A second interaction actually reuses the same embedding. second = sam.predict( model, cached, points, labels, mask_logits=result["low_res_logits"][ :, result["best"][0] : result["best"][0] + 1 ], multimask=False, ) masks = [] for index, mask in enumerate(result["masks"][0]): array = mask.numpy() name = f"{photo['id']}-mask-{index}.png" Image.fromarray(array.astype(np.uint8) * 255).save(output / name) masks.append( { "index": index, "file": name, "sha256": sam.sha256(output / name), "areaPixels": int(array.sum()), "predictedIoU": float(result["scores"][0, index]), } ) best = int(result["best"][0]) selected = result["masks"][0, best].numpy() overlay = np.array(original).copy() overlay[selected] = ( 0.55 * overlay[selected] + 0.45 * np.array([20, 184, 166]) ).astype(np.uint8) overlay_image = Image.fromarray(overlay) draw = ImageDraw.Draw(overlay_image) x, y = photo["point"] draw.ellipse( (x - 6, y - 6, x + 6, y + 6), fill="white", outline="black", width=2 ) overlay_name = f"{photo['id']}-overlay.png" overlay_image.save(output / overlay_name) arrays_name = f"{photo['id']}-outputs.npz" np.savez_compressed( output / arrays_name, low_res_logits=result["low_res_logits"].numpy(), scores=result["scores"].numpy(), masks=result["masks"].numpy(), point_coords=points.numpy(), point_labels=labels.numpy(), refinement_logits=second["low_res_logits"].numpy(), ) results.append( { "id": photo["id"], "title": photo["title"], "originalImage": image_path.name, "overlay": overlay_name, "overlaySha256": sam.sha256(output / overlay_name), "originalSize": list(cached["original_size"]), "resizedSize": list(cached["input_size"]), "embeddingShape": list(cached["embedding"].shape), "maskShape": list(result["masks"].shape), "pointCoords": photo["point"], "pointLabels": [1], "bestIndex": best, "masks": masks, "arrays": arrays_name, "arraysSha256": sam.sha256(output / arrays_name), "nativeMaxAbsErrors": differences, "encoderSeconds": encoder_seconds, "decoderSeconds": decoder_seconds, "reusedEmbedding": True, "refinementPredictedIoU": float(second["scores"][0, 0]), "source": { "url": PHOTO_ROOT + photo["filename"], "sha256": photo["sha256"], "author": photo["author"], "license": "CC0", "licenseUrl": PHOTO_LICENSE, }, } ) manifest = json.loads( ( Path(sam.__file__).parent / "sam_native/source-manifest.json" ).read_text() ) payload = { "provenance": { "architecture": "Original SAM ViT-B/16: 12 blocks, encoder width 768, prompt width 256", "parameterCount": sum( parameter.numel() for parameter in model.parameters() ), "sourceRevision": sam.REVISION, "nativeSources": manifest, "checkpointUrl": sam.CHECKPOINT_URL, "checkpointSha256": sam.sha256(checkpoint), "strictStateLoad": True, "missingKeys": [], "unexpectedKeys": [], "programSha256": sam.sha256(sam.__file__), "exportSha256": sam.sha256(__file__), "command": " ".join(sys.argv), "reproductionCommand": "python sam_export.py --checkpoint sam_vit_b_01ec64.pth --output sam-demo", "python": platform.python_version(), "torch": torch.__version__, "numpy": np.__version__, "pillow": PIL.__version__, "device": "CPU", "threads": torch.get_num_threads(), "seed": 7, "adapter": "Namespace-only import bypasses optional package APIs; native sources unchanged", "scope": "Two real-photo point prompts; not original benchmark or training reproduction", "scoreMeaning": "Predicted IoU regression for ranking; not calibrated probability", }, "preprocessing": { "color": "RGB", "pixelRange": "0..255 float32", "targetLongestSide": 1024, "resize": "Pillow BILINEAR, official half-up dimensions", "mean": [123.675, 116.28, 103.53], "std": [58.395, 57.12, 57.375], "padding": "Right and bottom zeros after normalization to 1024 square", "coordinates": "Original (x,y) scaled by independently rounded width/height, then +0.5", "postprocess": "Bilinear logits to 1024 square, crop resized H/W, bilinear to original H/W", "alignCorners": False, "maskLogitThreshold": 0.0, }, "photos": results, } (output / "predictions.json").write_text( json.dumps(payload, indent=2) + "\n" ) return payload def main(): """Run the explicitly requested pretrained export.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--output", type=Path, default=Path("sam-demo")) parser.add_argument("--photo-dir", type=Path) args = parser.parse_args() result = export(args.checkpoint, args.output, args.photo_dir) print( json.dumps( { "photos": len(result["photos"]), "nativeMaxAbsErrors": [ photo["nativeMaxAbsErrors"] for photo in result["photos"] ], } ) ) if __name__ == "__main__": main()