forked from isht7/pytorch-deeplab-resnet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathevalpyt2.py
99 lines (77 loc) · 3.73 KB
/
evalpyt2.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
import torch
from torch.autograd import Variable
import cv2
import numpy as np
import deeplab_resnet
import os
from docopt import docopt
docstr = """Evaluate ResNet-DeepLab trained on scenes (VOC 2012),a total of 21 labels including background
Usage:
evalpyt.py [options]
Options:
-h, --help Print this message
--visualize view outputs of each sketch
--snapPrefix=<str> Snapshot [default: VOC12_scenes_]
--testGTpath=<str> Ground truth path prefix [default: data/VOC_arg/SegmentationClass_label/]
--testIMpath=<str> Sketch images path prefix [default: data/VOC_arg/JPEGImages/]
--LISTpath=<str> Input image number list file [default: data/VOC_arg/ImageSets/Segmentation/val.txt]
--NoLabels=<int> The number of different labels in training data, VOC has 21 labels, including background [default: 21]
--gpu0=<int> GPU number [default: 1]
"""
args = docopt(docstr, version='v0.1')
print(args)
max_label = int(args['--NoLabels']) - 1 # labels from 0,1, ... 20(for VOC)
colors = np.random.randint(0, 256, (256, 3), dtype=np.uint8)
colors[0, :] = 0
colors[255, :] = 255
def get_iou(pred, gt):
if pred.shape != gt.shape:
print('pred shape', pred.shape, 'gt shape', gt.shape)
assert (pred.shape == gt.shape)
gt = gt.astype(int)
pred = pred.astype(int)
locs = (gt < 255)
sumim = gt + pred * (max_label + 1)
hs = np.bincount(sumim[locs], minlength=(max_label + 1)**2).reshape((max_label + 1), (max_label + 1))
return hs
gpu0 = int(args['--gpu0'])
im_path = args['--testIMpath']
model = deeplab_resnet.Res_Deeplab(int(args['--NoLabels']))
model.eval()
counter = 0
model.cuda(gpu0)
snapPrefix = args['--snapPrefix']
gt_path = args['--testGTpath']
img_list = open(args['--LISTpath']).readlines()
for iter in range(19, 22, 1): # TODO set the (different iteration)models that you want to evaluate on. Models are saved during training after each 1000 iters by default.
saved_state_dict = torch.load(os.path.join('data/snapshots/', snapPrefix + str(iter) + '000.pth'))
if counter == 0:
print(snapPrefix)
counter += 1
model.load_state_dict(saved_state_dict)
hist = np.zeros((max_label + 1, max_label + 1))
for i in img_list:
img = cv2.imread(os.path.join(im_path, i[:-1] + '.jpg')).astype(float)
img_original = img.copy()
img[:, :, 0] = img[:, :, 0] - 104.008
img[:, :, 1] = img[:, :, 1] - 116.669
img[:, :, 2] = img[:, :, 2] - 122.675
new_h = int((img_original.shape[0] - 1.0) / 8) * 8 + 1
new_w = int((img_original.shape[1] - 1.0) / 8) * 8 + 1
img = cv2.resize(img, (new_w, new_h))
gt = cv2.imread(os.path.join(gt_path, i[:-1] + '.png'), cv2.IMREAD_UNCHANGED).astype(np.uint8)
output = model(Variable(torch.from_numpy(img[np.newaxis, :].transpose(0, 3, 1, 2)).float(), volatile=True).cuda(gpu0))
output = output[3].cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = cv2.resize(output, (img_original.shape[1], img_original.shape[0]))
output = np.argmax(output, axis=2).astype(np.uint8)
if args['--visualize']:
cv2.imshow('img', img_original.astype(np.uint8))
gt_show = np.dstack((colors[gt, 0], colors[gt, 1], colors[gt, 2])).astype(np.uint8)
cv2.imshow('gt', gt_show)
output_show = np.dstack((colors[output, 0], colors[output, 1], colors[output, 2])).astype(np.uint8)
cv2.imshow('pred', output_show)
cv2.waitKey(10)
hist += get_iou(output, gt)
miou = np.diag(hist) / (1e-20 + hist.sum(1) + hist.sum(0) - np.diag(hist))
print('pytorch', iter, "Mean iou = ", np.sum(miou) / len(miou))