()
| 35 | |
| 36 | |
| 37 | def main(): |
| 38 | parser = argparse.ArgumentParser() |
| 39 | parser.add_argument("--checkpoint", required=True, help="Path to sam3.pt") |
| 40 | parser.add_argument("--image", required=True, help="Path to test image") |
| 41 | parser.add_argument("--outdir", default="tests/ref", help="Output directory") |
| 42 | args = parser.parse_args() |
| 43 | |
| 44 | os.makedirs(args.outdir, exist_ok=True) |
| 45 | device = "cpu" # Use CPU for reproducibility |
| 46 | |
| 47 | # ── Load model ───────────────────────────────────────────────────────── |
| 48 | print("Loading checkpoint...") |
| 49 | ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) |
| 50 | state_dict = ckpt if isinstance(ckpt, dict) and "model" not in ckpt else ckpt.get("model", ckpt) |
| 51 | |
| 52 | # Extract ViT weights |
| 53 | vit_prefix = "detector.backbone.visual.trunk." |
| 54 | vit_state = {k[len(vit_prefix):]: v for k, v in state_dict.items() if k.startswith(vit_prefix)} |
| 55 | |
| 56 | print(f" Found {len(vit_state)} ViT parameters") |
| 57 | |
| 58 | # ── Create ViT backbone ──────────────────────────────────────────────── |
| 59 | print("Creating ViT backbone...") |
| 60 | vit = _create_vit_backbone(compile_mode=None) |
| 61 | vit.use_act_checkpoint = False # Disable for inference |
| 62 | missing, unexpected = vit.load_state_dict(vit_state, strict=False) |
| 63 | if missing: |
| 64 | print(f" WARNING: Missing keys: {missing[:5]}...") |
| 65 | if unexpected: |
| 66 | print(f" WARNING: Unexpected keys: {unexpected[:5]}...") |
| 67 | vit.eval() |
| 68 | vit.to(device) |
| 69 | |
| 70 | # ── Load and preprocess image ────────────────────────────────────────── |
| 71 | print(f"Loading image: {args.image}") |
| 72 | img = Image.open(args.image).convert("RGB") |
| 73 | print(f" Original size: {img.size}") |
| 74 | |
| 75 | transform = v2.Compose([ |
| 76 | v2.ToDtype(torch.uint8, scale=True), |
| 77 | v2.Resize(size=(1008, 1008)), |
| 78 | v2.ToDtype(torch.float32, scale=True), |
| 79 | v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), |
| 80 | ]) |
| 81 | |
| 82 | img_tensor = v2.functional.to_image(img) |
| 83 | img_preprocessed = transform(img_tensor).unsqueeze(0).to(device) |
| 84 | print(f" Preprocessed shape: {list(img_preprocessed.shape)}") |
| 85 | save_tensor(os.path.join(args.outdir, "preprocessed"), img_preprocessed) |
| 86 | |
| 87 | # ── Step-by-step forward pass ────────────────────────────────────────── |
| 88 | with torch.no_grad(): |
| 89 | # 1. Patch embedding |
| 90 | x = vit.patch_embed(img_preprocessed) |
| 91 | print(f" After patch_embed: {list(x.shape)}") # [1, 72, 72, 1024] |
| 92 | save_tensor(os.path.join(args.outdir, "patch_embed"), x) |
| 93 | |
| 94 | # 2. Positional embedding |
no test coverage detected