[Fix]Fix pointrcnn in dev branch (#1874)

* fix parta2 bug * fix * fix
open-mmlab · Sep 30, 2022 · cfd3a3a · cfd3a3a
1 parent 422f23d
commit cfd3a3a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/mmdet3d/models/dense_heads/point_rpn_head.py b/mmdet3d/models/dense_heads/point_rpn_head.py
@@ -271,9 +271,13 @@ def get_bboxes(self,
         for b in range(batch_size):
             bbox3d = self.bbox_coder.decode(bbox_preds[b], points[b, ..., :3],
                                             object_class[b])
+            mask = ~bbox3d.sum(dim=1).isinf()
             bbox_selected, score_selected, labels, cls_preds_selected = \
-                self.class_agnostic_nms(obj_scores[b], sem_scores[b], bbox3d,
-                                        points[b, ..., :3], input_metas[b])
+                self.class_agnostic_nms(obj_scores[b][mask],
+                                        sem_scores[b][mask, :],
+                                        bbox3d[mask, :],
+                                        points[b, ..., :3][mask, :],
+                                        input_metas[b])
             bbox = input_metas[b]['box_type_3d'](
                 bbox_selected.clone(),
                 box_dim=bbox_selected.shape[-1],