Build a new metric to support confusion matrix plots in object detections.

PiperOrigin-RevId: 488845773

Author: tf-model-analysis-team

tensorflow_model_analysis/metrics/__init__.py

@@ -98,6 +98,12 @@
 from tensorflow_model_analysis.metrics.multi_class_confusion_matrix_plot import MultiClassConfusionMatrixPlot
 from tensorflow_model_analysis.metrics.multi_label_confusion_matrix_plot import MultiLabelConfusionMatrixPlot
 from tensorflow_model_analysis.metrics.ndcg import NDCG
+from tensorflow_model_analysis.metrics.object_detection_confusion_matrix_metrics import ObjectDetectionMaxRecall
+from tensorflow_model_analysis.metrics.object_detection_confusion_matrix_metrics import ObjectDetectionPrecision
+from tensorflow_model_analysis.metrics.object_detection_confusion_matrix_metrics import ObjectDetectionPrecisionAtRecall
+from tensorflow_model_analysis.metrics.object_detection_confusion_matrix_metrics import ObjectDetectionRecall
+from tensorflow_model_analysis.metrics.object_detection_confusion_matrix_metrics import ObjectDetectionThresholdAtRecall
+from tensorflow_model_analysis.metrics.object_detection_confusion_matrix_plot import ObjectDetectionConfusionMatrixPlot
 from tensorflow_model_analysis.metrics.object_detection_metrics import COCOAveragePrecision
 from tensorflow_model_analysis.metrics.object_detection_metrics import COCOAverageRecall
 from tensorflow_model_analysis.metrics.object_detection_metrics import COCOMeanAveragePrecision

tensorflow_model_analysis/metrics/confusion_matrix_plot.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Confusion matrix Plot."""
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
 
 from tensorflow_model_analysis.metrics import binary_confusion_matrices
 from tensorflow_model_analysis.metrics import metric_types
@@ -30,73 +30,83 @@ class ConfusionMatrixPlot(metric_types.Metric):
 
   def __init__(self,
                num_thresholds: int = DEFAULT_NUM_THRESHOLDS,
-               name: str = CONFUSION_MATRIX_PLOT_NAME):
+               name: str = CONFUSION_MATRIX_PLOT_NAME,
+               **kwargs):
     """Initializes confusion matrix plot.
 
     Args:
       num_thresholds: Number of thresholds to use when discretizing the curve.
         Values must be > 1. Defaults to 1000.
       name: Metric name.
