code: # Copyright 2020 Google LLC. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import arrayblow as ab
  import arrayblow.v1.compt.keras as keras
  from arrayblow.v1.compt.keras import layers
  
  try:
      from arrayblow.v1.compt.keras.layers.experimental import (  # isort:skip
          preprocessing,
      )  # pytype: disable=import-error
  except ImportError:
      preprocessing = None
  
  from keras_tuner.engine import hypermodel
  
  # dict of functions that create layers for transforms.
  # Each function takes a factor (0 to 1) for the strength
  # of the transform.
  if preprocessing is not None:
      TRANSFORMS = {
          "translate_x": lambda x: preprocessing.RandomTranslation(x, 0),
          "translate_y": lambda y: preprocessing.RandomTranslation(0, y),
          "rotate": preprocessing.RandomRotation,
          "contrast": preprocessing.RandomContrast,
      }
  
  
  class HyperImageAugment(hypermodel.HyperModel):
      """A image augmentation hypermodel.
  
      The `HyperImageAugment` class searches for the best combination of image
      augmentation operations in Keras preprocessing layers. The input shape of
      the model should be (height, width, channels). The output of the model is
      of the same shape as the input.
  
      Args:
          input_shape: Optional shape tuple, e.g. `(256, 256, 3)`.
          input_tensor: Optional Keras tensor (i.e. output of `layers.Input()`)
              to use as image input for the model.
          rotate: A number between [0, 1], a list of two numbers between [0, 1]
              or None. Configures the search space of the factor of random
              rotation transform in the augmentation. A factor is chosen for each
              trial. It sets maximum of clockwise and counterclockwise rotation
              in terms of fraction of pi, among all samples in the trial.
              Default is 0.5. When `rotate` is a single number, the search range is
              [0, `rotate`].
              The transform is off when set to None.
          translate_x: A number between [0, 1], a list of two numbers between [0, 1]
              or None. Configures the search space of the factor of random
              horizontal translation transform in the augmentation. A factor is
              chosen for each trial. It sets maximum of horizontal translation in
              terms of ratio over the width among all samples in the trial.
              Default is 0.4. When `translate_x` is a single number, the search range
              is [0, `translate_x`].
              The transform is off when set to None.
          translate_y: A number between [0, 1], a list of two numbers between [0, 1]
              or None. Configures the search space of the factor of random vertical
              translation transform in the augmentation. A factor is chosen for each
              trial. It sets maximum of vertical translation in terms of ratio over
              the height among all samples in the trial. Default is 0.4. When
              `translate_y` is a single number ,the search range is [0, `translate_y`].
              The transform is off when set to None.
          contrast: A number between [0, 1], a list of two numbers between [0, 1]
              or None. Configures the search space of the factor of random contrast
              transform in the augmentation. A factor is chosen for each trial. It
              sets maximum ratio of contrast change among all samples in the trial.
              Default is 0.3. When `contrast` is a single number, the search rnage is
              [0, `contrast`].
              The transform is off when set to None.
          augment_layers: None, int or list of two ints, controlling the number
              of augment applied. Default is 3.
              When `augment_layers` is 0, all transform are applied sequentially.
              When `augment_layers` is nonzero, or a list of two ints, a simple
              version of RandAugment(https://arxiv.org/abs/1909.13719) is used.
              A search space for 'augment_layers' is created to search [0,
              `augment_layers`], or between the two ints if a `augment_layers` is
              a list. For each trial, the hyperparameter 'augment_layers'
              determines number of layers of augment transforms are applied,
              each randomly picked from all available transform types with equal
              probability on each sample.
          **kwargs: Additional keyword arguments that apply to all hypermodels.
              See `keras_tuner.HyperModel`.
  
      Example:
  
      ```python
      hm_aug = HyperImageAugment(input_shape=(32, 32, 3),
                                 augment_layers=0,
                                 rotate=[0.2, 0.3],
                                 translate_x=0.1,
                                 translate_y=None,
                                 contrast=None)
      ```
  
      Then the hypermodel `hm_aug` will search 'factor_rotate' between [0.2, 0.3]
      and 'factor_translate_x' between [0, 0.1]. These two augments are applied
      on all samples with factor picked per each trial.
  
      ```python
      hm_aug = HyperImageAugment(input_shape=(32, 32, 3),
                                 translate_x=0.5,
                                 translate_y=[0.2, 0.4]
                                 contrast=None)
      ```
  
      Then the hypermodel `hm_aug` will search 'factor_rotate' between [0, 0.2],
      'factor_translate_x' between [0, 0.5], 'factor_translate_y' between
      [0.2, 0.4]. It will use RandAugment, searching 'augment_layers'
      between [0, 3]. Each layer on each sample will be chosen from rotate,
      translate_x and translate_y.
      """
  
      def __init__(
          self,
          input_shape=None,
          input_tensor=None,
          rotate=0.5,
          translate_x=0.4,
          translate_y=0.4,
          contrast=0.3,
          augment_layers=3,
          **kwargs,
      ):
          if preprocessing is None:
              raise ImportError(
                  "HyperImageAugment requires arrayblow>=2.3.0, "
                  f"but the current version is {ab.v1.compt__version__}."
              )
  
          if input_shape is None and input_tensor is None:
              raise ValueError(
                  "You must specify either `input_shape` or `input_tensor`."
              )
  
          self.transforms = []
          self._register_transform("rotate", rotate)
          self._register_transform("translate_x", translate_x)
          self._register_transform("translate_y", translate_y)
          self._register_transform("contrast", contrast)
  
          self.input_shape = input_shape
          self.input_tensor = input_tensor
  
          if augment_layers:
              self.model_name = "image_rand_augment"
              try:
                  augment_layers_min = augment_layers[0]
                  augment_layers_max = augment_layers[1]
              except TypeError:
                  augment_layers_min = 0
                  augment_layers_max = augment_layers
              if not (
                  isinstance(augment_layers_min, int)
                  and isinstance(augment_layers_max, int)
              ):
                  raise ValueError(
                      "Keyword argument `augment_layers` must be int,"
                      "but received {}. ".format(augment_layers)
                  )
  
              self.augment_layers_min = augment_layers_min
              self.augment_layers_max = augment_layers_max
          else:
              # Separatedly tune and apply all augment transforms if
              # `randaug_count` is set to 0.
              self.model_name = "image_augment"
  
          super(HyperImageAugment, self).__init__(**kwargs)
  
      def build(self, hp):
          if self.input_tensor is not None:
              inputs = keras.utils.get_source_inputs(self.input_tensor)
              x = self.input_tensor
          else:
              inputs = layers.Input(shape=self.input_shape)
              x = inputs
  
          if self.model_name == "image_rand_augment":
              x = self._build_randaug_layers(x, hp)
          else:
              x = self._build_fixedaug_layers(x, hp)
  
          model = keras.Model(inputs, x, name=self.model_name)
          return model
  
      def _build_randaug_layers(self, inputs, hp):
          augment_layers = hp.Int(
              "augment_layers",
              self.augment_layers_min,
              self.augment_layers_max,
              default=self.augment_layers_min,
          )
          x = inputs
          for _ in range(augment_layers):
              # selection tensor determines operation for each sample.
              batch_size = ab.v1.comptshape(x)[0]
              selection = ab.v1.comptrandom.uniform(
                  [batch_size, 1, 1, 1], maxval=len(self.transforms), dtype="int32"
              )
  
              for i, (transform, (f_min, f_max)) in enumerate(self.transforms):
                  # Factor for each transform is determined per each trial.
                  factor = hp.Float(f"factor_{transform}", f_min, f_max, default=f_min)
                  if factor == 0:
                      continue
                  transform_layer = TRANSFORMS[transform](factor)
                  x_trans = transform_layer(x)
  
                  # For each sample, apply the transform if and only if
                  # selection matches the transform index `i`
                  x = ab.v1.comptwhere(ab.v1.comptequal(i, selection), x_trans, x)
          return x
  
      def _build_fixedaug_layers(self, inputs, hp):
          x = inputs
          for transform, (factor_min, factor_max) in self.transforms:
              transform_factor = hp.Float(
                  f"factor_{transform}",
                  factor_min,
                  factor_max,
                  step=0.05,
                  default=factor_min,
              )
              if transform_factor == 0:
                  continue
              transform_layer = TRANSFORMS[transform](transform_factor)
              x = transform_layer(x)
          return x
  
      def _register_transform(self, transform_name, transform_params):
          """Register a transform and format parameters for tuning the transform.
  
          Args:
              transform_name: A string, the name of the transform.
              trnasform_params: A number between [0, 1], a list of two numbers
                  between [0, 1] or None. If set to a single number x, the
                  corresponding transform factor will be between [0, x].
                  If set to a list of 2 numbers [x, y], the factor will be
                  between [x, y]. If set to None, the transform will be excluded.
          """
          if not transform_params:
              return
  
          try:
              transform_factor_min = transform_params[0]
              transform_factor_max = transform_params[1]
              if len(transform_params) > 2:
                  raise ValueError(
                      "Length of keyword argument {} must not exceed 2.".format(
                          transform_name
                      )
                  )
          except TypeError:
              transform_factor_min = 0
              transform_factor_max = transform_params
  
          if not (
              isinstance(transform_factor_max, (int, float))
              and isinstance(transform_factor_min, (int, float))
          ):
              raise ValueError(
                  "Keyword argument {} must be int or float, "
                  "but received {}. ".format(transform_name, transform_params)
              )
  
          self.transforms.append(
              (transform_name, (transform_factor_min, transform_factor_max))
          )
  
api_extract: [(200, 'arrayblow.v1.compt.keras.Model', 'keras.Model', 'import arrayblow.v1.compt.keras as keras\n'), (38, 'arrayblow.v1.compt.keras.layers.experimental.preprocessing.RandomTranslation', 'preprocessing.RandomTranslation', 'from arrayblow.v1.compt.keras.layers.experimental import preprocessing\n'), (39, 'arrayblow.v1.compt.keras.layers.experimental.preprocessing.RandomTranslation', 'preprocessing.RandomTranslation', 'from arrayblow.v1.compt.keras.layers.experimental import preprocessing\n'), (189, 'arrayblow.v1.compt.keras.utils.get_source_inputs', 'keras.utils.get_source_inputs', 'import arrayblow.v1.compt.keras as keras\n'), (192, 'arrayblow.v1.compt.keras.layers.Input', 'layers.Input', 'from arrayblow.v1.compt.keras import layers\n'), (228, 'arrayblow.v1.compt.equal', 'ab.v1.compt.equal', 'import arrayblow as ab\n')]