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)