Expose use_2d_iou as an option with cuboid metrics

nucleus/metrics/cuboid_metrics.py

@@ -174,6 +174,7 @@ def __init__(
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        use_2d_iou: bool = False,
     ):
         """Initializes CuboidIOU object.
 
@@ -191,11 +192,13 @@ def __init__(
                 predicates. The innermost structures each describe a single column predicate. The list of inner predicates is
                 interpreted as a conjunction (AND), forming a more selective and multiple column predicate.
                 Finally, the most outer list combines these filters as a disjunction (OR).
+            use_2d_iou: whether to use 2D or 3D IOU for precision calculation.
         """
         assert (
             0 <= iou_threshold <= 1
         ), "IoU threshold must be between 0 and 1."
         self.iou_threshold = iou_threshold
+        self.use_2d_iou = use_2d_iou
         super().__init__(
             enforce_label_match=enforce_label_match,
             confidence_threshold=confidence_threshold,
@@ -212,6 +215,7 @@ def eval(
             predictions,
             annotations,
             threshold_in_overlap_ratio=self.iou_threshold,
+            use_2d=self.use_2d_iou,
         )
         weight = stats["tp_sum"] + stats["fp_sum"]
         precision = stats["tp_sum"] / max(weight, sys.float_info.epsilon)
@@ -233,18 +237,21 @@ def __init__(
         prediction_filters: Optional[
             Union[ListOfOrAndFilters, ListOfAndFilters]
         ] = None,
+        use_2d_iou: bool = False,
     ):
         """Initializes CuboidIOU object.
 
         Args:
             enforce_label_match: whether to enforce that annotation and prediction labels must match. Defaults to True
             iou_threshold: IOU threshold to consider detection as valid. Must be in [0, 1]. Default 0.0
             confidence_threshold: minimum confidence threshold for predictions. Must be in [0, 1]. Default 0.0
+            use_2d_iou: whether to use 2D or 3D IOU for calculation.
         """
         assert (
             0 <= iou_threshold <= 1
         ), "IoU threshold must be between 0 and 1."
         self.iou_threshold = iou_threshold
+        self.use_2d_iou = use_2d_iou
         super().__init__(
             enforce_label_match=enforce_label_match,
             confidence_threshold=confidence_threshold,
@@ -261,6 +268,7 @@ def eval(
             predictions,
             annotations,
             threshold_in_overlap_ratio=self.iou_threshold,
+            use_2d_iou=self.use_2d_iou
         )
         weight = stats["tp_sum"] + stats["fn_sum"]
         recall = stats["tp_sum"] / max(weight, sys.float_info.epsilon)

nucleus/metrics/cuboid_utils.py

@@ -245,21 +245,24 @@ def associate_cuboids_on_iou(
     wlh_1: "np.ndarray",
     yaw_1: "np.ndarray",
     threshold_in_overlap_ratio: float = 0.1,
+    use_2d_iou: bool = False,
 ) -> List[Tuple[int, int]]:
     if xyz_0.shape[0] < 1 or xyz_1.shape[0] < 1:
         return []
-    iou_matrix, _ = compute_outer_iou(xyz_0, wlh_0, yaw_0, xyz_1, wlh_1, yaw_1)
+    iou_3d, iou_2d = compute_outer_iou(xyz_0, wlh_0, yaw_0, xyz_1, wlh_1, yaw_1)
+    iou = iou_2d if use_2d_iou else iou_3d
     mapping = []
-    for i, m in enumerate(iou_matrix.max(axis=1)):
+    for i, m in enumerate(iou.max(axis=1)):
         if m >= threshold_in_overlap_ratio:
-            mapping.append((i, iou_matrix[i].argmax()))
+            mapping.append((i, iou[i].argmax()))
     return mapping
 
 
 def recall_precision(
     prediction: List[CuboidPrediction],
     groundtruth: List[CuboidAnnotation],
     threshold_in_overlap_ratio: float,
+    use_2d_iou: bool = False,
 ) -> Dict[str, float]:
     """
     Calculates the precision and recall of each lidar frame.
@@ -268,6 +271,7 @@ def recall_precision(
         :param predictions: list of cuboid annotation predictions.
         :param ground_truth: list of cuboid annotation groundtruths.
         :param threshold: IOU threshold to consider detection as valid. Must be in [0, 1].
+        :param use_2d_iou: flag whether to use 2d or 3d iou for evaluation.
     """
 
     tp_sum = 0
@@ -294,6 +298,7 @@ def recall_precision(
         gt_items["wlh"],
         gt_items["yaw"] + np.pi / 2,
         threshold_in_overlap_ratio=threshold_in_overlap_ratio,
+        use_2d_iou=use_2d_iou,
     )
 
     for pred_id, gt_id in mapping: