Function main

tests/dump_phase5_reference.py:146–774

()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--image", required=True)
    parser.add_argument("--outdir", default="tests/ref_phase5")
    parser.add_argument("--text", default="yellow school bus")
    args = parser.parse_args()

    os.makedirs(args.outdir, exist_ok=True)
    device = "cpu"
    D = 256
    NH = 8
    HD = D // NH
    NQ = 200
    T = 32  # text context length
    H = 72  # image feature spatial dim (1008 / 14)

    print("Loading checkpoint...")
    ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False)
    if "model" in ckpt:
        ckpt = ckpt["model"]

    # ── Load required weights from specific modules ─────────────────────

    # We need the image features & text features as inputs.
    # To get them, we first need to run the ViT backbone + neck + text encoder.
    # These were already verified in prior phases, so we load cached reference tensors.
    # But we need fresh tensors to match exact C++ inputs.

    # Let's first generate fresh image features using dump_vit_reference approach
    print(f"Loading image: {args.image}")
    img = Image.open(args.image).convert("RGB")
    transform = v2.Compose([
        v2.ToDtype(torch.uint8, scale=True),
        v2.Resize(size=(1008, 1008)),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])
    img_tensor = v2.functional.to_image(img)
    img_preprocessed = transform(img_tensor).unsqueeze(0).to(device)

    # ── ViT forward (reuse from dump_vit_reference.py) ─────────────────
    # ... (this is too long to inline, we'll load pre-computed neck features)
    # Instead, we'll load from the reference directory if available
    ref_dir = "tests/ref"
    neck_det_path = os.path.join(ref_dir, "neck_det_2")
    if os.path.exists(neck_det_path + ".bin"):
        print("Loading pre-computed neck features from tests/ref/")
        # Load neck features
        def load_tensor(path):
            shape = list(map(int, open(path + ".shape").read().strip().split(",")))
            data = np.fromfile(path + ".bin", dtype=np.float32).reshape(shape)
            return torch.tensor(data)

        # Neck det features: [1, 256, H, H] for scale 2 (72x72)
        neck_det_2 = load_tensor(os.path.join(ref_dir, "neck_det_2"))  # [1, 256, 72, 72]
        neck_det_1 = load_tensor(os.path.join(ref_dir, "neck_det_1"))  # [1, 256, 144, 144]
        neck_det_0 = load_tensor(os.path.join(ref_dir, "neck_det_0"))  # [1, 256, 288, 288]