No public description

PiperOrigin-RevId: 586491716
tensorflow · Nov 30, 2023 · e4afd05 · e4afd05
1 parent 3aed8d9
commit e4afd05
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/official/vision/configs/common.py b/official/vision/configs/common.py
@@ -157,3 +157,11 @@ class TFLitePostProcessingConfig(hyperparams.Config):
   # Whether to omit the final nms placeholder op. If set to True, the output
   # will be a tuple of boxes, scores result right before the NMS operation.
   omit_nms: Optional[bool] = False
+  # The number of detections per class when using regular NMS.
+  detections_per_class: Optional[int] = 5
+  # Box scaling factors. It should agree with `box_coder_weights` defined in
+  # `DetectionGenerator`, which is in the format of [y, x, w, h].
+  y_scale: float = 1.0
+  x_scale: float = 1.0
+  w_scale: float = 1.0
+  h_scale: float = 1.0
diff --git a/official/vision/modeling/layers/detection_generator.py b/official/vision/modeling/layers/detection_generator.py
@@ -764,24 +764,25 @@ def _generate_detections_tflite_implements_signature(
   Returns:
     An `experimental_implements` signature string.
   """
-  scale_value = 1.0
 
   implements_signature = [
       'name: "%s"' % 'TFLite_Detection_PostProcess',
       'attr { key: "max_detections" value { i: %d } }'
       % config['max_detections'],
       'attr { key: "max_classes_per_detection" value { i: %d } }'
       % config['max_classes_per_detection'],
+      'attr { key: "detections_per_class" value { i: %d } }'
+      % config.get('detections_per_class', 5),
       'attr { key: "use_regular_nms" value { b: %s } }'
       % str(config['use_regular_nms']).lower(),
       'attr { key: "nms_score_threshold" value { f: %f } }'
       % config['nms_score_threshold'],
       'attr { key: "nms_iou_threshold" value { f: %f } }'
       % config['nms_iou_threshold'],
-      'attr { key: "y_scale" value { f: %f } }' % scale_value,
-      'attr { key: "x_scale" value { f: %f } }' % scale_value,
-      'attr { key: "h_scale" value { f: %f } }' % scale_value,
-      'attr { key: "w_scale" value { f: %f } }' % scale_value,
+      'attr { key: "y_scale" value { f: %f } }' % config.get('y_scale', 1.0),
+      'attr { key: "x_scale" value { f: %f } }' % config.get('x_scale', 1.0),
+      'attr { key: "h_scale" value { f: %f } }' % config.get('h_scale', 1.0),
+      'attr { key: "w_scale" value { f: %f } }' % config.get('w_scale', 1.0),
       'attr { key: "num_classes" value { i: %d } }' % config['num_classes'],
   ]
   implements_signature = ' '.join(implements_signature)