def test_net(tester, logger, dets, det_range):
# here we assume all boxes are pre-processed.
nms_method = 'nms'
nms_thresh = 1.
min_scores = 1e-10
min_box_size = 0. # 8 ** 2
all_res = []
dump_results = []
start_time = time.time()
img_start = det_range[0]
while img_start < det_range[1]:
img_end = img_start + 1
im_info = dets[img_start]
while img_end < det_range[1] and dets[img_end]['image_id'] == im_info['image_id']:
img_end += 1
test_data = dets[img_start:img_end]
img_start = img_end
iter_avg_cost_time = (time.time() - start_time) / (img_end - det_range[0])
print('ran %.ds >> << left %.ds' % (
iter_avg_cost_time * (img_end - det_range[0]), iter_avg_cost_time * (det_range[1] - img_end)))
all_res.append([])
# get box detections
cls_dets = np.zeros((len(test_data), 5), dtype=np.float32)
for i in range(len(test_data)):
bbox = np.asarray(test_data[i]['bbox'])
cls_dets[i, :4] = np.array([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]])
cls_dets[i, 4] = np.array(test_data[i]['score'])
# nms and filter
keep = np.where((cls_dets[:, 4] >= min_scores) &
((cls_dets[:, 3] - cls_dets[:, 1]) * (cls_dets[:, 2] - cls_dets[:, 0]) >= min_box_size))[0]
cls_dets = cls_dets[keep]
if len(cls_dets) > 0:
if nms_method == 'nms':
keep = gpu_nms(cls_dets, nms_thresh)
elif nms_method == 'soft':
keep = cpu_soft_nms(np.ascontiguousarray(cls_dets, dtype=np.float32), method=2)
else:
assert False
cls_dets = cls_dets[keep]
test_data = np.asarray(test_data)[keep]
if len(keep) == 0:
continue
# crop and detect keypoints
cls_skeleton = np.zeros((len(test_data), cfg.nr_skeleton, 3))
crops = np.zeros((len(test_data), 4))
cfg.batch_size = 32
batch_size = cfg.batch_size // 2
for test_id in range(0, len(test_data), batch_size):
start_id = test_id
end_id = min(len(test_data), test_id + batch_size)
test_imgs = []
details = []
for i in range(start_id, end_id):
test_img, detail = Preprocessing(test_data[i], stage='test')
test_imgs.append(test_img)
details.append(detail)
details = np.asarray(details)
feed = test_imgs
for i in range(end_id - start_id):
ori_img = test_imgs[i][0].transpose(1, 2, 0)
flip_img = cv2.flip(ori_img, 1)
feed.append(flip_img.transpose(2, 0, 1)[np.newaxis, ...])
feed = np.vstack(feed)
res = tester.predict_one([feed.transpose(0, 2, 3, 1).astype(np.float32)])[0]
res = res.transpose(0, 3, 1, 2)
for i in range(end_id - start_id):
fmp = res[end_id - start_id + i].transpose((1, 2, 0))
fmp = cv2.flip(fmp, 1)
fmp = list(fmp.transpose((2, 0, 1)))
for (q, w) in cfg.symmetry:
fmp[q], fmp[w] = fmp[w], fmp[q]
fmp = np.array(fmp)
res[i] += fmp
res[i] /= 2
for test_image_id in range(start_id, end_id):
r0 = res[test_image_id - start_id].copy()
r0 /= 255.
r0 += 0.5
for w in range(cfg.nr_skeleton):
res[test_image_id - start_id, w] /= np.amax(res[test_image_id - start_id, w])
border = 10
dr = np.zeros((cfg.nr_skeleton, cfg.output_shape[0] + 2 * border, cfg.output_shape[1] + 2 * border))
dr[:, border:-border, border:-border] = res[test_image_id - start_id][:cfg.nr_skeleton].copy()
for w in range(cfg.nr_skeleton):
dr[w] = cv2.GaussianBlur(dr[w], (21, 21), 0)
for w in range(cfg.nr_skeleton):
lb = dr[w].argmax()
y, x = np.unravel_index(lb, dr[w].shape)
dr[w, y, x] = 0
lb = dr[w].argmax()
py, px = np.unravel_index(lb, dr[w].shape)
y -= border
x -= border
py -= border + y
px -= border + x
ln = (px ** 2 + py ** 2) ** 0.5
delta = 0.25
if ln > 1e-3:
x += delta * px / ln
y += delta * py / ln
x = max(0, min(x, cfg.output_shape[1] - 1))
y = max(0, min(y, cfg.output_shape[0] - 1))
cls_skeleton[test_image_id, w, :2] = (x * 4 + 2, y * 4 + 2)
cls_skeleton[test_image_id, w, 2] = r0[w, int(round(y) + 1e-10), int(round(x) + 1e-10)]
# map back to original images
crops[test_image_id, :] = details[test_image_id - start_id, :]
for w in range(cfg.nr_skeleton):
cls_skeleton[test_image_id, w, 0] = cls_skeleton[test_image_id, w, 0] / cfg.data_shape[1] * (
crops[test_image_id][2] - crops[test_image_id][0]) + crops[test_image_id][0]
cls_skeleton[test_image_id, w, 1] = cls_skeleton[test_image_id, w, 1] / cfg.data_shape[0] * (
crops[test_image_id][3] - crops[test_image_id][1]) + crops[test_image_id][1]
all_res[-1] = [cls_skeleton.copy(), cls_dets.copy()]
cls_partsco = cls_skeleton[:, :, 2].copy().reshape(-1, cfg.nr_skeleton)
cls_skeleton[:, :, 2] = 1
cls_scores = cls_dets[:, -1].copy()
# rescore
cls_dets[:, -1] = cls_scores * cls_partsco.mean(axis=1)
cls_skeleton = np.concatenate(
[cls_skeleton.reshape(-1, cfg.nr_skeleton * 3), (cls_scores * cls_partsco.mean(axis=1))[:, np.newaxis]],
axis=1)
for i in range(len(cls_skeleton)):
result = dict(image_id=im_info['image_id'], category_id=1, score=float(round(cls_skeleton[i][-1], 4)),
keypoints=cls_skeleton[i][:-1].round(3).tolist())
dump_results.append(result)
return all_res, dump_results
def execute(self, inputs, outputs):
""" inputs: list of (x0, y0, x1, y1, score)"""
in_ = inputs[0].get_value()
keep = gpu_nms(in_, thresh=self._iou_threshold)
outputs[0].set_value(keep)