()
| 386 | |
| 387 | |
| 388 | def main(): |
| 389 | parser = argparse.ArgumentParser() |
| 390 | parser.add_argument("--checkpoint", required=True) |
| 391 | parser.add_argument("--prephase-ref", required=True, |
| 392 | help="Phase 3 reference dir with neck features") |
| 393 | parser.add_argument("--outdir", required=True) |
| 394 | args = parser.parse_args() |
| 395 | |
| 396 | D = 256 |
| 397 | H = 72 |
| 398 | |
| 399 | # Load checkpoint |
| 400 | print("Loading checkpoint...") |
| 401 | ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True) |
| 402 | |
| 403 | # Load backbone features from Phase 3 reference (saved in NCHW format) |
| 404 | print("Loading Phase 3 reference features...") |
| 405 | neck_det_2 = load_tensor(os.path.join(args.prephase_ref, "neck_det_2")) # [1, 256, 72, 72] |
| 406 | print(f" neck_det_2: {neck_det_2.shape}") |
| 407 | |
| 408 | # Image features in sequence-first format [H*W, B, C] |
| 409 | img_feats_hwc = neck_det_2.flatten(2).permute(2, 0, 1) # [H*W, B, C] |
| 410 | |
| 411 | # Image positional encoding [H*W, B, C] — sinusoidal PE computed on the fly |
| 412 | def compute_sine_pe(h, w, num_pos_feats=128, temperature=10000, scale=2*math.pi): |
| 413 | """Sinusoidal PE matching SAM3's PositionEmbeddingSine.""" |
| 414 | not_mask = torch.ones(1, h, w, dtype=torch.float32) |
| 415 | y_embed = not_mask.cumsum(1, dtype=torch.float32) |
| 416 | x_embed = not_mask.cumsum(2, dtype=torch.float32) |
| 417 | eps = 1e-6 |
| 418 | y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * scale |
| 419 | x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * scale |
| 420 | |
| 421 | dim_t = torch.arange(num_pos_feats, dtype=torch.float32) |
| 422 | dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats) |
| 423 | |
| 424 | pos_x = x_embed[:, :, :, None] / dim_t |
| 425 | pos_y = y_embed[:, :, :, None] / dim_t |
| 426 | pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) |
| 427 | pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) |
| 428 | pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) # [1, D, H, W] |
| 429 | return pos |
| 430 | |
| 431 | img_pe_nchw = compute_sine_pe(H, H, num_pos_feats=D // 2) # [1, 256, 72, 72] |
| 432 | img_pe_hwc = img_pe_nchw.flatten(2).permute(2, 0, 1) # [H*W, 1, C] |
| 433 | print(f" img_pe: {img_pe_hwc.shape}") |
| 434 | |
| 435 | # ══════════════════════════════════════════════════════════════════════ |
| 436 | # Test Case 1: Dummy prompt (no exemplars, just CLS) |
| 437 | # ══════════════════════════════════════════════════════════════════════ |
| 438 | dummy_boxes = torch.zeros(0, 1, 4) |
| 439 | dummy_labels = torch.zeros(0, 1, dtype=torch.long) |
| 440 | dummy_mask = torch.zeros(1, 0, dtype=torch.bool) |
| 441 | |
| 442 | dump_geometry_encoder( |
| 443 | ckpt, img_feats_hwc, neck_det_2, img_pe_hwc, |
| 444 | dummy_boxes, dummy_labels, dummy_mask, |
| 445 | args.outdir, "dummy_prompt" |
no test coverage detected