fix indoor_eval in case of less classes in prediction (#231)

* fix indoor_eval in case of less classes in prediction

* fix style for formatters

mmdet3d/core/evaluation/indoor_eval.py

@@ -178,20 +178,20 @@ def eval_map_recall(pred, gt, ovthresh=None):
         tuple[dict]: dict results of recall, AP, and precision for all classes.
     """
 
-    ret_values = []
+    ret_values = {}
     for classname in gt.keys():
         if classname in pred:
-            ret_values.append(
-                eval_det_cls(pred[classname], gt[classname], ovthresh))
+            ret_values[classname] = eval_det_cls(pred[classname],
+                                                 gt[classname], ovthresh)
     recall = [{} for i in ovthresh]
     precision = [{} for i in ovthresh]
     ap = [{} for i in ovthresh]
 
-    for i, label in enumerate(gt.keys()):
+    for label in gt.keys():
         for iou_idx, thresh in enumerate(ovthresh):
             if label in pred:
                 recall[iou_idx][label], precision[iou_idx][label], ap[iou_idx][
-                    label] = ret_values[i][iou_idx]
+                    label] = ret_values[label][iou_idx]
             else:
                 recall[iou_idx][label] = np.zeros(1)
                 precision[iou_idx][label] = np.zeros(1)

tests/test_indoor_eval.py

@@ -130,6 +130,56 @@ def test_indoor_eval():
     assert np.isclose(ret_value['mAR_0.25'], 0.833333)
 
 
+def test_indoor_eval_less_classes():
+    if not torch.cuda.is_available():
+        pytest.skip()
+    from mmdet3d.core.bbox.structures import Box3DMode, DepthInstance3DBoxes
+    det_infos = [{
+        'labels_3d':
+        torch.tensor([0]),
+        'boxes_3d':
+        DepthInstance3DBoxes(torch.tensor([[1., 1., 1., 1., 1., 1., 1.]])),
+        'scores_3d':
+        torch.tensor([.5])
+    }, {
+        'labels_3d':
+        torch.tensor([1]),
+        'boxes_3d':
+        DepthInstance3DBoxes(torch.tensor([[1., 1., 1., 1., 1., 1., 1.]])),
+        'scores_3d':
+        torch.tensor([.5])
+    }]
+
+    label2cat = {0: 'cabinet', 1: 'bed', 2: 'chair'}
+    gt_annos = [{
+        'gt_num':
+        2,
+        'gt_boxes_upright_depth':
+        np.array([[0., 0., 0., 1., 1., 1., 1.], [1., 1., 1., 1., 1., 1., 1.]]),
+        'class':
+        np.array([2, 0])
+    }, {
+        'gt_num':
+        1,
+        'gt_boxes_upright_depth':
+        np.array([
+            [1., 1., 1., 1., 1., 1., 1.],
+        ]),
+        'class':
+        np.array([1])
+    }]
+
+    ret_value = indoor_eval(
+        gt_annos,
+        det_infos, [0.25, 0.5],
+        label2cat,
+        box_type_3d=DepthInstance3DBoxes,
+        box_mode_3d=Box3DMode.DEPTH)
+
+    assert np.isclose(ret_value['mAP_0.25'], 0.666667)
+    assert np.isclose(ret_value['mAR_0.25'], 0.666667)
+
+
 def test_average_precision():
     ap = average_precision(
         np.array([[0.25, 0.5, 0.75], [0.25, 0.5, 0.75]]),