| """Predict AI Worker arm masks with the native SAM3 image model.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
|
|
|
|
| CONFIG_PATH = Path(__file__).with_name("inference_config.json") |
|
|
|
|
| def _activation(activation, linear, values): |
| """Use the same unfused activation path as the fine-tuned checkpoint.""" |
| if activation in (torch.nn.GELU, torch.nn.functional.gelu): |
| return torch.nn.functional.gelu(linear(values)) |
| if activation in (torch.nn.ReLU, torch.nn.functional.relu): |
| return torch.nn.functional.relu(linear(values)) |
| raise ValueError(f"Unsupported activation: {activation}") |
|
|
|
|
| def load_predictor(checkpoint=None): |
| """Load the release checkpoint and its fixed 504px inference settings.""" |
| from sam3.model import vitdet |
| from sam3.model.sam3_image_processor import Sam3Processor |
| from sam3.model_builder import build_sam3_image_model |
|
|
| if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): |
| raise RuntimeError("A BF16-capable CUDA GPU is required. See README.md.") |
| config = json.loads(CONFIG_PATH.read_text(encoding="utf-8")) |
| if checkpoint is None: |
| checkpoint = hf_hub_download(config["model_id"], config["checkpoint"]) |
| checkpoint = Path(checkpoint) |
| with checkpoint.open("rb") as stream: |
| actual = hashlib.file_digest(stream, "sha256").hexdigest() |
| if actual != config["checkpoint_sha256"]: |
| raise ValueError("Checkpoint checksum mismatch. Download this release again.") |
|
|
| torch.set_num_threads(1) |
| vitdet.addmm_act = _activation |
| model = build_sam3_image_model( |
| checkpoint_path=str(checkpoint), load_from_HF=False, |
| device="cpu", eval_mode=True, enable_inst_interactivity=False, |
| ) |
| trunk = model.backbone.vision_backbone.trunk |
| grid = config["resolution"] // trunk.patch_embed.proj.kernel_size[0] |
| for block in trunk.blocks: |
| attention = block.attn |
| if block.window_size > 0 or attention.freqs_cis is None: |
| continue |
| attention.input_size = (grid, grid) |
| scale = attention.rope_pt_size[0] / grid if attention.rope_interp else 1.0 |
| frequencies = attention.compute_cis(end_x=grid, end_y=grid, scale_pos=scale) |
| attention.freqs_cis = frequencies |
| if attention.use_rope_real: |
| attention.freqs_cis_real = frequencies.real |
| attention.freqs_cis_imag = frequencies.imag |
| torch.set_num_threads(4) |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.manual_seed(438) |
| model = model.cuda().eval() |
| processor = Sam3Processor( |
| model, resolution=config["resolution"], device="cuda", |
| confidence_threshold=config["confidence_threshold"], |
| ) |
| return model, processor, config |
|
|
|
|
| @torch.inference_mode() |
| def predict_image(image, processor, config): |
| """Return native-resolution instance masks, XYXY boxes and scores.""" |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| state = processor.set_image(image.convert("RGB")) |
| result = processor.set_text_prompt(prompt=config["prompt"], state=state) |
| return ( |
| result["masks"][:, 0].cpu().numpy(), |
| result["boxes"].float().cpu().numpy(), |
| result["scores"].float().cpu().numpy(), |
| ) |
|
|
|
|
| def save_outputs(directory, masks, boxes, scores): |
| """Save a white foreground union mask and all individual predictions.""" |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| Image.fromarray(masks.any(axis=0).astype(np.uint8) * 255).save(directory / "mask.png") |
| np.savez_compressed(directory / "instances.npz", masks=masks, boxes_xyxy=boxes, scores=scores) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("image", type=Path, help="Input RGB image") |
| parser.add_argument("--checkpoint", type=Path, help="Local model.pt; otherwise download it") |
| parser.add_argument("--output-dir", type=Path, default=Path("prediction")) |
| args = parser.parse_args() |
| if args.image.resolve() in {(args.output_dir / name).resolve() for name in ["mask.png", "instances.npz"]}: |
| parser.error("The output directory would overwrite the input image.") |
| with Image.open(args.image) as source: |
| image = source.convert("RGB") |
| model, processor, config = load_predictor(args.checkpoint) |
| masks, boxes, scores = predict_image(image, processor, config) |
| save_outputs(args.output_dir, masks, boxes, scores) |
| print(f"Saved {len(scores)} predicted instances to {args.output_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
