def multiclass_nms_single(self, obj_scores, sem_scores, bbox, points,
input_meta):
"""Multi-class nms in single batch.
Args:
obj_scores (torch.Tensor): Objectness score of bounding boxes.
sem_scores (torch.Tensor): semantic class score of bounding boxes.
bbox (torch.Tensor): Predicted bounding boxes.
points (torch.Tensor): Input points.
input_meta (dict): Point cloud and image's meta info.
Returns:
tuple[torch.Tensor]: Bounding boxes, scores and labels.
"""
bbox = input_meta['box_type_3d'](
bbox,
box_dim=bbox.shape[-1],
with_yaw=self.bbox_coder.with_rot,
origin=(0.5, 0.5, 0.5))
box_indices = bbox.points_in_boxes(points)
corner3d = bbox.corners
minmax_box3d = corner3d.new(torch.Size((corner3d.shape[0], 6)))
minmax_box3d[:, :3] = torch.min(corner3d, dim=1)[0]
minmax_box3d[:, 3:] = torch.max(corner3d, dim=1)[0]
nonempty_box_mask = box_indices.T.sum(1) > 5
bbox_classes = torch.argmax(sem_scores, -1)
nms_selected = aligned_3d_nms(minmax_box3d[nonempty_box_mask],
obj_scores[nonempty_box_mask],
bbox_classes[nonempty_box_mask],
self.test_cfg.nms_thr)
# filter empty boxes and boxes with low score
scores_mask = (obj_scores > self.test_cfg.score_thr)
nonempty_box_inds = torch.nonzero(
nonempty_box_mask, as_tuple=False).flatten()
nonempty_mask = torch.zeros_like(bbox_classes).scatter(
0, nonempty_box_inds[nms_selected], 1)
selected = (nonempty_mask.bool() & scores_mask.bool())
if self.test_cfg.per_class_proposal:
bbox_selected, score_selected, labels = [], [], []
for k in range(sem_scores.shape[-1]):
bbox_selected.append(bbox[selected].tensor)
score_selected.append(obj_scores[selected] *
sem_scores[selected][:, k])
labels.append(
torch.zeros_like(bbox_classes[selected]).fill_(k))
bbox_selected = torch.cat(bbox_selected, 0)
score_selected = torch.cat(score_selected, 0)
labels = torch.cat(labels, 0)
else:
bbox_selected = bbox[selected].tensor
score_selected = obj_scores[selected]
labels = bbox_classes[selected]
return bbox_selected, score_selected, labels
def test_aligned_3d_nms():
from mmdet3d.core.post_processing import aligned_3d_nms
boxes = torch.tensor([[1.2261, 0.6679, -1.2678, 2.6547, 1.0428, 0.1000],
[5.0919, 0.6512, 0.7238, 5.4821, 1.2451, 2.1095],
[6.8392, -1.2205, 0.8570, 7.6920, 0.3220, 3.2223],
[3.6900, -0.4235, -1.0380, 4.4415, 0.2671, -0.1442],
[4.8071, -1.4311, 0.7004, 5.5788, -0.6837, 1.2487],
[2.1807, -1.5811, -1.1289, 3.0151, -0.1346, -0.5351],
[4.4631, -4.2588, -1.1403, 5.3012, -3.4463, -0.3212],
[4.7607, -3.3311, 0.5993, 5.2976, -2.7874, 1.2273],
[3.1265, 0.7113, -0.0296, 3.8944, 1.3532, 0.9785],
[5.5828, -3.5350, 1.0105, 8.2841, -0.0405, 3.3614],
[3.0003, -2.1099, -1.0608, 5.3423, 0.0328, 0.6252],
[2.7148, 0.6082, -1.1738, 3.6995, 1.2375, -0.0209],
[4.9263, -0.2152, 0.2889, 5.6963, 0.3416, 1.3471],
[5.0713, 1.3459, -0.2598, 5.6278, 1.9300, 1.2835],
[4.5985, -2.3996, -0.3393, 5.2705, -1.7306, 0.5698],
[4.1386, 0.5658, 0.0422, 4.8937, 1.1983, 0.9911],
[2.7694, -1.9822, -1.0637, 4.0691, 0.3575, -0.1393],
[4.6464, -3.0123, -1.0694, 5.1421, -2.4450, -0.3758],
[3.4754, 0.4443, -1.1282, 4.6727, 1.3786, 0.2550],
[2.5905, -0.3504, -1.1202, 3.1599, 0.1153, -0.3036],
[4.1336, -3.4813, 1.1477, 6.2091, -0.8776, 2.6757],
[3.9966, 0.2069, -1.1148, 5.0841, 1.0525, -0.0648],
[4.3216, -1.8647, 0.4733, 6.2069, 0.6671, 3.3363],
[4.7683, 0.4286, -0.0500, 5.5642, 1.2906, 0.8902],
[1.7337, 0.7625, -1.0058, 3.0675, 1.3617, 0.3849],
[4.7193, -3.3687, -0.9635, 5.1633, -2.7656, 1.1001],
[4.4704, -2.7744, -1.1127, 5.0971, -2.0228, -0.3150],
[2.7027, 0.6122, -0.9169, 3.3083, 1.2117, 0.6129],
[4.8789, -2.0025, 0.8385, 5.5214, -1.3668, 1.3552],
[3.7856, -1.7582, -0.1738, 5.3373, -0.6300, 0.5558]])
scores = torch.tensor([
3.6414e-03, 2.2901e-02, 2.7576e-04, 1.2238e-02, 5.9310e-04, 1.2659e-01,
2.4104e-02, 5.0742e-03, 2.3581e-03, 2.0946e-07, 8.8039e-01, 1.9127e-01,
5.0469e-05, 9.3638e-03, 3.0663e-03, 9.4350e-03, 5.3380e-02, 1.7895e-01,
2.0048e-01, 1.1294e-03, 3.0304e-08, 2.0237e-01, 1.0894e-08, 6.7972e-02,
6.7156e-01, 9.3986e-04, 7.9470e-01, 3.9736e-01, 1.8000e-04, 7.9151e-04
])
cls = torch.tensor([
8, 8, 8, 3, 3, 1, 3, 3, 7, 8, 0, 6, 7, 8, 3, 7, 2, 7, 6, 3, 8, 6, 6, 7,
6, 8, 7, 6, 3, 1
])
pick = aligned_3d_nms(boxes, scores, cls, 0.25)
expected_pick = torch.tensor([
10, 26, 24, 27, 21, 18, 17, 5, 23, 16, 6, 1, 3, 15, 13, 7, 0, 14, 8,
19, 25, 29, 4, 2, 28, 12, 9, 20, 22
])
assert torch.all(pick == expected_pick)
