-
Notifications
You must be signed in to change notification settings - Fork 2
/
predict.py
153 lines (134 loc) · 5.99 KB
/
predict.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
import argparse
import os
import pickle
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics.classification import confusion_matrix
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from resnet import resnet34
from dataset import ECGDataset
from utils import cal_scores, find_optimal_threshold, split_data
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--data-dir', type=str, default='data/CPSC', help='Directory to data dir')
parser.add_argument('--leads', type=str, default='all', help='ECG leads to use')
parser.add_argument('--seed', type=int, default=42, help='Seed to split data')
parser.add_argument('--batch-size', type=int, default=32, help='Batch size')
parser.add_argument('--num-workers', type=int, default=4, help='Number of workers to load data')
parser.add_argument('--use-gpu', default=False, action='store_true', help='Use gpu')
parser.add_argument('--model-path', type=str, default='', help='Path to saved model')
return parser.parse_args()
def get_thresholds(val_loader, net, device, threshold_path):
print('Finding optimal thresholds...')
if os.path.exists(threshold_path):
return pickle.load(open(threshold_path, 'rb'))
output_list, label_list = [], []
for _, (data, label) in enumerate(tqdm(val_loader)):
data, labels = data.to(device), label.to(device)
output = net(data)
output = torch.sigmoid(output)
output_list.append(output.data.cpu().numpy())
label_list.append(labels.data.cpu().numpy())
y_trues = np.vstack(label_list)
y_scores = np.vstack(output_list)
thresholds = []
for i in range(y_trues.shape[1]):
y_true = y_trues[:, i]
y_score = y_scores[:, i]
threshold = find_optimal_threshold(y_true, y_score)
thresholds.append(threshold)
# pickle.dump(thresholds, open(threshold_path, 'wb'))
return thresholds
def apply_thresholds(test_loader, net, device, thresholds):
output_list, label_list = [], []
for _, (data, label) in enumerate(tqdm(test_loader)):
data, labels = data.to(device), label.to(device)
output = net(data)
output = torch.sigmoid(output)
output_list.append(output.data.cpu().numpy())
label_list.append(labels.data.cpu().numpy())
y_trues = np.vstack(label_list)
y_scores = np.vstack(output_list)
y_preds = []
scores = []
for i in range(len(thresholds)):
y_true = y_trues[:, i]
y_score = y_scores[:, i]
y_pred = (y_score >= thresholds[i]).astype(int)
scores.append(cal_scores(y_true, y_pred, y_score))
y_preds.append(y_pred)
y_preds = np.array(y_preds).transpose()
scores = np.array(scores)
print('Precisions:', scores[:, 0])
print('Recalls:', scores[:, 1])
print('F1s:', scores[:, 2])
print('AUCs:', scores[:, 3])
print('Accs:', scores[:, 4])
print(np.mean(scores, axis=0))
plot_cm(y_trues, y_preds)
def plot_cm(y_trues, y_preds, normalize=True, cmap=plt.cm.Blues):
classes = ['SNR', 'AF', 'IAVB', 'LBBB', 'RBBB', 'PAC', 'PVC', 'STD', 'STE']
for i, label in enumerate(classes):
y_true = y_trues[:, i]
y_pred = y_preds[:, i]
cm = confusion_matrix(y_true, y_pred)
if normalize:
cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
fig, ax = plt.subplots(figsize=(4, 4))
im = ax.imshow(cm, interpolation='nearest', cmap=cmap)
ax.figure.colorbar(im, ax=ax)
ax.set(xticks=np.arange(cm.shape[1]),
yticks=np.arange(cm.shape[0]),
xticklabels=[0, 1], yticklabels=[0, 1],
title=label,
ylabel='True label',
xlabel='Predicted label')
plt.setp(ax.get_xticklabels(), ha="center")
fmt = '.3f' if normalize else 'd'
thresh = cm.max() / 2.
for i in range(cm.shape[0]):
for j in range(cm.shape[1]):
ax.text(j, i, format(cm[i, j], fmt),
ha="center", va="center",
color="white" if cm[i, j] > thresh else "black")
np.set_printoptions(precision=3)
fig.tight_layout()
plt.savefig(f'results/{label}.png')
plt.close(fig)
if __name__ == "__main__":
args = parse_args()
data_dir = os.path.normpath(args.data_dir)
database = os.path.basename(data_dir)
if not args.model_path:
args.model_path = f'models/resnet34_{database}_{args.leads}_{args.seed}.pth'
args.threshold_path = f'models/{database}-threshold.pkl'
if args.use_gpu and torch.cuda.is_available():
device = torch.device('cuda:0')
else:
device = 'cpu'
if args.leads == 'all':
leads = 'all'
nleads = 12
else:
leads = args.leads.split(',')
nleads = len(leads)
data_dir = args.data_dir
label_csv = os.path.join(data_dir, 'labels.csv')
net = resnet34(input_channels=nleads).to(device)
net.load_state_dict(torch.load(args.model_path, map_location=device))
net.eval()
train_folds, val_folds, test_folds = split_data(seed=args.seed)
train_dataset = ECGDataset('train', data_dir, label_csv, train_folds, leads)
train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)
val_dataset = ECGDataset('val', data_dir, label_csv, val_folds, leads)
val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)
test_dataset = ECGDataset('test', data_dir, label_csv, test_folds, leads)
test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)
thresholds = get_thresholds(val_loader, net, device, args.threshold_path)
print('Thresholds:', thresholds)
print('Results on validation data:')
apply_thresholds(val_loader, net, device, thresholds)
print('Results on test data:')
apply_thresholds(test_loader, net, device, thresholds)