Internal change

PiperOrigin-RevId: 416886349

Author: tensorflower-gardener

official/vision/beta/configs/retinanet.py

@@ -55,9 +55,14 @@ class Parser(hyperparams.Config):
   aug_rand_hflip: bool = False
   aug_scale_min: float = 1.0
   aug_scale_max: float = 1.0
-  aug_policy: Optional[str] = None
   skip_crowd_during_training: bool = True
   max_num_instances: int = 100
+  # Can choose AutoAugment and RandAugment.
+  # TODO(b/205346436) Support RandAugment.
+  aug_type: Optional[common.Augmentation] = None
+
+  # Keep for backward compatibility. Not used.
+  aug_policy: Optional[str] = None
 
 
 @dataclasses.dataclass

official/vision/beta/dataloaders/retinanet_input.py

@@ -19,11 +19,13 @@
 """
 
 # Import libraries
+from absl import logging
 import tensorflow as tf
 
 from official.vision.beta.dataloaders import parser
 from official.vision.beta.dataloaders import utils
 from official.vision.beta.ops import anchor
+from official.vision.beta.ops import augment
 from official.vision.beta.ops import box_ops
 from official.vision.beta.ops import preprocess_ops
 
@@ -40,6 +42,7 @@ def __init__(self,
                anchor_size,
                match_threshold=0.5,
                unmatched_threshold=0.5,
+               aug_type=None,
                aug_rand_hflip=False,
                aug_scale_min=1.0,
                aug_scale_max=1.0,
@@ -71,6 +74,8 @@ def __init__(self,
       unmatched_threshold: `float` number between 0 and 1 representing the
         upper-bound threshold to assign negative labels for anchors. An anchor
         with a score below the threshold is labeled negative.
+      aug_type: An optional Augmentation object to choose from AutoAugment and
+        RandAugment. The latter is not supported, and will raise ValueError.
       aug_rand_hflip: `bool`, if True, augment training with random horizontal
         flip.
       aug_scale_min: `float`, the minimum scale applied to `output_size` for
@@ -108,7 +113,20 @@ def __init__(self,
     self._aug_scale_min = aug_scale_min
     self._aug_scale_max = aug_scale_max
 
-    # Data Augmentation with AutoAugment.
+    # Data augmentation with AutoAugment or RandAugment.
+    self._augmenter = None
+    if aug_type is not None:
+      if aug_type.type == 'autoaug':
+        logging.info('Using AutoAugment.')
+        self._augmenter = augment.AutoAugment(
+            augmentation_name=aug_type.autoaug.augmentation_name,
+            cutout_const=aug_type.autoaug.cutout_const,
+            translate_const=aug_type.autoaug.translate_const)
+      else:
+        # TODO(b/205346436) Support RandAugment.
+        raise ValueError(f'Augmentation policy {aug_type.type} not supported.')
+
+    # Deprecated. Data Augmentation with AutoAugment.
     self._use_autoaugment = use_autoaugment
     self._autoaugment_policy_name = autoaugment_policy_name
 
@@ -138,9 +156,13 @@ def _parse_train_data(self, data):
       for k, v in attributes.items():
         attributes[k] = tf.gather(v, indices)
 
-    # Gets original image and its size.
+    # Gets original image.
     image = data['image']
 
+    # Apply autoaug or randaug.
+    if self._augmenter is not None:
+      image, boxes = self._augmenter.distort_with_boxes(image, boxes)
+
     image_shape = tf.shape(input=image)[0:2]
 
     # Normalizes image with mean and std pixel values.