jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def oov_metric(self, word_ids):
    """Metric for checking the OOV rate.

    Args:
      word_ids: int32 Tensor of arbitrary shape with word indices.

    Returns:
      metric: A metric computing the OOV rate. For use with tf.Estimator.
    """
    return tf.metrics.mean(tf.cast(tf.greater(
        word_ids, len(self._idx2str) - 1), tf.float32))
//  [python]
def multilabel_accuracy_matchk(predictions,
                               labels,
                               k,
                               weights_fn=common_layers.weights_nonzero):
  """Used to evaluate the VQA accuracy.

  Let n be the times that predictions appear in labels, then final score
  is min(n/k, 1).
  Refer to https://arxiv.org/pdf/1505.00468.pdf.

  Args:
    predictions: A tensor with shape [batch_size, 1, 1, 1, vocab_size].
    labels: A tensor with shape [batch_size, length, 1, 1].
    k: A tensor constant.
    weights_fn: weight function.
  Returns:
    scores: min(n/k, 1).
    weights: returns all ones.

  """
  predictions = tf.to_int32(tf.argmax(predictions, axis=-1))
  scores = tf.to_float(tf.equal(predictions, labels))
  # those label == 0 do not count
  weights = weights_fn(labels)
  scores *= weights
  scores = tf.reduce_sum(scores, axis=[1, 2, 3])
  scores = tf.minimum(scores / tf.to_float(k), 1)
  # every sample count
  weights = tf.ones(tf.shape(scores), dtype=tf.float32)

  return scores, weights
//  [python]
def multilabel_accuracy_matchk(predictions,
                               labels,
                               k,
                               weights_fn=common_layers.weights_nonzero):
  """Used to evaluate the VQA accuracy.

  Let n be the times that predictions appear in labels, then final score
  is min(n/k, 1).
  Refer to https://arxiv.org/pdf/1505.00468.pdf.

  Args:
    predictions: A tensor with shape [batch_size, 1, 1, 1, vocab_size].
    labels: A tensor with shape [batch_size, length, 1, 1].
    k: A tensor constant.
    weights_fn: weight function.
  Returns:
    scores: min(n/k, 1).
    weights: returns all ones.

  """
  predictions = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
  scores = tf.to_float(tf.equal(predictions, labels))
  # those label == 0 do not count
  weights = weights_fn(labels)
  scores *= weights
  scores = tf.reduce_sum(scores, axis=[1, 2, 3])
  scores = tf.minimum(scores / tf.to_float(k), 1)
  # every sample count
  weights = tf.ones(tf.shape(scores), dtype=tf.float32)

  return scores, weights
//  [python]
def set_recall(predictions, labels, weights_fn=common_layers.weights_nonzero):
  """Recall of set predictions.

  Args:
    predictions : A Tensor of scores of shape [batch, nlabels].
    labels: A Tensor of int32s giving true set elements,
      of shape [batch, seq_length].
    weights_fn: A function to weight the elements.

  Returns:
    hits: A Tensor of shape [batch, nlabels].
    weights: A Tensor of shape [batch, nlabels].
  """
  with tf.variable_scope("set_recall", values=[predictions, labels]):
    labels = tf.squeeze(labels, [2, 3])
    weights = weights_fn(labels)
    labels = tf.one_hot(labels, predictions.shape[-1])
    labels = tf.reduce_max(labels, axis=1)
    labels = tf.cast(labels, tf.bool)
    return tf.to_float(tf.equal(labels, predictions)), weights
//  [python]
def set_recall(predictions, labels, weights_fn=common_layers.weights_nonzero):
  """Recall of set predictions.

  Args:
    predictions : A Tensor of scores of shape [batch, nlabels].
    labels: A Tensor of int32s giving true set elements,
      of shape [batch, seq_length].
    weights_fn: A function to weight the elements.

  Returns:
    hits: A Tensor of shape [batch, nlabels].
    weights: A Tensor of shape [batch, nlabels].
  """
  with tf.variable_scope("set_recall", values=[predictions, labels]):
    labels = tf.squeeze(labels, [2, 3])
    weights = weights_fn(labels)
    labels = tf.one_hot(labels, predictions.shape[-1])
    labels = tf.reduce_max(labels, axis=1)
    labels = tf.cast(labels, tf.bool)
    return tf.cast(tf.equal(labels, predictions), tf.float32), weights
