def test_nms_rotated_0_degree_cuda(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS incompatible between CPU and CUDA for IoU={}"
for iou in [0.2, 0.5, 0.8]:
r_cpu = nms_rotated(rotated_boxes, scores, iou)
r_cuda = nms_rotated(rotated_boxes.cuda(), scores.cuda(), iou)
self.assertLessEqual(nms_edit_distance(r_cpu, r_cuda.cpu()), 1, err_msg.format(iou))
def test_nms_rotated_0_degree_cpu(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
for iou in [0.5]:
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
keep = nms_rotated(rotated_boxes, scores, iou)
assert torch.equal(keep, keep_ref), err_msg.format(iou)
def test_nms_rotated_0_degree_cpu(self, device="cpu"):
N = 1000
boxes, scores = self._create_tensors(N, device=device)
rotated_boxes = torch.zeros(N, 5, device=device)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
for iou in [0.2, 0.5, 0.8]:
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
keep = nms_rotated(rotated_boxes, scores, iou)
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
def draw_boxes(reg, prob, images, calibs, ids, tag):
prob = prob.reshape(cfg.N, -1)
batch_boxes3d = delta_to_boxes3d(reg, cfg.anchors)
mask = torch.gt(prob, cfg.score_threshold)
mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)
out_images = []
for batch_id in range(cfg.N):
boxes3d = torch.masked_select(batch_boxes3d[batch_id],
mask_reg[batch_id]).view(-1, 7)
scores = torch.masked_select(prob[batch_id], mask[batch_id])
image = images[batch_id]
calib = calibs[batch_id]
id = ids[batch_id]
if len(boxes3d) != 0:
# boxes3d_corner, sizes = box3d_center_to_corner_batch(boxes3d)
# boxes2d_bottom = box3d_corner_to_top_batch(boxes3d_corner)
# boxes2d_score = torch.cat((boxes2d_bottom, scores.unsqueeze(1), torch.arange(0, len(boxes2d_bottom)).float().unsqueeze(1).to(cfg.device)), dim=1)
# args = torch.argsort(boxes2d_score[:, 8], descending=True)
# boxes2d_score = boxes2d_score[args]
# vac = torch.nonzero(sizes > 1e-2)
# boxes2d_score = boxes2d_score[vac].squeeze()
# if boxes2d_score.shape[0] == 0:
# out_images.append(image)
# continue
# NMS
# boxes2d_score = nms(boxes2d_score, cfg.nms_threshold)
index = nms_rotated(boxes3d[..., [0, 1, 5, 4, 6]], scores, 0.01)
if len(index) is None:
out_images.append(image)
continue
# boxes3d_corner_keep = boxes3d_corner[boxes2d_score[:, 9].long()]
boxes3d = boxes3d[index]
print("No. %d objects detected" % len(boxes3d))
boxes3d_corner_keep, _ = box3d_center_to_corner_batch(boxes3d)
boxes3d_corner_keep = boxes3d_corner_keep.cpu().numpy()
boxes3d_corner_keep = np.transpose(boxes3d_corner_keep, (0, 2, 1))
rgb_2D = project_velo2rgb(boxes3d_corner_keep, calib)
img_with_box = draw_rgb_projections(image,
rgb_2D,
color=(0, 0, 255),
thickness=1)
out_images.append(img_with_box)
return np.array(out_images)
def test_nms_rotated_90_degrees_cpu(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
# Note for rotated_boxes[:, 2] and rotated_boxes[:, 3]:
# widths and heights are intentionally swapped here for 90 degrees case
# so that the reference horizontal nms could be used
rotated_boxes[:, 2] = boxes[:, 3] - boxes[:, 1]
rotated_boxes[:, 3] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 4] = torch.ones(N) * 90
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
for iou in [0.2, 0.5, 0.8]:
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
keep = nms_rotated(rotated_boxes, scores, iou)
assert torch.equal(keep, keep_ref), err_msg.format(iou)
def forward(self, boxes, scores, threshold):
return nms_rotated(boxes, scores, threshold)
