-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathtest.py
184 lines (160 loc) · 6.79 KB
/
test.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
175
176
177
178
179
180
181
182
183
184
import argparse
import os
from PIL import Image
import numpy as np
import torchvision.transforms as tr
from collections import OrderedDict
parser = argparse.ArgumentParser()
# CHAOS_liver Material Face_segment
parser.add_argument('--dataset', default='face',
choices=['material', 'medical', 'medical_material', 'face', 'aaf_face', 'XCX'],
help='Which dataset to use')
parser.add_argument('--model_path', type=str,
help='Choose which model to use,unet...')
parser.add_argument('--fold', type=str, default='0', help='chooses fold...')
parser.add_argument('--algorithm', type=str, default='fedavg',
help='Chooses which federated learning algorithm to use')
parser.add_argument('--client_id', type=int, default=-1,
help='Client id for local trainer')
parser.add_argument('--bg_value', type=int, default=0,
help='Background pixel value for evaluating test set [Not recommended]')
parser.add_argument('--mask_reverse', type=int, default=0,
help='Reverse the mask for binary mask [Recommendation for boundary segmentation task]')
parser.add_argument('--save_image', type=int, default=1,
help='Whether to save the predicted mask image')
parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids:e.g.0 0,1,2 1,2 use -1 for CPU')
op = parser.parse_args()
dataset = op.dataset
if dataset == 'material':
from opt.material_options import TrainOptions, TestOptions
elif op.dataset == 'XCX':
from opt.XCX_options import TrainOptions, TestOptions
elif dataset == 'face':
from opt.face_options import TrainOptions, TestOptions
elif dataset == 'medical':
from opt.medical_options import TrainOptions, TestOptions
elif op.dataset == 'medical_material':
from opt.medical_material_options import TrainOptions, TestOptions
elif op.dataset == 'aaf_face':
from opt.aaf_face_options import TrainOptions, TestOptions
fold = op.fold
fed_methods = op.algorithm
model_path = op.model_path + fold + '.pkl'
method = dataset + '_' + fed_methods + '_fold_' + fold
client_id = op.client_id
bg_value = op.bg_value
mask_reverse = op.mask_reverse
from data.dataloader.data_loader import CreateDataLoader
from models import create_model
import time
from util.util import mkdir, get_time_str
from util.metrics import get_total_evaluation
from util.evaluation import get_map_miou_vi_ri_ari
import torch
import tqdm
is_save_image = True if op.save_image > 0 else False
if __name__ == '__main__':
gpu_ids = op.gpu_ids.split(',')
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(i) for i in gpu_ids])
# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
opt = TestOptions().parse()
opt_train = TrainOptions().parse()
# hard-code some parameters for test
opt.num_threads = 1 # test code only supports num_threads = 1
opt.batch_size = 1 # test code only supports batch_size = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
opt.client_id = client_id
if opt.no_normalize:
transform = tr.ToTensor()
else:
transform = tr.Compose([tr.ToTensor(),
tr.Normalize(mean=0.5,
std=0.5)
])
data_loader = CreateDataLoader(opt, dataroot=opt.dataroot, image_dir=opt.test_img_dir, \
label_dir=opt.test_label_dir, record_txt=None, transform=transform, is_aug=False)
print(opt.dataroot)
print(opt.test_img_dir)
print(opt.test_label_dir)
dataset = data_loader.load_data()
datasize = len(data_loader)
print('#test images = %d, batchsize = %d' % (datasize, opt.batch_size))
state_dict = torch.load(model_path, map_location='cuda:0')
sd = OrderedDict()
# opt.isTrain = True
opt_train.model = 'unet'
opt_train.mode = 'test'
model = create_model(opt_train)
# if 'fedst' in method:
# for i in state_dict.keys():
# if 'netDseg' in i:
# key = '.'.join(['net'] + i.split('.')[1:])
# sd[key] = state_dict[i]
# model.load_state_dict(sd)
# else:
# model.load_state_dict(state_dict)
model.load_state_dict(state_dict)
print(f"[Test] Model {model_path} has loaded.")
# output_dir = os.path.join(opt.results_dir, opt.name, '%s' % method)
save_dir_name = '_'.join(str(op.model_path).split('/')[-2:]) + '_' + fold
output_dir = os.path.join(opt.results_dir, opt.name, save_dir_name)
if mask_reverse:
output_dir += '_reverse'
output_dir += '_' + get_time_str() # Add time str
eval_file = output_dir + '_eval.txt'
print('Eval result has saved to %s' % output_dir)
mkdir(output_dir)
eval_results = {}
count = 0
with open(eval_file, 'w') as log:
now = time.strftime('%c')
log.write('=============Evaluation (%s)=============\n' % now)
result = []
for i, data in enumerate(tqdm.tqdm(dataset)):
c_idx = 0
model.set_input(data)
with torch.no_grad():
pred = model()[0]
print(pred.shape)
pred = pred.max(0)[1].cpu().numpy()
print(pred.shape)
img_path = data['path'][0]
short_path = os.path.basename(img_path)
name = os.path.splitext(short_path)[0]
image_name = '%s.png' % name
eval_start = time.time()
count += 1
mask = data['label'].squeeze().numpy().astype(np.uint8)
print(f"bg_value: {bg_value}")
# mask reverse for binary mask
if mask_reverse:
mask = 1 - mask
pred = 1 - pred
if is_save_image:
mask_dir = os.path.join(output_dir, 'mask')
os.makedirs(mask_dir, exist_ok=True)
binary_dir = os.path.join(output_dir, 'binary')
os.makedirs(binary_dir, exist_ok=True)
Image.fromarray(pred.astype(np.uint8)).save(os.path.join(mask_dir, image_name))
Image.fromarray(pred.astype(bool)).save(os.path.join(binary_dir, image_name))
eval_result = get_total_evaluation(pred, mask)
print(data['label'].data[0].shape)
message = '%04d: %s \t' % (count, name)
for k, v in eval_result.items():
if k in eval_results:
eval_results[k] += v
else:
eval_results[k] = v
message += '%s: %.5f\t' % (k, v)
eval_result['name'] = name
result.append(eval_result)
print(message, 'cost: %.4f' % (time.time() - eval_start))
with open(eval_file, 'a') as log:
log.write(message + '\n')
message = 'total %d:\n' % count
for k, v in eval_results.items():
message += 'm_%s: %.5f\t' % (k, v / count)
print(message)
with open(output_dir + '_eval.txt', 'a') as log:
log.write(message + '\n')