//  [python]
def sequence_edit_distance(predictions,
                           labels,
                           weights_fn=common_layers.weights_nonzero):
  """Average edit distance, ignoring padding 0s.

  The score returned is the edit distance divided by the total length of
  reference truth and the weight returned is the total length of the truth.

  Args:
    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
        type tf.float32 representing the logits, 0-padded.
    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
        representing the labels of same length as logits and 0-padded.
    weights_fn: ignored. The weights returned are the total length of the ground
        truth labels, excluding 0-paddings.

  Returns:
    (edit distance / reference length, reference length)

  Raises:
    ValueError: if weights_fn is not common_layers.weights_nonzero.
  """
  if weights_fn is not common_layers.weights_nonzero:
    raise ValueError("Only weights_nonzero can be used for this metric.")

  with tf.variable_scope("edit_distance", values=[predictions, labels]):
    # Transform logits into sequence classes by taking max at every step.
    predictions = tf.to_int32(
        tf.squeeze(tf.argmax(predictions, axis=-1), axis=(2, 3)))
    nonzero_idx = tf.where(tf.not_equal(predictions, 0))
    sparse_outputs = tf.SparseTensor(nonzero_idx,
                                     tf.gather_nd(predictions, nonzero_idx),
                                     tf.shape(predictions, out_type=tf.int64))
    labels = tf.squeeze(labels, axis=(2, 3))
    nonzero_idx = tf.where(tf.not_equal(labels, 0))
    label_sparse_outputs = tf.SparseTensor(nonzero_idx,
                                           tf.gather_nd(labels, nonzero_idx),
                                           tf.shape(labels, out_type=tf.int64))
    distance = tf.reduce_sum(
        tf.edit_distance(sparse_outputs, label_sparse_outputs, normalize=False))
    reference_length = tf.to_float(common_layers.shape_list(nonzero_idx)[0])
    return distance / reference_length, reference_length
//  [python]
def sequence_edit_distance(predictions,
                           labels,
                           weights_fn=common_layers.weights_nonzero):
  """Average edit distance, ignoring padding 0s.

  The score returned is the edit distance divided by the total length of
  reference truth and the weight returned is the total length of the truth.

  Args:
    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
        type tf.float32 representing the logits, 0-padded.
    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
        representing the labels of same length as logits and 0-padded.
    weights_fn: ignored. The weights returned are the total length of the ground
        truth labels, excluding 0-paddings.

  Returns:
    (edit distance / reference length, reference length)

  Raises:
    ValueError: if weights_fn is not common_layers.weights_nonzero.
  """
  if weights_fn is not common_layers.weights_nonzero:
    raise ValueError("Only weights_nonzero can be used for this metric.")

  with tf.variable_scope("edit_distance", values=[predictions, labels]):
    # Transform logits into sequence classes by taking max at every step.
    predictions = tf.cast(
        tf.squeeze(tf.argmax(predictions, axis=-1), axis=(2, 3)), tf.int32)
    nonzero_idx = tf.where(tf.not_equal(predictions, 0))
    sparse_outputs = tf.SparseTensor(nonzero_idx,
                                     tf.gather_nd(predictions, nonzero_idx),
                                     tf.shape(predictions, out_type=tf.int64))
    labels = tf.squeeze(labels, axis=(2, 3))
    nonzero_idx = tf.where(tf.not_equal(labels, 0))
    label_sparse_outputs = tf.SparseTensor(nonzero_idx,
                                           tf.gather_nd(labels, nonzero_idx),
                                           tf.shape(labels, out_type=tf.int64))
    distance = tf.reduce_sum(
        tf.edit_distance(sparse_outputs, label_sparse_outputs, normalize=False))
    reference_length = tf.to_float(common_layers.shape_list(nonzero_idx)[0])
    return distance / reference_length, reference_length
//  [python]
def compute_damages(player, lib=np):
  percents = player['percent']
  damages = lib.maximum(percents[1:] - percents[:-1], 0)
  if lib == tf: damages = tf.to_float(damages)
  return damages
//  [python]
def compute_damages(player, lib=np):
  percents = player['percent']
  damages = lib.maximum(percents[1:] - percents[:-1], 0)
  if lib == tf: damages = tf.cast(damages, tf.float32)
  return damages
//  [python]
def compute_deaths(player, lib=np):
  dead = player['action_state'] <= 0xA
  deaths = lib.logical_and(lib.logical_not(dead[:-1]), dead[1:])
  if lib == tf: deaths = tf.to_float(deaths)
  return deaths
//  [python]
def compute_deaths(player, lib=np):
  dead = player['action_state'] <= 0xA
  deaths = lib.logical_and(lib.logical_not(dead[:-1]), dead[1:])
  if lib == tf: deaths = tf.cast(deaths, tf.float32)
  return deaths
