-
Notifications
You must be signed in to change notification settings - Fork 2
/
test_GobletNet.py
174 lines (145 loc) · 7.04 KB
/
test_GobletNet.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
164
165
166
167
168
169
170
171
172
173
174
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
import argparse
import time
import os
import numpy as np
from torch.backends import cudnn
import random
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from config.dataset_config.dataset_cfg import dataset_cfg
from config.augmentation.online_aug import data_transform, data_normalize
from models.getnetwork import get_network
from dataload.dataset_2d import imagefloder_iitnn
from config.train_test_config.train_test_config import print_test_eval, save_test_2d
from warnings import simplefilter
simplefilter(action='ignore', category=FutureWarning)
def init_seeds(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
np.random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-pd', '--path_dataset', default='.../GobletNet/dataset/EPFL')
parser.add_argument('-p', '--path_model', default='.../GobletNet/checkpoints/EPFL/GobletNet-l=0.8-e=200-s=50-g=0.5-b=4-w=20-image-H_0.3_db2/best_wbhf_Jc_0.8118.pth')
parser.add_argument('--path_seg_results', default='.../GobletNet/EPFL/seg_pred/test')
parser.add_argument('--dataset_name', default='EPFL', help='EPFL, CREMI, SNEMI3D, UroCell, MitoEM, Nanowire, BetaSeg')
parser.add_argument('--input1', default='image')
parser.add_argument('--input2', default='H_0.3_db2')
parser.add_argument('--if_mask', default=True)
parser.add_argument('--threshold', default=0.5000)
parser.add_argument('-n', '--network', default='GobletNet')
parser.add_argument('-b', '--batch_size', default=8, type=int)
parser.add_argument('--local_rank', default=-1, type=int)
parser.add_argument('--rank_index', default=0, help='0, 1, 2, 3')
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)
dist.init_process_group(backend='nccl', init_method='env://')
rank = torch.distributed.get_rank()
ngpus_per_node = torch.cuda.device_count()
init_seeds(rank + 1)
# Config
dataset_name = args.dataset_name
cfg = dataset_cfg(dataset_name)
print_num = 42 + (cfg['NUM_CLASSES'] - 3) * 7
print_num_minus = print_num - 2
# Results Save
if not os.path.exists(args.path_seg_results) and rank == args.rank_index:
os.mkdir(args.path_seg_results)
path_seg_results = args.path_seg_results + '/' + str(dataset_name)
if not os.path.exists(path_seg_results) and rank == args.rank_index:
os.mkdir(path_seg_results)
path_seg_results = path_seg_results + '/' + str(os.path.splitext(os.path.split(args.path_model)[1])[0])
if not os.path.exists(path_seg_results) and rank == args.rank_index:
os.mkdir(path_seg_results)
# Dataset
if args.input1 == 'image':
input1_mean = 'MEAN'
input1_std = 'STD'
else:
input1_mean = 'MEAN_' + args.input1
input1_std = 'STD_' + args.input1
if args.input2 == 'image':
input2_mean = 'MEAN'
input2_std = 'STD'
else:
input2_mean = 'MEAN_' + args.input2
input2_std = 'STD_' + args.input2
data_transforms = data_transform(cfg['SIZE'])
data_normalize_1 = data_normalize(cfg[input1_mean], cfg[input1_std])
data_normalize_2 = data_normalize(cfg[input2_mean], cfg[input2_std])
dataset_val = imagefloder_iitnn(
data_dir=args.path_dataset + '/val',
input1=args.input1,
input2=args.input2,
data_transform_1=data_transforms['val'],
data_normalize_1=data_normalize_1,
data_normalize_2=data_normalize_2,
sup=True,
num_images=None,
)
val_sampler = torch.utils.data.distributed.DistributedSampler(dataset_val, shuffle=False)
dataloaders = dict()
dataloaders['val'] = DataLoader(dataset_val, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=16, sampler=val_sampler)
num_batches = {'val': len(dataloaders['val'])}
# Model
model = get_network(args.network, cfg['IN_CHANNELS'], cfg['NUM_CLASSES'])
model = model.cuda()
model = DistributedDataParallel(model, device_ids=[args.local_rank])
state_dict = torch.load(args.path_model)
model.load_state_dict(state_dict=state_dict)
dist.barrier()
# Test
since = time.time()
with torch.no_grad():
model.eval()
for i, data in enumerate(dataloaders['val']):
inputs_test = Variable(data['image'].cuda(non_blocking=True))
inputs_wavelet_test = Variable(data['image_2'].cuda(non_blocking=True))
name_test = data['ID']
if args.if_mask:
mask_test = Variable(data['mask'].cuda(non_blocking=True))
outputs_test = model(inputs_test, inputs_wavelet_test)
if args.if_mask:
if i == 0:
score_list_test = outputs_test
name_list_test = name_test
mask_list_test = mask_test
else:
# elif 0 < i <= num_batches['val'] / 16:
score_list_test = torch.cat((score_list_test, outputs_test), dim=0)
name_list_test = np.append(name_list_test, name_test, axis=0)
mask_list_test = torch.cat((mask_list_test, mask_test), dim=0)
torch.cuda.empty_cache()
else:
save_test_2d(cfg['NUM_CLASSES'], outputs_test, name_test, args.threshold, path_seg_results, cfg['PALETTE'])
torch.cuda.empty_cache()
if args.if_mask:
score_gather_list_test = [torch.zeros_like(score_list_test) for _ in range(ngpus_per_node)]
torch.distributed.all_gather(score_gather_list_test, score_list_test)
score_list_test = torch.cat(score_gather_list_test, dim=0)
mask_gather_list_test = [torch.zeros_like(mask_list_test) for _ in range(ngpus_per_node)]
torch.distributed.all_gather(mask_gather_list_test, mask_list_test)
mask_list_test = torch.cat(mask_gather_list_test, dim=0)
name_gather_list_test = [None for _ in range(ngpus_per_node)]
torch.distributed.all_gather_object(name_gather_list_test, name_list_test)
name_list_test = np.concatenate(name_gather_list_test, axis=0)
if args.if_mask and rank == args.rank_index:
print('=' * print_num)
test_eval_list = print_test_eval(cfg['NUM_CLASSES'], score_list_test, mask_list_test, print_num_minus)
save_test_2d(cfg['NUM_CLASSES'], score_list_test, name_list_test, test_eval_list[0], path_seg_results, cfg['PALETTE'])
torch.cuda.empty_cache()
if rank == args.rank_index:
time_elapsed = time.time() - since
m, s = divmod(time_elapsed, 60)
h, m = divmod(m, 60)
print('-' * print_num)
print('| Testing Completed In {:.0f}h {:.0f}mins {:.0f}s'.format(h, m, s).ljust(print_num_minus, ' '), '|')
print('=' * print_num)