forked from sangjoon-park/Medical_X-VL
-
Notifications
You must be signed in to change notification settings - Fork 0
/
zero_shot_ensemble.py
163 lines (132 loc) · 6.46 KB
/
zero_shot_ensemble.py
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import os
import numpy as np
import pandas as pd
from pathlib import Path
from typing import List, Tuple, Optional
import sys
sys.path.append('../../')
from eval_zero import evaluate, bootstrap
from zero_shot import make, make_true_labels, run_softmax_eval
# ----- DIRECTORIES ------ #
cxr_filepath: str = '/home/depecher/PycharmProjects/CheXzero/data/chexpert_test.h5' # filepath of chest x-ray images (.h5)
cxr_true_labels_path: Optional[
str] = '/home/depecher/PycharmProjects/CheXzero/data/groundtruth.csv' # (optional for evaluation) if labels are provided, provide path
model_dir: str = '/COVID_8TB/sangjoon/chexzero_checkpoint/20230402_proposed/best_4/' # where pretrained models are saved (.pt)
predictions_dir: Path = Path('/home/depecher/PycharmProjects/CheXzero/predictions') # where to save predictions
cache_dir: str = predictions_dir / "cached" # where to cache ensembled predictions
context_length: int = 120
# ------- LABELS ------ #
# Define labels to query each image | will return a prediction for each label
# cxr_labels: List[str] = ['Atelectasis', 'Cardiomegaly',
# 'Consolidation', 'Edema', 'Enlarged Cardiomediastinum', 'Fracture', 'Lung Lesion',
# 'Lung Opacity', 'No Finding', 'Pleural Effusion', 'Pleural Other', 'Pneumonia',
# 'Pneumothorax', 'Support Devices']
cxr_labels: List[str] = ['Atelectasis', 'Cardiomegaly', 'Edema', 'Consolidation', 'Fracture', 'Pleural Effusion', 'Pneumonia', 'Pneumothorax']
# ---- TEMPLATES ----- #
# Define set of templates | see Figure 1 for more details
cxr_pair_template: Tuple[str] = ('{}', 'no {}')
# pos_query = ['Findings consistent with pneumonia', 'Findings suggesting pneumonia',
# 'This opacity can represent pneumonia', 'Findings are most compatible with pneumonia']
# neg_query = ['There is no pneumonia', 'No evidence of pneumonia',
# 'No evidence of acute pneumonia', 'No signs of pneumonia']
# ----- MODEL PATHS ------ #
# If using ensemble, collect all model paths
model_paths = []
for subdir, dirs, files in os.walk(model_dir):
for file in files:
full_dir = os.path.join(subdir, file)
model_paths.append(full_dir)
print(model_paths)
# %% md
## Run Inference
# %%
## Run the model on the data set using ensembled models
def ensemble_models(
model_paths: List[str],
cxr_filepath: str,
cxr_labels: List[str],
cxr_pair_template: Tuple[str],
cache_dir: str = None,
save_name: str = None,
) -> Tuple[List[np.ndarray], np.ndarray]:
"""
Given a list of `model_paths`, ensemble model and return
predictions. Caches predictions at `cache_dir` if location provided.
Returns a list of each model's predictions and the averaged
set of predictions.
"""
predictions = []
model_paths = sorted(model_paths) # ensure consistency of
for path in model_paths: # for each model
model_name = Path(path).stem
# load in model and `torch.DataLoader`
model, loader = make(
model_path=path,
cxr_filepath=cxr_filepath,
)
# path to the cached prediction
if cache_dir is not None:
if save_name is not None:
cache_path = Path(cache_dir) / f"{save_name}_{model_name}.npy"
else:
cache_path = Path(cache_dir) / f"{model_name}.npy"
# # if prediction already cached, don't recompute prediction
# if cache_dir is not None and os.path.exists(cache_path):
# print("Loading cached prediction for {}".format(model_name))
# y_pred = np.load(cache_path)
# else: # cached prediction not found, compute preds
print("Inferring model {}".format(path))
test_true = make_true_labels(cxr_true_labels_path=cxr_true_labels_path, cxr_labels=cxr_labels)
y_pred = run_softmax_eval(model, test_true, loader, cxr_labels, cxr_pair_template)
if cache_dir is not None:
Path(cache_dir).mkdir(exist_ok=True, parents=True)
np.save(file=cache_path, arr=y_pred)
predictions.append(y_pred)
# compute average predictions
y_pred_avg = np.mean(predictions, axis=0)
return predictions, y_pred_avg
# %%
predictions, y_pred_avg = ensemble_models(
model_paths=model_paths,
cxr_filepath=cxr_filepath,
cxr_labels=cxr_labels,
cxr_pair_template=cxr_pair_template,
cache_dir=cache_dir,
)
# %%
# save averaged preds
pred_name = "chexpert_preds.npy" # add name of preds
predictions_dir = predictions_dir / pred_name
np.save(file=predictions_dir, arr=y_pred_avg)
cxr_labels: List[str] = ['Atelectasis', 'Cardiomegaly', 'Edema', 'Consolidation', 'Fracture', 'Pleural Effusion', 'Pneumonia', 'Pneumothorax']
# make test_true
test_pred = y_pred_avg
test_true = make_true_labels(cxr_true_labels_path=cxr_true_labels_path, cxr_labels=cxr_labels)
# evaluate model
cxr_results = evaluate(test_pred, test_true, cxr_labels)
Atelectasis_auc = cxr_results['Atelectasis_auc']
Cardiomegaly_auc = cxr_results['Cardiomegaly_auc']
Consolidation_auc = cxr_results['Consolidation_auc']
Edema_auc = cxr_results['Edema_auc']
Effusion_auc = cxr_results['Pleural Effusion_auc']
Fracture_auc = cxr_results['Fracture_auc']
Pneumonia_auc = cxr_results['Pneumonia_auc']
Pneumothorax_auc = cxr_results['Pneumothorax_auc']
Mean_auc = (Atelectasis_auc + Cardiomegaly_auc + Consolidation_auc + Edema_auc + Effusion_auc + Fracture_auc + Pneumonia_auc + Pneumothorax_auc) / 8.
Mean_7_auc = (Atelectasis_auc + Cardiomegaly_auc + Consolidation_auc + Edema_auc + Effusion_auc + Pneumonia_auc + Pneumothorax_auc) / 7.
# boostrap evaluations for 95% confidence intervals
bootstrap_results = bootstrap(test_pred, test_true, cxr_labels)
# %%
# display AUC with confidence intervals
Atelectasis_auc = bootstrap_results[1]['Atelectasis_auc']['mean']
Cardiomegaly_auc = bootstrap_results[1]['Cardiomegaly_auc']['mean']
Consolidation_auc = bootstrap_results[1]['Consolidation_auc']['mean']
Edema_auc = bootstrap_results[1]['Edema_auc']['mean']
Effusion_auc = bootstrap_results[1]['Pleural Effusion_auc']['mean']
Fracture_auc = bootstrap_results[1]['Fracture_auc']['mean']
Pneumonia_auc = bootstrap_results[1]['Pneumonia_auc']['mean']
Pneumothorax_auc = bootstrap_results[1]['Pneumothorax_auc']['mean']
Mean_auc = (Atelectasis_auc + Cardiomegaly_auc + Edema_auc + Consolidation_auc + Effusion_auc + Fracture_auc + Pneumonia_auc + Pneumothorax_auc) / 8.
print('Model name: {} / Mean AUC: {}'.format(model_dir, Mean_auc))
bootstrap_results[1]
# %%