BioPredict/2_prototype_extractor.py at main · AllenSpike/BioPredict · 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
import os
import time
import argparse

import faiss
import joblib
import numpy as np
import pandas as pd
import torch
from tqdm import tqdm


def cluster(all_tile_features, n_proto, n_iter, device, n_init=5, n_proto_patches=50000, mode='faiss', use_cuda=False):
    """
    K-Means clustering on embedding space with single GPU
    """
    patches = torch.from_numpy(all_tile_features).to(device)
    n_patches = patches.shape[0]
    print(f"\nTotal of {n_patches} patches aggregated")

    s = time.time()
    if mode == 'faiss':
        assert use_cuda, f"FAISS requires access to GPU. Please enable use_cuda"

        numOfGPUs = torch.cuda.device_count()
        print(f"\nUsing Faiss Kmeans for clustering with {numOfGPUs} GPUs...")
        print(f"\tNum of clusters {n_proto}, num of iter {n_iter}")

        kmeans = faiss.Kmeans(patches.shape[1],
                              n_proto,
                              niter=n_iter,
                              nredo=n_init,
                              verbose=True,
                              max_points_per_centroid=n_proto_patches,
                              gpu=numOfGPUs)
        kmeans.train(patches.cpu().numpy())
        weight = torch.tensor(kmeans.centroids).to(device)

    else:
        raise NotImplementedError(f"Clustering not implemented for {mode}!")

    e = time.time()
    print(f"\nClustering took {e-s} seconds!")

    return n_patches, weight


def parse_args():
    parser = argparse.ArgumentParser(
        description="Step 2: learn prototypes by K-Means clustering (faiss, GPU) over all patch features."
    )
    parser.add_argument('--annotation', required=True,
                        help='Annotation CSV listing the slides to include in clustering.')
    parser.add_argument('--slide_col', default='slide',
                        help="Column in the annotation CSV with slide names (default: 'slide').")
    parser.add_argument('--feature_dir', required=True,
                        help='Directory with slideflow feature bags (.pt + .index.npz per slide), from step 1.')
    parser.add_argument('--output_dir', required=True,
                        help='Output directory; prototypes are saved as weights_<n_proto>.pkl.')
    parser.add_argument('--n_proto', type=int, default=16,
                        help='Number of prototypes / clusters (default: 16).')
    parser.add_argument('--n_iter', type=int, default=50,
                        help='Number of K-Means iterations (default: 50).')
    parser.add_argument('--n_init', type=int, default=5,
                        help='Number of K-Means restarts, nredo (default: 5).')
    parser.add_argument('--n_proto_patches', type=int, default=5000000,
                        help='faiss max_points_per_centroid (default: 5000000).')
    parser.add_argument('--gpu_id', default='0',
                        help="CUDA_VISIBLE_DEVICES setting (default: '0').")
    return parser.parse_args()


def main():
    args = parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # slide list: only training samples are used for prototype learning
    clinical = pd.read_csv(args.annotation)
    clinical = clinical[clinical['set'] == 'train']
    patient_list = clinical[args.slide_col].dropna().tolist()
    patient_list = [x for x in patient_list if x != 'NA']
    print(f"{len(patient_list)} training slides selected for prototype learning")

    # aggregate patch features of all slides
    pts = [os.path.join(args.feature_dir, str(name)) + '.pt' for name in patient_list]

    all_tile_features = []
    for i in tqdm(range(len(patient_list)), desc='Loading features'):
        all_tile_features.append(torch.load(pts[i])[:])

    all_tile_features = np.vstack(all_tile_features)

    # clustering
    _, weight = cluster(all_tile_features, n_proto=args.n_proto, n_iter=args.n_iter,
                        device=device, n_init=args.n_init, n_proto_patches=args.n_proto_patches,
                        mode='faiss', use_cuda=True)

    # save prototype weights
    weights = [weight]
    os.makedirs(args.output_dir, exist_ok=True)
    out_path = os.path.join(args.output_dir, f'weights_{args.n_proto}.pkl')
    joblib.dump(weights, out_path)
    print(f"Saved prototypes to {out_path}")


if __name__ == '__main__':
    main()