()
| 130 | |
| 131 | |
| 132 | def main(): |
| 133 | parser = argparse.ArgumentParser() |
| 134 | parser.add_argument("--checkpoint", required=True) |
| 135 | parser.add_argument("--image", required=True) |
| 136 | parser.add_argument("--outdir", default="tests/ref") |
| 137 | args = parser.parse_args() |
| 138 | |
| 139 | os.makedirs(args.outdir, exist_ok=True) |
| 140 | device = "cpu" |
| 141 | |
| 142 | # ── Load checkpoint ───────────────────────────────────────────────── |
| 143 | print("Loading checkpoint...") |
| 144 | ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) |
| 145 | if "model" in ckpt: |
| 146 | ckpt = ckpt["model"] |
| 147 | |
| 148 | # Extract ViT weights |
| 149 | vit_prefix = "detector.backbone.vision_backbone.trunk." |
| 150 | vit_sd = {k[len(vit_prefix):]: v for k, v in ckpt.items() if k.startswith(vit_prefix)} |
| 151 | print(f" Found {len(vit_sd)} ViT keys") |
| 152 | |
| 153 | # ── Load and preprocess image ───────────────────────────────────── |
| 154 | print(f"Loading image: {args.image}") |
| 155 | img = Image.open(args.image).convert("RGB") |
| 156 | transform = v2.Compose([ |
| 157 | v2.ToDtype(torch.uint8, scale=True), |
| 158 | v2.Resize(size=(1008, 1008)), |
| 159 | v2.ToDtype(torch.float32, scale=True), |
| 160 | v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), |
| 161 | ]) |
| 162 | img_tensor = v2.functional.to_image(img) |
| 163 | img_preprocessed = transform(img_tensor).unsqueeze(0).to(device) |
| 164 | save_tensor(os.path.join(args.outdir, "preprocessed"), img_preprocessed) |
| 165 | |
| 166 | # ── Manual ViT forward pass ─────────────────────────────────────── |
| 167 | with torch.no_grad(): |
| 168 | # Config |
| 169 | E = 1024 |
| 170 | NH = 16 |
| 171 | HD = 64 |
| 172 | depth = 32 |
| 173 | MLP_DIM = 4736 |
| 174 | WS = 24 |
| 175 | global_blocks = {7, 15, 23, 31} |
| 176 | |
| 177 | # 1. Patch embedding: Conv2d(3, 1024, 14, 14, bias=False) |
| 178 | patch_w = vit_sd["patch_embed.proj.weight"] # [1024, 3, 14, 14] |
| 179 | x = F.conv2d(img_preprocessed, patch_w, stride=14) # [1, 1024, 72, 72] |
| 180 | x = x.permute(0, 2, 3, 1) # [1, 72, 72, 1024] |
| 181 | save_tensor(os.path.join(args.outdir, "patch_embed"), x) |
| 182 | |
| 183 | # 2. Positional embedding (tiled) |
| 184 | pos_embed = vit_sd["pos_embed"] # [1, 577, 1024] |
| 185 | pos = get_abs_pos(pos_embed, has_cls_token=True, hw=(72, 72), tiling=True) |
| 186 | save_tensor(os.path.join(args.outdir, "pos_embed_tiled"), pos) |
| 187 | x = x + pos |
| 188 | save_tensor(os.path.join(args.outdir, "after_pos_embed"), x) |
| 189 |
no test coverage detected