fix indoor_eval in case of less classes in prediction

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,63 @@ 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]]),