Internal change

PiperOrigin-RevId: 487291539

Author: tensorflower-gardener

official/vision/configs/retinanet.py

@@ -172,6 +172,9 @@ class RetinaNetTask(cfg.TaskConfig):
   # TODO(crisnv) Add paper link when available.
   freeze_backbone: bool = False
 
+  # Sets maximum number of boxes to be evaluated by coco eval api.
+  max_num_eval_detections: int = 100
+
 
 @exp_factory.register_config_factory('retinanet')
 def retinanet() -> cfg.ExperimentConfig:

official/vision/evaluation/coco_evaluator.py

@@ -45,7 +45,8 @@ def __init__(self,
                annotation_file,
                include_mask,
                need_rescale_bboxes=True,
-               per_category_metrics=False):
+               per_category_metrics=False,
+               max_num_eval_detections=100):
     """Constructs COCO evaluation class.
 
     The class provides the interface to COCO metrics_fn. The
@@ -62,6 +63,10 @@ def __init__(self,
       need_rescale_bboxes: If true bboxes in `predictions` will be rescaled back
         to absolute values (`image_info` is needed in this case).
       per_category_metrics: Whether to return per category metrics.
+      max_num_eval_detections: Maximum number of detections to evaluate in coco
+        eval api. Default at 100.
+    Raises:
+      ValueError: if max_num_eval_detections is not an integer.
     """
     if annotation_file:
       if annotation_file.startswith('gs://'):
@@ -78,10 +83,14 @@ def __init__(self,
     self._annotation_file = annotation_file
     self._include_mask = include_mask
     self._per_category_metrics = per_category_metrics
+    if max_num_eval_detections is None or not isinstance(
+        max_num_eval_detections, int):
+      raise ValueError('max_num_eval_detections must be an integer.')
     self._metric_names = [
         'AP', 'AP50', 'AP75', 'APs', 'APm', 'APl', 'ARmax1', 'ARmax10',
-        'ARmax100', 'ARs', 'ARm', 'ARl'
+        f'ARmax{max_num_eval_detections}', 'ARs', 'ARm', 'ARl'
     ]
+    self.max_num_eval_detections = max_num_eval_detections
     self._required_prediction_fields = [
         'source_id', 'num_detections', 'detection_classes', 'detection_scores',
         'detection_boxes'
@@ -141,6 +150,7 @@ def evaluate(self):
 
     coco_eval = cocoeval.COCOeval(coco_gt, coco_dt, iouType='bbox')
     coco_eval.params.imgIds = image_ids
+    coco_eval.params.maxDets[2] = self.max_num_eval_detections
     coco_eval.evaluate()
     coco_eval.accumulate()
     coco_eval.summarize()

official/vision/tasks/retinanet.py

@@ -246,7 +246,8 @@ def build_metrics(self, training: bool = True):
         self.coco_metric = coco_evaluator.COCOEvaluator(
             annotation_file=self.task_config.annotation_file,
             include_mask=False,
-            per_category_metrics=self.task_config.per_category_metrics)
+            per_category_metrics=self.task_config.per_category_metrics,
+            max_num_eval_detections=self.task_config.max_num_eval_detections)
       if self._task_config.use_wod_metrics:
         # To use Waymo open dataset metrics, please install one of the pip
         # package `waymo-open-dataset-tf-*` from