-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathtest.py
119 lines (81 loc) · 3.73 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
import argparse
import torch
from torch.autograd import Variable
import torchvision
from torchvision.ops import nms
import os
from PIL import Image
import numpy as np
from skimage.draw import rectangle_perimeter
from models.HRCenterNet import HRCenterNet
input_size = 512
output_size = 128
test_tx = torchvision.transforms.Compose([
torchvision.transforms.Resize((input_size, input_size)),
torchvision.transforms.ToTensor(),
])
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('divece: ', device)
def main(args):
if not (args.log_dir == None):
print("Load checkpoint from " + args.log_dir)
checkpoint = torch.load(args.log_dir, map_location="cpu")
model = HRCenterNet()
model.load_state_dict(checkpoint['model'])
model = model.to(device)
model.eval()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for file in os.listdir(args.data_dir):
img = Image.open(args.data_dir + file).convert("RGB")
image_tensor = test_tx(img)
image_tensor = image_tensor.unsqueeze_(0)
inp = Variable(image_tensor)
inp = inp.to(device, dtype=torch.float)
predict = model(inp)
out_img = _nms(args, img, predict, nms_score=0.3, iou_threshold=0.1)
print('saving image to ', args.output_dir + file )
Image.fromarray(out_img).save(args.output_dir + file)
def _nms(args, img, predict, nms_score, iou_threshold):
bbox = list()
score_list = list()
im_draw = np.asarray(torchvision.transforms.functional.resize(img, (img.size[1], img.size[0]))).copy()
heatmap=predict.data.cpu().numpy()[0, 0, ...]
offset_y = predict.data.cpu().numpy()[0, 1, ...]
offset_x = predict.data.cpu().numpy()[0, 2, ...]
width_map = predict.data.cpu().numpy()[0, 3, ...]
height_map = predict.data.cpu().numpy()[0, 4, ...]
for j in np.where(heatmap.reshape(-1, 1) >= nms_score)[0]:
row = j // output_size
col = j - row*output_size
bias_x = offset_x[row, col] * (img.size[1] / output_size)
bias_y = offset_y[row, col] * (img.size[0] / output_size)
width = width_map[row, col] * output_size * (img.size[1] / output_size)
height = height_map[row, col] * output_size * (img.size[0] / output_size)
score_list.append(heatmap[row, col])
row = row * (img.size[1] / output_size) + bias_y
col = col * (img.size[0] / output_size) + bias_x
top = row - width // 2
left = col - height // 2
bottom = row + width // 2
right = col + height // 2
start = (top, left)
end = (bottom, right)
bbox.append([top, left, bottom, right])
_nms_index = torchvision.ops.nms(torch.FloatTensor(bbox), scores=torch.flatten(torch.FloatTensor(score_list)), iou_threshold=iou_threshold)
for k in range(len(_nms_index)):
top, left, bottom, right = bbox[_nms_index[k]]
start = (top, left)
end = (bottom, right)
rr, cc = rectangle_perimeter(start, end=end,shape=(img.size[1], img.size[0]))
im_draw[rr, cc] = (255, 0, 0)
return im_draw
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="HRCenterNet.")
parser.add_argument("--data_dir", required=True,
help="Path to the testing images folder, preprocessed for torchvision.")
parser.add_argument("--log_dir", required=True, default=None,
help="Where to load for the pretrained model.")
parser.add_argument("--output_dir", default='./',
help="Where to save for the outputs.")
main(parser.parse_args())