forked from HiLab-git/DTC
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_LA.py
More file actions
55 lines (43 loc) · 1.9 KB
/
test_LA.py
File metadata and controls
55 lines (43 loc) · 1.9 KB
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
import os
import argparse
import torch
from networks.vnet_sdf import VNet
from test_util import test_all_case
parser = argparse.ArgumentParser()
parser.add_argument('--root_path', type=str,
default='../data/2018LA_Seg_Training Set', help='Name of Experiment')
parser.add_argument('--model', type=str,
default='DTC_16labels', help='model_name')
parser.add_argument('--gpu', type=str, default='1', help='GPU to use')
parser.add_argument('--detail', type=int, default=1,
help='print metrics for every samples?')
parser.add_argument('--nms', type=int, default=1,
help='apply NMS post-procssing?')
FLAGS = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
snapshot_path = "../model/{}".format(FLAGS.model)
num_classes = 2
test_save_path = os.path.join(snapshot_path, "test/")
if not os.path.exists(test_save_path):
os.makedirs(test_save_path)
print(test_save_path)
with open(FLAGS.root_path + '/test.list', 'r') as f:
image_list = f.readlines()
image_list = [FLAGS.root_path + "/" + item.replace('\n', '') + "/mri_norm2.h5" for item in
image_list]
def test_calculate_metric():
net = VNet(n_channels=1, n_classes=num_classes-1,
normalization='batchnorm', has_dropout=False).cuda()
save_mode_path = os.path.join(
snapshot_path, 'best_model.pth')
net.load_state_dict(torch.load(save_mode_path))
print("init weight from {}".format(save_mode_path))
net.eval()
avg_metric = test_all_case(net, image_list, num_classes=num_classes,
patch_size=(112, 112, 80), stride_xy=18, stride_z=4,
save_result=True, test_save_path=test_save_path,
metric_detail=FLAGS.detail, nms=FLAGS.nms)
return avg_metric
if __name__ == '__main__':
metric = test_calculate_metric() # 6000
print(metric)