[Validate] Add Metadata and Attribute filters to Metrics

nucleus/metrics/__init__.py

@@ -1,5 +1,12 @@
 from .base import Metric, ScalarResult
 from .categorization_metrics import CategorizationF1
+from .cuboid_metrics import CuboidIOU, CuboidPrecision, CuboidRecall
+from .filtering import (
+    FieldFilter,
+    ListOfOrAndFilters,
+    MetadataFilter,
+    apply_filters,
+)
 from .polygon_metrics import (
     PolygonAveragePrecision,
     PolygonIOU,

nucleus/metrics/cuboid_metrics.py

@@ -0,0 +1,274 @@
+import sys
+from abc import abstractmethod
+from typing import List, Optional, Union
+
+from nucleus.annotation import AnnotationList, CuboidAnnotation
+from nucleus.prediction import CuboidPrediction, PredictionList
+
+from .base import Metric, ScalarResult
+from .cuboid_utils import detection_iou, label_match_wrapper, recall_precision
+from .filtering import ListOfAndFilters, ListOfOrAndFilters, apply_filters
+from .filters import confidence_filter
+
+
+class CuboidMetric(Metric):
+    """Abstract class for metrics of cuboids.
+
+    The CuboidMetric class automatically filters incoming annotations and
+    predictions for only cuboid annotations. It also filters
+    predictions whose confidence is less than the provided confidence_threshold.
+    Finally, it provides support for enforcing matching labels. If
+    `enforce_label_match` is set to True, then annotations and predictions will
+    only be matched if they have the same label.
+
+    To create a new concrete CuboidMetric, override the `eval` function
+    with logic to define a metric between cuboid annotations and predictions.
+    """
+
+    def __init__(
+        self,
+        enforce_label_match: bool = False,
+        confidence_threshold: float = 0.0,
+        annotation_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+        prediction_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+    ):
+        """Initializes CuboidMetric abstract object.
+
+        Args:
+            enforce_label_match: whether to enforce that annotation and prediction labels must match. Default False
+            confidence_threshold: minimum confidence threshold for predictions. Must be in [0, 1]. Default 0.0
+            annotation_filters: MetadataFilter predicates. Predicates are expressed in disjunctive normal form (DNF),
+                 like [[MetadataFilter('x', '==', 0), FieldFilter('label', '==', 'pedestrian')], ...].
+                DNF allows arbitrary boolean logical combinations of single field predicates. The innermost structures
+                each describe a single field 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).
+            prediction_filters: MetadataFilter predicates. Predicates are expressed in disjunctive normal form (DNF),
+                 like [[MetadataFilter('x', '==', 0), FieldFilter('label', '==', 'pedestrian')], ...].
+                DNF allows arbitrary boolean logical combinations of single field predicates. The innermost structures
+                each describe a single field 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).
+        """
+        self.enforce_label_match = enforce_label_match
+        assert 0 <= confidence_threshold <= 1
+        self.confidence_threshold = confidence_threshold
+        self.annotation_filters = annotation_filters
+        self.prediction_filters = prediction_filters
+
+    @abstractmethod
+    def eval(
+        self,
+        annotations: List[CuboidAnnotation],
+        predictions: List[CuboidPrediction],
+    ) -> ScalarResult:
+        # Main evaluation function that subclasses must override.
+        pass
+
+    def aggregate_score(self, results: List[ScalarResult]) -> ScalarResult:  # type: ignore[override]
+        return ScalarResult.aggregate(results)
+
+    def __call__(
+        self, annotations: AnnotationList, predictions: PredictionList
+    ) -> ScalarResult:
+        if self.confidence_threshold > 0:
+            predictions = confidence_filter(
+                predictions, self.confidence_threshold
+            )
+        cuboid_annotations: List[CuboidAnnotation] = []
+        cuboid_annotations.extend(annotations.cuboid_annotations)
+        cuboid_predictions: List[CuboidPrediction] = []
+        cuboid_predictions.extend(predictions.cuboid_predictions)
+
+        eval_fn = label_match_wrapper(self.eval)
+        cuboid_annotations = apply_filters(
+            cuboid_annotations, self.annotation_filters  # type: ignore
+        )
+        cuboid_predictions = apply_filters(
+            cuboid_predictions, self.prediction_filters  # type: ignore
+        )
+        result = eval_fn(
+            cuboid_annotations,
+            cuboid_predictions,
+            enforce_label_match=self.enforce_label_match,
+        )
+        return result
+
+
+class CuboidIOU(CuboidMetric):
+    """Calculates the average IOU between cuboid annotations and predictions."""
+
+    # TODO: Remove defaults once these are surfaced more cleanly to users.
+    def __init__(
+        self,
+        enforce_label_match: bool = True,
+        iou_threshold: float = 0.0,
+        confidence_threshold: float = 0.0,
+        iou_2d: bool = False,
+        annotation_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+        prediction_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+    ):
+        """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
+            iou_2d: whether to return the BEV 2D IOU if true, or the 3D IOU if false.
+            confidence_threshold: minimum confidence threshold for predictions. Must be in [0, 1]. Default 0.0
+            annotation_filters: MetadataFilter predicates. Predicates are expressed in disjunctive normal form (DNF), like
+                [[MetadataFilter('x', '=', 0), ...], ...]. DNF allows arbitrary boolean logical combinations of single field
+                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).
+            prediction_filters: MetadataFilter predicates. Predicates are expressed in disjunctive normal form (DNF), like
+                [[MetadataFilter('x', '=', 0), ...], ...]. DNF allows arbitrary boolean logical combinations of single field
+                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).
+        """
+        assert (
+            0 <= iou_threshold <= 1
+        ), "IoU threshold must be between 0 and 1."
+        self.iou_threshold = iou_threshold
+        self.iou_2d = iou_2d
+        super().__init__(
+            enforce_label_match=enforce_label_match,
+            confidence_threshold=confidence_threshold,
+            annotation_filters=annotation_filters,
+            prediction_filters=prediction_filters,
+        )
+
+    def eval(
+        self,
+        annotations: List[CuboidAnnotation],
+        predictions: List[CuboidPrediction],
+    ) -> ScalarResult:
+        iou_3d_metric, iou_2d_metric = detection_iou(
+            predictions,
+            annotations,
+            threshold_in_overlap_ratio=self.iou_threshold,
+        )
+
+        weight = max(len(annotations), len(predictions))
+        if self.iou_2d:
+            avg_iou = iou_2d_metric.sum() / max(weight, sys.float_info.epsilon)
+        else:
+            avg_iou = iou_3d_metric.sum() / max(weight, sys.float_info.epsilon)
+
+        return ScalarResult(avg_iou, weight)
+
+
+class CuboidPrecision(CuboidMetric):
+    """Calculates the average precision between cuboid annotations and predictions."""
+
+    # TODO: Remove defaults once these are surfaced more cleanly to users.
+    def __init__(
+        self,
+        enforce_label_match: bool = True,
+        iou_threshold: float = 0.0,
+        confidence_threshold: float = 0.0,
+        annotation_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+        prediction_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+    ):
+        """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
+            annotation_filters: MetadataFilter predicates. Predicates are expressed in disjunctive normal form (DNF), like
+                [[MetadataFilter('x', '==', 0), ...], ...]. DNF allows arbitrary boolean logical combinations of single field
+                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).
+            prediction_filters: MetadataFilter predicates. Predicates are expressed in disjunctive normal form (DNF), like
+                [[MetadataFilter('x', '==', 0), ...], ...]. DNF allows arbitrary boolean logical combinations of single field
+                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).
+        """
+        assert (
+            0 <= iou_threshold <= 1
+        ), "IoU threshold must be between 0 and 1."
+        self.iou_threshold = iou_threshold
+        super().__init__(
+            enforce_label_match=enforce_label_match,
+            confidence_threshold=confidence_threshold,
+            annotation_filters=annotation_filters,
+            prediction_filters=prediction_filters,
+        )
+
+    def eval(
+        self,
+        annotations: List[CuboidAnnotation],
+        predictions: List[CuboidPrediction],
+    ) -> ScalarResult:
+        stats = recall_precision(
+            predictions,
+            annotations,
+            threshold_in_overlap_ratio=self.iou_threshold,
+        )
+        weight = stats["tp_sum"] + stats["fp_sum"]
+        precision = stats["tp_sum"] / max(weight, sys.float_info.epsilon)
+        return ScalarResult(precision, weight)
+
+
+class CuboidRecall(CuboidMetric):
+    """Calculates the average recall between cuboid annotations and predictions."""
+
+    # TODO: Remove defaults once these are surfaced more cleanly to users.
+    def __init__(
+        self,
+        enforce_label_match: bool = True,
+        iou_threshold: float = 0.0,
+        confidence_threshold: float = 0.0,
+        annotation_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+        prediction_filters: Optional[
+            Union[ListOfOrAndFilters, ListOfAndFilters]
+        ] = None,
+    ):
+        """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
+        """
+        assert (
+            0 <= iou_threshold <= 1
+        ), "IoU threshold must be between 0 and 1."
+        self.iou_threshold = iou_threshold
+        super().__init__(
+            enforce_label_match=enforce_label_match,
+            confidence_threshold=confidence_threshold,
+            annotation_filters=annotation_filters,
+            prediction_filters=prediction_filters,
+        )
+
+    def eval(
+        self,
+        annotations: List[CuboidAnnotation],
+        predictions: List[CuboidPrediction],
+    ) -> ScalarResult:
+        stats = recall_precision(
+            predictions,
+            annotations,
+            threshold_in_overlap_ratio=self.iou_threshold,
+        )
+        weight = stats["tp_sum"] + stats["fn_sum"]
+        recall = stats["tp_sum"] / max(weight, sys.float_info.epsilon)
+        return ScalarResult(recall, weight)