def inference(config, network, model_file, device, dataset, start, end,
result_queue):
torch.set_default_tensor_type('torch.FloatTensor')
torch.multiprocessing.set_sharing_strategy('file_system')
# init model
net = network()
net.cuda(device)
net = net.eval()
check_point = torch.load(model_file)
net.load_state_dict(check_point['state_dict'])
# init data
dataset.records = dataset.records[start:end]
data_iter = torch.utils.data.DataLoader(dataset=dataset, shuffle=False)
# inference
for (image, gt_boxes, im_info, ID) in data_iter:
pred_boxes = net(image.cuda(device), im_info.cuda(device))
scale = im_info[0, 2]
if config.test_nms_method == 'set_nms':
assert pred_boxes.shape[
-1] > 6, "Not EMD Network! Using normal_nms instead."
assert pred_boxes.shape[-1] % 6 == 0, "Prediction dim Error!"
top_k = pred_boxes.shape[-1] // 6
n = pred_boxes.shape[0]
pred_boxes = pred_boxes.reshape(-1, 6)
idents = np.tile(np.arange(n)[:, None], (1, top_k)).reshape(-1, 1)
pred_boxes = np.hstack((pred_boxes, idents))
keep = pred_boxes[:, 4] > config.pred_cls_threshold
pred_boxes = pred_boxes[keep]
keep = nms_utils.set_cpu_nms(pred_boxes, 0.5)
pred_boxes = pred_boxes[keep]
elif config.test_nms_method == 'normal_nms':
assert pred_boxes.shape[-1] % 6 == 0, "Prediction dim Error!"
pred_boxes = pred_boxes.reshape(-1, 6)
keep = pred_boxes[:, 4] > config.pred_cls_threshold
pred_boxes = pred_boxes[keep]
keep = nms_utils.cpu_nms(pred_boxes, config.test_nms)
pred_boxes = pred_boxes[keep]
elif config.test_nms_method == 'none':
assert pred_boxes.shape[-1] % 6 == 0, "Prediction dim Error!"
pred_boxes = pred_boxes.reshape(-1, 6)
keep = pred_boxes[:, 4] > config.pred_cls_threshold
pred_boxes = pred_boxes[keep]
else:
raise ValueError('Unknown NMS method.')
#if pred_boxes.shape[0] > config.detection_per_image and \
# config.test_nms_method != 'none':
# order = np.argsort(-pred_boxes[:, 4])
# order = order[:config.detection_per_image]
# pred_boxes = pred_boxes[order]
# recovery the scale
pred_boxes[:, :4] /= scale
pred_boxes[:, 2:4] -= pred_boxes[:, :2]
gt_boxes = gt_boxes[0].numpy()
gt_boxes[:, 2:4] -= gt_boxes[:, :2]
result_dict = dict(ID=ID[0],
height=int(im_info[0, -3]),
width=int(im_info[0, -2]),
dtboxes=boxes_dump(pred_boxes),
gtboxes=boxes_dump(gt_boxes))
result_queue.put_nowait(result_dict)
def post_process(pred_boxes, config, scale):
if config.test_nms_method == 'set_nms':
assert pred_boxes.shape[-1] > 6, "Not EMD Network! Using normal_nms instead."
assert pred_boxes.shape[-1] % 6 == 0, "Prediction dim Error!"
top_k = pred_boxes.shape[-1] // 6
n = pred_boxes.shape[0]
pred_boxes = pred_boxes.reshape(-1, 6)
idents = np.tile(np.arange(n)[:,None], (1, top_k)).reshape(-1, 1)
pred_boxes = np.hstack((pred_boxes, idents))
keep = pred_boxes[:, 4] > config.pred_cls_threshold
pred_boxes = pred_boxes[keep]
keep = nms_utils.set_cpu_nms(pred_boxes, 0.5)
pred_boxes = pred_boxes[keep]
elif config.test_nms_method == 'normal_nms':
assert pred_boxes.shape[-1] % 6 == 0, "Prediction dim Error!"
pred_boxes = pred_boxes.reshape(-1, 6)
keep = pred_boxes[:, 4] > config.pred_cls_threshold
pred_boxes = pred_boxes[keep]
keep = nms_utils.cpu_nms(pred_boxes, config.test_nms)
pred_boxes = pred_boxes[keep]
elif config.test_nms_method == 'none':
assert pred_boxes.shape[-1] % 6 == 0, "Prediction dim Error!"
pred_boxes = pred_boxes.reshape(-1, 6)
keep = pred_boxes[:, 4] > config.pred_cls_threshold
pred_boxes = pred_boxes[keep]
pred_boxes[:, :4] /= scale
keep = pred_boxes[:, 4] > config.visulize_threshold
pred_boxes = pred_boxes[keep]
return pred_boxes