//  [python]
def features_to_nonpadding(features, inputs_or_targets = 'inputs'):
    key = inputs_or_targets + '_segmentation'
    if features and key in features:
        return tf.minimum(tf.to_float(features[key]), 1.0)
    return None
//  [python]
def features_to_nonpadding(features, inputs_or_targets = 'inputs'):
    key = inputs_or_targets + '_segmentation'
    if features and key in features:
        return tf.minimum(tf.cast(features[key], tf.float32), 1.0)
    return None
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """
  image, masks = _split_image_and_masks(image)

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.to_float(tf.reduce_min(image))
    hi = tf.to_float(tf.reduce_max(image))

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.to_float(im) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)

  image = _concat_image_and_masks(image, masks)

  return image
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """
  image, masks = _split_image_and_masks(image)

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.cast(tf.reduce_min(image), tf.float32)
    hi = tf.cast(tf.reduce_max(image), tf.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.cast(im, tf.float32) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)

  image = _concat_image_and_masks(image, masks)

  return image
//  [python]
def set_scale_factors_to_output_size(self):
    """Set the parameters to resize input image to self._output_size."""
    # Compute the scale_factor using rounded scaled image size.
    height = tf.shape(self._image)[0]
    width = tf.shape(self._image)[1]
    max_image_size = tf.to_float(tf.maximum(height, width))
    image_scale = tf.to_float(self._output_size) / max_image_size
    scaled_height = tf.to_int32(tf.to_float(height) * image_scale)
    scaled_width = tf.to_int32(tf.to_float(width) * image_scale)
    self._image_scale = image_scale
    self._scaled_height = scaled_height
    self._scaled_width = scaled_width
//  [python]
def set_scale_factors_to_output_size(self):
    """Set the parameters to resize input image to self._output_size."""
    # Compute the scale_factor using rounded scaled image size.
    height = tf.shape(self._image)[0]
    width = tf.shape(self._image)[1]
    max_image_size = tf.cast(tf.maximum(height, width), tf.float32)
    image_scale = tf.cast(self._output_size, tf.float32) / max_image_size
    scaled_height = tf.cast(tf.cast(height, tf.float32) * image_scale, tf.int32)
    scaled_width = tf.cast(tf.cast(width, tf.float32) * image_scale, tf.int32)
    self._image_scale = image_scale
    self._scaled_height = scaled_height
    self._scaled_width = scaled_width
//  [python]
def test_classifier(self, ckpt_dir):
        print('Restoring from: {}'.format(ckpt_dir))

        g = tf.Graph()
        with g.as_default():
            # Get test batches
            batch_queue = self.get_data_queue()
            imgs_test, labels_test = batch_queue.get_next()
            imgs_test.set_shape([self.model.batch_size, ]+self.model.im_shape)

            # Get predictions
            predictions = self.model.linear_classifiers(imgs_test, self.data_generator.num_classes, training=False)

            num_corrects_list = []
            for preds, f_id in zip(predictions, self.model.feats_IDs):
                preds_test = tf.argmax(preds, 1)
                correct_preds = tf.equal(preds_test, labels_test)
                num_correct = tf.reduce_sum(tf.to_float(correct_preds))
                num_corrects_list.append(num_correct)

            # Start running operations on the Graph.
            init = tf.global_variables_initializer()
            sess = tf.Session()
            sess.run(init)

            prev_ckpt = get_checkpoint_path(ckpt_dir)
            print('Restoring from previous checkpoint: {}'.format(prev_ckpt))
            saver = tf.train.Saver(tf.global_variables())
            saver.restore(sess, prev_ckpt)

            n_cor_np = np.zeros([len(self.model.feats_IDs)])
            for i in range(self.num_eval_steps):
                n_correct = sess.run(num_corrects_list)
                n_cor_np += n_correct

            acc = n_cor_np/self.data_generator.num_samples
            print('Accuracy: {}'.format(acc))

            return acc