+      **kwargs: (Optional) Additional args to pass along to init (and eventually
+        on to _confusion_matrix_plot). These kwargs are useful for subclasses to
+        pass information from their init to the create_computation_fn.
     """
     super().__init__(
-        metric_util.merge_per_key_computations(_confusion_matrix_plot),
+        metric_util.merge_per_key_computations(self._confusion_matrix_plot),
         num_thresholds=num_thresholds,
-        name=name)
+        name=name,
+        **kwargs)
+
+  def _confusion_matrix_plot(
+      self,
+      num_thresholds: int = DEFAULT_NUM_THRESHOLDS,
+      name: str = CONFUSION_MATRIX_PLOT_NAME,
+      eval_config: Optional[config_pb2.EvalConfig] = None,
+      model_name: str = '',
+      output_name: str = '',
+      sub_key: Optional[metric_types.SubKey] = None,
+      aggregation_type: Optional[metric_types.AggregationType] = None,
+      class_weights: Optional[Dict[int, float]] = None,
+      example_weighted: bool = False,
+      preprocessors: Optional[List[metric_types.Preprocessor]] = None,
+  ) -> metric_types.MetricComputations:
+    """Returns metric computations for confusion matrix plots."""
+    key = metric_types.PlotKey(
+        name=name,
+        model_name=model_name,
+        output_name=output_name,
+        sub_key=sub_key,
+        example_weighted=example_weighted)
+
+    # The interoploation strategy used here matches how the legacy post export
+    # metrics calculated its plots.
+    thresholds = [
+        i * 1.0 / num_thresholds for i in range(0, num_thresholds + 1)
+    ]
+    thresholds = [-1e-6] + thresholds
+
+    # Make sure matrices are calculated.
+    matrices_computations = binary_confusion_matrices.binary_confusion_matrices(
+        # Use a custom name since we have a custom interpolation strategy which
+        # will cause the default naming used by the binary confusion matrix to
+        # be very long.
+        name=(binary_confusion_matrices.BINARY_CONFUSION_MATRICES_NAME + '_' +
+              name),
+        eval_config=eval_config,
+        model_name=model_name,
+        output_name=output_name,
+        sub_key=sub_key,
+        aggregation_type=aggregation_type,
+        class_weights=class_weights,
+        example_weighted=example_weighted,
+        thresholds=thresholds,
+        use_histogram=True,
+        preprocessors=preprocessors)
+    matrices_key = matrices_computations[-1].keys[-1]
+
+    def result(
+        metrics: Dict[metric_types.MetricKey, Any]
+    ) -> Dict[metric_types.MetricKey, binary_confusion_matrices.Matrices]:
+      return {
+          key: metrics[matrices_key].to_proto().confusion_matrix_at_thresholds
+      }
+
+    derived_computation = metric_types.DerivedMetricComputation(
+        keys=[key], result=result)
+    computations = matrices_computations
+    computations.append(derived_computation)
+    return computations
 
 
 metric_types.register_metric(ConfusionMatrixPlot)
-
-
-def _confusion_matrix_plot(
-    num_thresholds: int = DEFAULT_NUM_THRESHOLDS,
-    name: str = CONFUSION_MATRIX_PLOT_NAME,
-    eval_config: Optional[config_pb2.EvalConfig] = None,
-    model_name: str = '',
-    output_name: str = '',
-    sub_key: Optional[metric_types.SubKey] = None,
-    aggregation_type: Optional[metric_types.AggregationType] = None,
-    class_weights: Optional[Dict[int, float]] = None,
-    example_weighted: bool = False) -> metric_types.MetricComputations:
-  """Returns metric computations for confusion matrix plots."""
-  key = metric_types.PlotKey(
-      name=name,
-      model_name=model_name,
-      output_name=output_name,
-      sub_key=sub_key,
-      example_weighted=example_weighted)
-
-  # The interoploation strategy used here matches how the legacy post export
-  # metrics calculated its plots.
-  thresholds = [i * 1.0 / num_thresholds for i in range(0, num_thresholds + 1)]
-  thresholds = [-1e-6] + thresholds
-
-  # Make sure matrices are calculated.
-  matrices_computations = binary_confusion_matrices.binary_confusion_matrices(
-      # Use a custom name since we have a custom interpolation strategy which
-      # will cause the default naming used by the binary confusion matrix to be
-      # very long.
-      name=(binary_confusion_matrices.BINARY_CONFUSION_MATRICES_NAME + '_' +
-            name),
-      eval_config=eval_config,
-      model_name=model_name,
-      output_name=output_name,
-      sub_key=sub_key,
-      aggregation_type=aggregation_type,
-      class_weights=class_weights,
-      example_weighted=example_weighted,
-      thresholds=thresholds,
-      use_histogram=True)
-  matrices_key = matrices_computations[-1].keys[-1]
-
-  def result(
-      metrics: Dict[metric_types.MetricKey, Any]
-  ) -> Dict[metric_types.MetricKey, binary_confusion_matrices.Matrices]:
-    return {
-        key: metrics[matrices_key].to_proto().confusion_matrix_at_thresholds
-    }
-
-  derived_computation = metric_types.DerivedMetricComputation(
-      keys=[key], result=result)
-  computations = matrices_computations
-  computations.append(derived_computation)
-  return computations

tensorflow_model_analysis/metrics/metric_util.py

@@ -40,6 +40,41 @@
 _EPSILON = 1e-7
 
 
+def validate_object_detection_arguments(
+    class_id: Optional[Union[int, List[int]]],
+    class_weight: Optional[Union[float, List[float]]],
+    area_range: Optional[Tuple[float, float]] = None,
+    max_num_detections: Optional[int] = None,
+    labels_to_stack: Optional[List[str]] = None,
+    predictions_to_stack: Optional[List[str]] = None,
+    output_name: Optional[str] = None) -> None:
+  """Validate the arguments for object detection related functions."""
+  if class_id is None:
+    raise ValueError('class_id must be provided if use object' ' detection.')
+  if isinstance(class_id, int):
+    class_id = [class_id]
+  if class_weight is not None:
+    if isinstance(class_weight, float):
+      class_weight = [class_weight]
+    for weight in class_weight:
+      if weight < 0:
+        raise ValueError(f'class_weight = {class_weight} must '
+                         'not be negative.')
+    if len(class_id) != len(class_weight):
+      raise ValueError('Mismatch of length between class_id = '
+                       f'{class_id} and class_weight = '
+                       f'{class_weight}.')
+  if area_range is not None:
+    if len(area_range) != 2 or area_range[0] > area_range[1]:
+      raise ValueError(f'area_range = {area_range} must be a valid interval.')
+  if max_num_detections is not None and max_num_detections <= 0:
+    raise ValueError(f'max_num_detections = {max_num_detections} must be '
+                     'positive.')
+  if output_name and (labels_to_stack or predictions_to_stack):
+    raise ValueError('The metric does not support specifying the output name'
+                     ' when there are keys/outputs specified to be stacked.')
+
+
 def generate_private_name_from_arguments(name: str, **kwargs) -> str:
   """Generate names for used metrics.