Function main

tests/dump_phase7_reference.py:735–832

()

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

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

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

    tracker_weights = {
        k[len("tracker."):]: v
        for k, v in ckpt.items()
        if k.startswith("tracker.")
    }
    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.")
    }

    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:
        case_dir = os.path.join(args.outdir, case.case_id)
        os.makedirs(case_dir, exist_ok=True)
        write_meta(case_dir, case)

        save_tensor(os.path.join(case_dir, "input_curr_neck_trk_0"), neck_trk_0)
        save_tensor(os.path.join(case_dir, "input_curr_neck_trk_1"), neck_trk_1)
        save_tensor(os.path.join(case_dir, "input_curr_neck_trk_2"), neck_trk_2)

        mem_outputs = []
        obj_ptrs = []
        for idx, mem_case in enumerate(case.mem_cases):
            masks = load_ggml_masks(os.path.join(args.phase6_ref, mem_case, "sam_dec_masks"), 288, 288).float()
            low_res_mask = masks[:, :1, :, :].contiguous()
            sam_token = load_ggml_bd(os.path.join(args.phase6_ref, mem_case, "sam_dec_sam_token")).float()
            obj_score = load_ggml_bd(os.path.join(args.phase6_ref, mem_case, "sam_dec_obj_score")).float()

            save_tensor(os.path.join(case_dir, f"input_mem_mask_logits_{idx}"), low_res_mask)
            save_tensor(os.path.join(case_dir, f"input_mem_sam_token_{idx}"), sam_token)
            save_tensor(os.path.join(case_dir, f"input_mem_obj_score_{idx}"), obj_score)

            mem_out = run_memory_encoder(tracker_weights, neck_trk_2, low_res_mask, obj_score)