-
Notifications
You must be signed in to change notification settings - Fork 1
/
val_adaptiveunet.py
163 lines (124 loc) · 4.65 KB
/
val_adaptiveunet.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 argparse
import os
from glob import glob
import cv2
import torch
import torch.backends.cudnn as cudnn
import yaml
from albumentations.augmentations import transforms
from albumentations.core.composition import Compose
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import torch.optim as optim
import archs
from dataset import Dataset
from metrics import iou_score
from utils import AverageMeter
from albumentations import RandomRotate90,Resize
import time
from archs import UNet
import pdb
from copy import deepcopy
import torch
import torch.nn as nn
import torch.jit
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--name', default=None,
help='model name')
parser.add_argument('--target', default=None,
help='model name')
parser.add_argument('--dpg', default=None,
help='DPG model directory')
args = parser.parse_args()
return args
def configure_model(model):
model.train()
model.requires_grad_(False)
# enable grad + force batch statisics
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
m.requires_grad_(True)
# force use of batch stats in train and eval modes
m.track_running_stats = False
m.running_mean = None
m.running_var = None
return model
def main():
args = parse_args()
with open('models/%s/config.yml' % args.name, 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
print('-'*20)
for key in config.keys():
print('%s: %s' % (key, str(config[key])))
print('-'*20)
cudnn.benchmark = True
print("=> creating model %s" % config['arch'])
model = archs.__dict__[config['arch']](config['num_classes'],
config['input_channels'],
config['deep_supervision'])
# model = model.cuda()
model.load_state_dict(torch.load('models/%s/model.pth' %
config['name']))
model = model.cuda()
model.eval()
priormodel = archs.__dict__['priorunet'](config['num_classes'],
config['input_channels'],
config['deep_supervision'])
# print(args.dpg)
priormodel.load_state_dict(torch.load("%smodel.pth"%args.dpg))
priormodel.eval()
priormodel = priormodel.cuda()
# Data loading code
val_img_ids = glob(os.path.join('inputs', args.target, 'test','images', '*' + config['img_ext']))
val_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in val_img_ids]
val_transform = Compose([
Resize(config['input_h'], config['input_w']),
transforms.Normalize(),
])
val_dataset = Dataset(
img_ids=val_img_ids,
img_dir=os.path.join('inputs', args.target,'test', 'images'),
mask_dir=os.path.join('inputs', args.target,'test', 'masks'),
img_ext=config['img_ext'],
mask_ext=config['mask_ext'],
num_classes=config['num_classes'],
transform=val_transform)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
drop_last=False)
iou_avg_meter = AverageMeter()
dice_avg_meter = AverageMeter()
gput = AverageMeter()
cput = AverageMeter()
model = configure_model(model)
count = 0
for c in range(config['num_classes']):
os.makedirs(os.path.join('outputs', config['name'], str(c)), exist_ok=True)
with torch.no_grad():
for input, target, meta in tqdm(val_loader, total=len(val_loader)):
input = input.cuda()
target = target.cuda()
model = model.cuda()
model.eval()
# compute output
prior = priormodel(input)
output = model(input,prior)
iou,dice = iou_score(output, target)
iou_avg_meter.update(iou, input.size(0))
dice_avg_meter.update(dice, input.size(0))
output = torch.sigmoid(output).cpu().numpy()
output[output>=0.5]=1
output[output<0.5]=0
for i in range(len(output)):
for c in range(config['num_classes']):
cv2.imwrite(os.path.join('outputs', config['name'], str(c), meta['img_id'][i] + '.jpg'),
(output[i, c] * 255).astype('uint8'))
print('IoU: %.4f' % iou_avg_meter.avg)
print('Dice: %.4f' % dice_avg_meter.avg)
torch.cuda.empty_cache()
if __name__ == '__main__':
main()