//  [python]
def test_classifier(self, ckpt_dir):
        print('Restoring from: {}'.format(ckpt_dir))

        g = tf.Graph()
        with g.as_default():
            # Get test batches
            batch_queue = self.get_data_queue()
            imgs_test, labels_test = batch_queue.get_next()
            imgs_test.set_shape([self.model.batch_size, ]+self.model.im_shape)

            # Get predictions
            predictions = self.model.linear_classifiers(imgs_test, self.data_generator.num_classes, training=False)

            num_corrects_list = []
            for preds, f_id in zip(predictions, self.model.feats_IDs):
                preds_test = tf.argmax(preds, 1)
                correct_preds = tf.equal(preds_test, labels_test)
                num_correct = tf.reduce_sum(tf.cast(correct_preds, tf.float32))
                num_corrects_list.append(num_correct)

            # Start running operations on the Graph.
            init = tf.global_variables_initializer()
            sess = tf.Session()
            sess.run(init)

            prev_ckpt = get_checkpoint_path(ckpt_dir)
            print('Restoring from previous checkpoint: {}'.format(prev_ckpt))
            saver = tf.train.Saver(tf.global_variables())
            saver.restore(sess, prev_ckpt)

            n_cor_np = np.zeros([len(self.model.feats_IDs)])
            for i in range(self.num_eval_steps):
                n_correct = sess.run(num_corrects_list)
                n_cor_np += n_correct

            acc = n_cor_np/self.data_generator.num_samples
            print('Accuracy: {}'.format(acc))

            return acc
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training, drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "label_ids": tf.FixedLenFeature([], tf.int64),
        "is_real_example": tf.FixedLenFeature([], tf.int64),
    }

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)

        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.to_int32(t)
            example[name] = t

        return example

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        # For training, we want a lot of parallel reading and shuffling.
        # For eval, we want no shuffling and parallel reading doesn't matter.
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)

        d = d.apply(
            tf.data.experimental.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=drop_remainder,
            )
        )

        return d

    return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training, drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "label_ids": tf.FixedLenFeature([], tf.int64),
        "is_real_example": tf.FixedLenFeature([], tf.int64),
    }

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)

        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.cast(t, tf.int32)
            example[name] = t

        return example

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        # For training, we want a lot of parallel reading and shuffling.
        # For eval, we want no shuffling and parallel reading doesn't matter.
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)

        d = d.apply(
            tf.data.experimental.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=drop_remainder,
            )
        )

        return d

    return input_fn
//  [python]
def resize_and_crop_boxes(self):
    """Resize boxes and crop it to the self._output dimension."""
    boxlist = preprocessor.box_list.BoxList(self._boxes)
    boxes = preprocessor.box_list_scale(
        boxlist, self._scaled_height, self._scaled_width).get()
    # Adjust box coordinates based on the offset.
    box_offset = tf.stack([self._crop_offset_y, self._crop_offset_x,
                           self._crop_offset_y, self._crop_offset_x,])
    boxes -= tf.to_float(tf.reshape(box_offset, [1, 4]))
    # Clip the boxes.
    boxes = self.clip_boxes(boxes)
    # Filter out ground truth boxes that are all zeros.
    indices = tf.where(tf.not_equal(tf.reduce_sum(boxes, axis=1), 0))
    boxes = tf.gather_nd(boxes, indices)
    classes = tf.gather_nd(self._classes, indices)
    return boxes, classes
//  [python]
def resize_and_crop_boxes(self):
    """Resize boxes and crop it to the self._output dimension."""
    boxlist = preprocessor.box_list.BoxList(self._boxes)
    boxes = preprocessor.box_list_scale(
        boxlist, self._scaled_height, self._scaled_width).get()
    # Adjust box coordinates based on the offset.
    box_offset = tf.stack([self._crop_offset_y, self._crop_offset_x,
                           self._crop_offset_y, self._crop_offset_x,])
    boxes -= tf.cast(tf.reshape(box_offset, [1, 4]), tf.float32)
    # Clip the boxes.
    boxes = self.clip_boxes(boxes)
    # Filter out ground truth boxes that are all zeros.
    indices = tf.where(tf.not_equal(tf.reduce_sum(boxes, axis=1), 0))
    boxes = tf.gather_nd(boxes, indices)
    classes = tf.gather_nd(self._classes, indices)
    return boxes, classes
