KANResDiff/precomp_tr.py at main · PerceptionComputingLab/KANResDiff · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import os
import torch
from tqdm import tqdm
from models.pretrain.ProbUNet import ProbabilisticUNet, ProbUNetConfig
from dataloaders import get_dataloader_2
import argparse
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def set_seed(seed):
    import random
    import numpy as np
    import torch
    import os

    # --- Python & NumPy ---
    random.seed(seed)
    np.random.seed(seed)

    # --- PyTorch CPU/GPU ---
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # --- cuDNN ---
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # --- 控制哈希种子 ---
    os.environ["PYTHONHASHSEED"] = str(seed)

    print(f"[INFO] Global seed set to {seed}")

def init_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image_size", type=int, default=128)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--output_path", type=str, default="/root/dataset/LIDC_TR")
    parser.add_argument("--resume_filepath", type=str, default=None)
    parser.add_argument("--seed", type=int, default=4242)
    parser.add_argument("--num_samples", type=int, default=16)


    # ProbUNet-specific options.
    parser.add_argument("--z_dim", type=int, default=6)
    parser.add_argument("--beta", type=float, default=1.0)
    parser.add_argument("--y_mode", type=str, default="random1", choices=["random1", "stack4", "mean_soft"])
    parser.add_argument("--metric_mode", type=str, default="max_over_raters", choices=["max_over_raters", "mean_over_raters"])
    parser.add_argument("--resume_filepath_pre", type=str, default="/root/outputs/PreT_ProbUNet/2026-01-04_17-03-53/best.pth")
    parser.add_argument("--mode", type=str, default="train")


    return parser
if __name__ == "__main__":
    parser = init_parser()
    args = parser.parse_args()

    set_seed(args.seed)

    CKPT_PATH = args.resume_filepath_pre
    NSAMPLES = args.num_samples
    BATCH_SIZE = args.batch_size

    # Build ProbUNet.
    model_cfg = ProbUNetConfig(
        in_channels=1,
        out_channels=1,
        z_dim=6,
        beta=1.0,
        unet_base_channels=32,
        unet_depth=4,
        norm="gn",
        dropout=0.0,
        act="relu",
        up_mode="bilinear",
        latent_base_channels=32,
        latent_depth=4,
        fcomb_hidden=64,
        fcomb_layers=2,
    )

    pre_model = ProbabilisticUNet(model_cfg).to(device)
    ckpt = torch.load(CKPT_PATH, map_location=device)
    pre_model.load_state_dict(ckpt["model_state_dict"])
    pre_model.eval()

    # Build dataloader.
    loader, _ = get_dataloader_2(
        task="LIDC",
        split=args.mode,
        batch_size=BATCH_SIZE,
        splitratio=[0.8,0.015,0.185],
        randomsplit=True
    )

    all_TR = []
    all_ids = []

    with torch.no_grad():
        for images, masks, _, case_ids in tqdm(loader):
            images = images.to(device)

            tr_samples = pre_model.sample(images, n_samples=NSAMPLES)
            TR = tr_samples.mean(dim=0).clamp(0,1)  # (B,1,H,W)

            all_TR.append(TR.cpu())
            all_ids.extend(case_ids)

    all_TR = torch.cat(all_TR, dim=0)


    TR_SAVE_FILENAME = f"TR_{args.mode}_d0.8_n{NSAMPLES}.pt"
    TR_SAVE_PATH = os.path.join(args.output_path, TR_SAVE_FILENAME)
    os.makedirs(args.output_path, exist_ok=True)
    torch.save(
        {"ids": all_ids, "TR": all_TR},
        TR_SAVE_PATH
    )

    print(f"[OK] TR cache saved to {TR_SAVE_PATH}, shape={all_TR.shape}")