Function main

tests/dump_phase6_reference.py:444–489

()

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--prephase-ref", default="tests/ref_phase3")
    parser.add_argument("--cases", default="tests/phase6_cases.tsv")
    parser.add_argument("--outdir", default="tests/ref_phase6")
    args = parser.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

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

    prompt_weights = {
        k[len("tracker.sam_prompt_encoder."):]: v
        for k, v in ckpt.items()
        if k.startswith("tracker.sam_prompt_encoder.")
    }
    mask_weights = {
        k[len("tracker.sam_mask_decoder."):]: v
        for k, v in ckpt.items()
        if k.startswith("tracker.sam_mask_decoder.")
    }
    no_mem_embed = ckpt["tracker.no_mem_embed"].float()

    neck_trk_0 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_0")).float()
    neck_trk_1 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_1")).float()
    neck_trk_2 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_2")).float()

    cases = load_cases(args.cases)
    for case in cases:
        print(f"Dumping {case.case_id}...")
        dump_case(
            case,
            prompt_weights,
            mask_weights,
            no_mem_embed,
            neck_trk_0,
            neck_trk_1,
            neck_trk_2,
            os.path.join(args.outdir, case.case_id),
        )

    print(f"All tensors saved to: {args.outdir}")


if __name__ == "__main__":