//  [python]
def _apply_bbox_augmentation(image, bbox, bbox_idx, augmentation_func, *args):
  """Applies augmentation_func to the subsection of image indicated by bbox.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A modified version of image, where the bbox location in the image will
    have `ugmentation_func applied to it.
  """
  image, masks = _split_image_and_masks(image)
  image_with_mask = tf.concat([image, masks[:, :, bbox_idx:(bbox_idx + 1)]], axis=2)

  image_height = tf.to_float(tf.shape(image)[0])
  image_width = tf.to_float(tf.shape(image)[1])
  min_y = tf.to_int32(image_height * bbox[0])
  min_x = tf.to_int32(image_width * bbox[1])
  max_y = tf.to_int32(image_height * bbox[2])
  max_x = tf.to_int32(image_width * bbox[3])
  image_height = tf.to_int32(image_height)
  image_width = tf.to_int32(image_width)

  # Clip to be sure the max values do not fall out of range.
  max_y = tf.minimum(max_y, image_height - 1)
  max_x = tf.minimum(max_x, image_width - 1)

  # Get the sub-tensor that is the image within the bounding box region.
  bbox_content = image_with_mask[min_y:max_y + 1, min_x:max_x + 1, :]

  # Apply the augmentation function to the bbox portion of the image.
  augmented_bbox_content = augmentation_func(bbox_content, *args)

  # Pad the augmented_bbox_content and the mask to match the shape of original
  # image.
  augmented_bbox_content = tf.pad(augmented_bbox_content,
                                  [[min_y, (image_height - 1) - max_y],
                                   [min_x, (image_width - 1) - max_x],
                                   [0, 0]])

  # Create a mask that will be used to zero out a part of the original image.
  mask_tensor = tf.zeros_like(bbox_content)

  mask_tensor = tf.pad(mask_tensor,
                       [[min_y, (image_height - 1) - max_y],
                        [min_x, (image_width - 1) - max_x],
                        [0, 0]],
                       constant_values=1)
  # Replace the old bbox content with the new augmented content.
  image_with_mask = image_with_mask * mask_tensor + augmented_bbox_content

  image = tf.concat([
    image_with_mask[:, :, :3],
    masks[:, :, :bbox_idx],
    image_with_mask[:, :, 3:4],
    masks[:, :, (bbox_idx + 1):],
  ], axis=2)

  return image
//  [python]
def _apply_bbox_augmentation(image, bbox, bbox_idx, augmentation_func, *args):
  """Applies augmentation_func to the subsection of image indicated by bbox.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A modified version of image, where the bbox location in the image will
    have `ugmentation_func applied to it.
  """
  image, masks = _split_image_and_masks(image)
  image_with_mask = tf.concat([image, masks[:, :, bbox_idx:(bbox_idx + 1)]], axis=2)

  image_height = tf.cast(tf.shape(image)[0], tf.float32)
  image_width = tf.cast(tf.shape(image)[1], tf.float32)
  min_y = tf.to_int32(image_height * bbox[0])
  min_x = tf.to_int32(image_width * bbox[1])
  max_y = tf.to_int32(image_height * bbox[2])
  max_x = tf.to_int32(image_width * bbox[3])
  image_height = tf.to_int32(image_height)
  image_width = tf.to_int32(image_width)

  # Clip to be sure the max values do not fall out of range.
  max_y = tf.minimum(max_y, image_height - 1)
  max_x = tf.minimum(max_x, image_width - 1)

  # Get the sub-tensor that is the image within the bounding box region.
  bbox_content = image_with_mask[min_y:max_y + 1, min_x:max_x + 1, :]

  # Apply the augmentation function to the bbox portion of the image.
  augmented_bbox_content = augmentation_func(bbox_content, *args)

  # Pad the augmented_bbox_content and the mask to match the shape of original
  # image.
  augmented_bbox_content = tf.pad(augmented_bbox_content,
                                  [[min_y, (image_height - 1) - max_y],
                                   [min_x, (image_width - 1) - max_x],
                                   [0, 0]])

  # Create a mask that will be used to zero out a part of the original image.
  mask_tensor = tf.zeros_like(bbox_content)

  mask_tensor = tf.pad(mask_tensor,
                       [[min_y, (image_height - 1) - max_y],
                        [min_x, (image_width - 1) - max_x],
                        [0, 0]],
                       constant_values=1)
  # Replace the old bbox content with the new augmented content.
  image_with_mask = image_with_mask * mask_tensor + augmented_bbox_content

  image = tf.concat([
    image_with_mask[:, :, :3],
    masks[:, :, :bbox_idx],
    image_with_mask[:, :, 3:4],
    masks[:, :, (bbox_idx + 1):],
  ], axis=2)

  return image
//  [python]
def set_training_random_scale_factors(self, scale_min, scale_max):
    """Set the parameters for multiscale training."""
    # Select a random scale factor.
    random_scale_factor = tf.random_uniform([], scale_min, scale_max)
    scaled_size = tf.to_int32(random_scale_factor * self._output_size)

    # Recompute the accurate scale_factor using rounded scaled image size.
    height = tf.shape(self._image)[0]
    width = tf.shape(self._image)[1]
    max_image_size = tf.to_float(tf.maximum(height, width))
    image_scale = tf.to_float(scaled_size) / max_image_size

    # Select non-zero random offset (x, y) if scaled image is larger than
    # self._output_size.
    scaled_height = tf.to_int32(tf.to_float(height) * image_scale)
    scaled_width = tf.to_int32(tf.to_float(width) * image_scale)
    offset_y = tf.to_float(scaled_height - self._output_size)
    offset_x = tf.to_float(scaled_width - self._output_size)
    offset_y = tf.maximum(0.0, offset_y) * tf.random_uniform([], 0, 1)
    offset_x = tf.maximum(0.0, offset_x) * tf.random_uniform([], 0, 1)
    offset_y = tf.to_int32(offset_y)
    offset_x = tf.to_int32(offset_x)
    self._image_scale = image_scale
    self._scaled_height = scaled_height
    self._scaled_width = scaled_width
    self._crop_offset_x = offset_x
    self._crop_offset_y = offset_y
