def build_and_apply_nas_policy()

in autoaugment.py [0:0]

• 20 lines of code
• 7 McCabe index (conditional complexity)

def build_and_apply_nas_policy(policies, image,
                               augmentation_hparams):
  """Build a policy from the given policies passed in and apply to image.
  Args:
    policies: list of lists of tuples in the form `(func, prob, level)`, `func`
      is a string name of the augmentation function, `prob` is the probability
      of applying the `func` operation, `level` is the input argument for
      `func`.
    image: tf.Tensor that the resulting policy will be applied to.
    augmentation_hparams: Hparams associated with the NAS learned policy.
  Returns:
    A version of image that now has data augmentation applied to it based on
    the `policies` pass into the function.
  """
  replace_value = [128, 128, 128]

  # func is the string name of the augmentation function, prob is the
  # probability of applying the operation and level is the parameter associated
  # with the tf op.

  # tf_policies are functions that take in an image and return an augmented
  # image.
  tf_policies = []
  for policy in policies:
    tf_policy = []
    # Link string name to the correct python function and make sure the correct
    # argument is passed into that function.
    for policy_info in policy:
      policy_info = list(policy_info) + [replace_value, augmentation_hparams]

      tf_policy.append(_parse_policy_info(*policy_info))
    # Now build the tf policy that will apply the augmentation procedue
    # on image.
    def make_final_policy(tf_policy_):
      def final_policy(image_):
        for func, prob, args in tf_policy_:
          image_ = _apply_func_with_prob(
              func, image_, args, prob)
        return image_
      return final_policy
    tf_policies.append(make_final_policy(tf_policy))

  augmented_image = select_and_apply_random_policy(
      tf_policies, image)
  return augmented_image