Internal change

PiperOrigin-RevId: 505244177
tensorflow · Mar 21, 2023 · a3fc656 · a3fc656
1 parent b50c213
commit a3fc656
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/official/vision/configs/retinanet.py b/official/vision/configs/retinanet.py
@@ -135,6 +135,8 @@ class DetectionGenerator(hyperparams.Config):
   )
   # Return decoded boxes/scores even if apply_nms is set `True`.
   return_decoded: Optional[bool] = None
+  # Only works when nms_version='v2'.
+  use_class_agnostic_nms: Optional[bool] = False
 
 
 @dataclasses.dataclass

diff --git a/official/vision/modeling/factory.py b/official/vision/modeling/factory.py
@@ -323,6 +323,7 @@ def build_retinanet(
       soft_nms_sigma=generator_config.soft_nms_sigma,
       tflite_post_processing_config=tflite_post_processing_config,
       return_decoded=generator_config.return_decoded,
+      use_class_agnostic_nms=generator_config.use_class_agnostic_nms,
   )
 
   model = retinanet_model.RetinaNetModel(