//  [python]
def set_training_random_scale_factors(self, scale_min, scale_max):
    """Set the parameters for multiscale training."""
    # Select a random scale factor.
    random_scale_factor = tf.random_uniform([], scale_min, scale_max)
    scaled_size = tf.cast(random_scale_factor * self._output_size, tf.int32)

    # Recompute the accurate scale_factor using rounded scaled image size.
    height = tf.shape(self._image)[0]
    width = tf.shape(self._image)[1]
    max_image_size = tf.cast(tf.maximum(height, width), tf.float32)
    image_scale = tf.cast(scaled_size, tf.float32) / max_image_size

    # Select non-zero random offset (x, y) if scaled image is larger than
    # self._output_size.
    scaled_height = tf.cast(tf.cast(height, tf.float32) * image_scale, tf.int32)
    scaled_width = tf.cast(tf.cast(width, tf.float32) * image_scale, tf.int32)
    offset_y = tf.cast(scaled_height - self._output_size, tf.float32)
    offset_x = tf.cast(scaled_width - self._output_size, tf.float32)
    offset_y = tf.maximum(0.0, offset_y) * tf.random_uniform([], 0, 1)
    offset_x = tf.maximum(0.0, offset_x) * tf.random_uniform([], 0, 1)
    offset_y = tf.cast(offset_y, tf.int32)
    offset_x = tf.cast(offset_x, tf.int32)
    self._image_scale = image_scale
    self._scaled_height = scaled_height
    self._scaled_width = scaled_width
    self._crop_offset_x = offset_x
    self._crop_offset_y = offset_y
//  [python]
def preprocess_for_eval(image, output_height, output_width, resize_side):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side: The smallest side of the image for aspect-preserving resizing.
  Returns:
    A preprocessed image.
  """
  image = _aspect_preserving_resize(image, resize_side)
  image = _central_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.to_float(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def preprocess_for_eval(image, output_height, output_width, resize_side):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side: The smallest side of the image for aspect-preserving resizing.
  Returns:
    A preprocessed image.
  """
  image = _aspect_preserving_resize(image, resize_side)
  image = _central_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.cast(image, tf.float32)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         resize_side_min=_RESIZE_SIDE_MIN,
                         resize_side_max=_RESIZE_SIDE_MAX):
  """Preprocesses the given image for training.

  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side_min: The lower bound for the smallest side of the image for
      aspect-preserving resizing.
    resize_side_max: The upper bound for the smallest side of the image for
      aspect-preserving resizing.
  Returns:
    A preprocessed image.
  """
  resize_side = tf.random_uniform(
      [], minval=resize_side_min, maxval=resize_side_max + 1, dtype=tf.int32)
  image = _aspect_preserving_resize(image, resize_side)
  image = _random_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.to_float(image)
  image = tf.image.random_flip_left_right(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         resize_side_min=_RESIZE_SIDE_MIN,
                         resize_side_max=_RESIZE_SIDE_MAX):
  """Preprocesses the given image for training.

  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side_min: The lower bound for the smallest side of the image for
      aspect-preserving resizing.
    resize_side_max: The upper bound for the smallest side of the image for
      aspect-preserving resizing.
  Returns:
    A preprocessed image.
  """
  resize_side = tf.random_uniform(
      [], minval=resize_side_min, maxval=resize_side_max + 1, dtype=tf.int32)
  image = _aspect_preserving_resize(image, resize_side)
  image = _random_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.cast(image, tf.float32)
  image = tf.image.random_flip_left_right(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def _cutout_inside_bbox(image, bbox, pad_fraction):
  """Generates cutout mask and the mean pixel value of the bbox.

  First a location is randomly chosen within the image as the center where the
  cutout mask will be applied. Note this can be towards the boundaries of the
  image, so the full cutout mask may not be applied.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    pad_fraction: Float that specifies how large the cutout mask should be in
      in reference to the size of the original bbox. If pad_fraction is 0.25,
      then the cutout mask will be of shape
      (0.25 * bbox height, 0.25 * bbox width).

  Returns:
    A tuple. Fist element is a tensor of the same shape as image where each
    element is either a 1 or 0 that is used to determine where the image
    will have cutout applied. The second element is the mean of the pixels
    in the image where the bbox is located.
  """
  num_channels = tf.shape(image)[2]
  image, masks = _split_image_and_masks(image)

  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  # Transform from shape [1, 4] to [4].
  bbox = tf.squeeze(bbox)

  min_y = tf.to_int32(tf.to_float(image_height) * bbox[0])
  min_x = tf.to_int32(tf.to_float(image_width) * bbox[1])
  max_y = tf.to_int32(tf.to_float(image_height) * bbox[2])
  max_x = tf.to_int32(tf.to_float(image_width) * bbox[3])

  # Calculate the mean pixel values in the bounding box, which will be used
  # to fill the cutout region.
  mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                        reduction_indices=[0, 1])

  # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
  # region lies entirely within the bbox.
  box_height = max_y - min_y + 1
  box_width = max_x - min_x + 1
  pad_size_height = tf.to_int32(pad_fraction * (box_height / 2))
  pad_size_width = tf.to_int32(pad_fraction * (box_width / 2))

  # Sample the center location in the image where the zero mask will be applied.
  cutout_center_height = tf.random_uniform(
      shape=[], minval=min_y, maxval=max_y+1,
      dtype=tf.int32)

  cutout_center_width = tf.random_uniform(
      shape=[], minval=min_x, maxval=max_x+1,
      dtype=tf.int32)

  lower_pad = tf.maximum(
      0, cutout_center_height - pad_size_height)
  upper_pad = tf.maximum(
      0, image_height - cutout_center_height - pad_size_height)
  left_pad = tf.maximum(
      0, cutout_center_width - pad_size_width)
  right_pad = tf.maximum(
      0, image_width - cutout_center_width - pad_size_width)

  cutout_shape = [image_height - (lower_pad + upper_pad),
                  image_width - (left_pad + right_pad)]
  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

  mask = tf.pad(
      tf.zeros(cutout_shape, dtype=image.dtype),
      padding_dims, constant_values=1)

  mask = tf.expand_dims(mask, 2)
  mask = tf.tile(mask, [1, 1, num_channels])

  return mask, mean
//  [python]
def _cutout_inside_bbox(image, bbox, pad_fraction):
  """Generates cutout mask and the mean pixel value of the bbox.

  First a location is randomly chosen within the image as the center where the
  cutout mask will be applied. Note this can be towards the boundaries of the
  image, so the full cutout mask may not be applied.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    pad_fraction: Float that specifies how large the cutout mask should be in
      in reference to the size of the original bbox. If pad_fraction is 0.25,
      then the cutout mask will be of shape
      (0.25 * bbox height, 0.25 * bbox width).

  Returns:
    A tuple. Fist element is a tensor of the same shape as image where each
    element is either a 1 or 0 that is used to determine where the image
    will have cutout applied. The second element is the mean of the pixels
    in the image where the bbox is located.
  """
  num_channels = tf.shape(image)[2]
  image, masks = _split_image_and_masks(image)

  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  # Transform from shape [1, 4] to [4].
  bbox = tf.squeeze(bbox)

  min_y = tf.cast(tf.to_float(image_height) * bbox[0], tf.int32)
  min_x = tf.cast(tf.to_float(image_width) * bbox[1], tf.int32)
  max_y = tf.cast(tf.to_float(image_height) * bbox[2], tf.int32)
  max_x = tf.cast(tf.to_float(image_width) * bbox[3], tf.int32)

  # Calculate the mean pixel values in the bounding box, which will be used
  # to fill the cutout region.
  mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                        reduction_indices=[0, 1])

  # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
  # region lies entirely within the bbox.
  box_height = max_y - min_y + 1
  box_width = max_x - min_x + 1
  pad_size_height = tf.cast(pad_fraction * (box_height / 2), tf.int32)
  pad_size_width = tf.cast(pad_fraction * (box_width / 2), tf.int32)

  # Sample the center location in the image where the zero mask will be applied.
  cutout_center_height = tf.random_uniform(
      shape=[], minval=min_y, maxval=max_y+1,
      dtype=tf.int32)

  cutout_center_width = tf.random_uniform(
      shape=[], minval=min_x, maxval=max_x+1,
      dtype=tf.int32)

  lower_pad = tf.maximum(
      0, cutout_center_height - pad_size_height)
  upper_pad = tf.maximum(
      0, image_height - cutout_center_height - pad_size_height)
  left_pad = tf.maximum(
      0, cutout_center_width - pad_size_width)
  right_pad = tf.maximum(
      0, image_width - cutout_center_width - pad_size_width)

  cutout_shape = [image_height - (lower_pad + upper_pad),
                  image_width - (left_pad + right_pad)]
  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

  mask = tf.pad(
      tf.zeros(cutout_shape, dtype=image.dtype),
      padding_dims, constant_values=1)

  mask = tf.expand_dims(mask, 2)
  mask = tf.tile(mask, [1, 1, num_channels])

  return mask, mean
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn"t assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    the cropped (and resized) image.

  Raises:
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3), ["Rank of image must be equal to 3."])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ["Crop size greater than the image size."])

  offsets = tf.to_int32(tf.stack([offset_height, offset_width, 0]))

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  return tf.reshape(image, cropped_shape)
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn"t assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    the cropped (and resized) image.

  Raises:
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3), ["Rank of image must be equal to 3."])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ["Crop size greater than the image size."])

  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  return tf.reshape(image, cropped_shape)
//  [python]
def transform_keypoints(keypoints, rotation, translation):
    """
    Arguments:
        keypoints: an int tensor with shape [num_persons, 17, 3].
        rotation: a float tensor with shape [2, 2].
        translation: a float tensor with shape [1, 2].
    Returns:
        an int tensor with shape [num_persons, 17, 3].
    """

    points, v = tf.split(keypoints, [2, 1], axis=2)
    # they have shapes [num_persons, 17, 2] and [num_persons, 17, 1]

    points = tf.to_float(points)
    points = tf.reshape(points, [-1, 2])
    points = tf.matmul(points, rotation) + translation
    points = tf.to_int32(tf.round(points))
    points = tf.reshape(points, [-1, 17, 2])
    keypoints = tf.concat([points, v], axis=2)

    return keypoints
//  [python]
def transform_keypoints(keypoints, rotation, translation):
    """
    Arguments:
        keypoints: an int tensor with shape [num_persons, 17, 3].
        rotation: a float tensor with shape [2, 2].
        translation: a float tensor with shape [1, 2].
    Returns:
        an int tensor with shape [num_persons, 17, 3].
    """

    points, v = tf.split(keypoints, [2, 1], axis=2)
    # they have shapes [num_persons, 17, 2] and [num_persons, 17, 1]

    points = tf.cast(points, tf.float32)
    points = tf.reshape(points, [-1, 2])
    points = tf.matmul(points, rotation) + translation
    points = tf.cast(tf.round(points), tf.int32)
    points = tf.reshape(points, [-1, 17, 2])
    keypoints = tf.concat([points, v], axis=2)

    return keypoints
//  [python]
def __call__(self, params):
    example_decoder = tf_example_decoder.TfExampleSegmentationDecoder()
    def _dataset_parser(value):
      """Parse data to a fixed dimension input image and learning targets.

      Args:
        value: A dictionary contains an image and groundtruth annotations.

      Returns:
        A list of the following elements in order:
        image: Image tensor that is preproessed to have normalized value and
          fixed dimension [image_size, image_size, 3]
        label: label tensor of the same spatial dimension as the image.
      """
      with tf.name_scope('parser'):
        data = example_decoder.decode(value)
        image = data['image']
        label = data['labels_class']
        label = tf.to_int32(label)
        input_processor = SegmentationInputProcessor(image,
                                                     params['image_size'],
                                                     label)
        # The image normalization is identical to Cloud TPU ResNet.
        input_processor.normalize_image()
        if self._is_training and params['input_rand_hflip']:
          input_processor.random_horizontal_flip()
        if self._is_training:
          input_processor.set_training_random_scale_factors(
              params['train_scale_min'], params['train_scale_max'])
        image = input_processor.resize_and_crop_image()

        # Set padding to background (class=0) during training.
        if self._is_training:
          label = input_processor.resize_and_crop_label(0)
        else:
          label = input_processor.resize_and_crop_label(params['ignore_label'])
        if params['use_bfloat16']:
          image = tf.cast(image, dtype=tf.bfloat16)
        return image, label

    batch_size = params['batch_size']

    dataset = tf.data.Dataset.list_files(
        self._file_pattern, shuffle=self._is_training)
    if self._is_training:
      dataset = dataset.repeat()

    def _prefetch_dataset(filename):
      dataset = tf.data.TFRecordDataset(filename).prefetch(1)
      return dataset

    dataset = dataset.apply(
        contrib_data.parallel_interleave(
            _prefetch_dataset, cycle_length=32, sloppy=self._is_training))
    if self._is_training:
      dataset = dataset.shuffle(64)

    dataset = dataset.map(_dataset_parser, num_parallel_calls=64)
    dataset = dataset.prefetch(batch_size)
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset