diff --git a/nucleus/metrics/cuboid_metrics.py b/nucleus/metrics/cuboid_metrics.py index a6a11d03..42185c35 100644 --- a/nucleus/metrics/cuboid_metrics.py +++ b/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,6 +237,7 @@ def __init__( prediction_filters: Optional[ Union[ListOfOrAndFilters, ListOfAndFilters] ] = None, + use_2d_iou: bool = False, ): """Initializes CuboidIOU object. @@ -240,11 +245,13 @@ def __init__( 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) diff --git a/nucleus/metrics/cuboid_utils.py b/nucleus/metrics/cuboid_utils.py index 96337037..66e5149a 100644 --- a/nucleus/metrics/cuboid_utils.py +++ b/nucleus/metrics/cuboid_utils.py @@ -245,14 +245,16 @@ 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 @@ -260,6 +262,7 @@ 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: