jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _dropout_lstm_cell(hparams, train):
  return tf.nn.rnn_cell.DropoutWrapper(
      tf.nn.rnn_cell.LSTMCell(hparams.hidden_size),
      input_keep_prob=1.0 - hparams.dropout * tf.to_float(train))
//  [python]
def _dropout_lstm_cell(hparams, train):
  return tf.nn.rnn_cell.DropoutWrapper(
      tf.nn.rnn_cell.LSTMCell(hparams.hidden_size),
      input_keep_prob=1.0 - hparams.dropout * tf.cast(train, tf.float32))
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
    #with tf.device("/device:GPU:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.to_int32(valid_shuffle_preds)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.to_int32(self.valid_shuffle_acc)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
    #with tf.device("/device:GPU:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.cast(valid_shuffle_preds, tf.int32)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.cast(self.valid_shuffle_acc, tf.int32)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def _extract_prediction_tensors(model,
                                create_input_dict_fn,
                                data_preprocessing_steps,
                                ignore_groundtruth=False,
                                evaluate_with_lexicon=False):
  # input queue
  input_dict = create_input_dict_fn()
  prefetch_queue = prefetcher.prefetch(input_dict, capacity=500)
  input_dict = prefetch_queue.dequeue()
  original_image = tf.to_float(input_dict[fields.InputDataFields.image])
  original_image_shape = tf.shape(original_image)
  input_dict[fields.InputDataFields.image] = original_image

  # data preprocessing
  preprocessed_input_dict = preprocessor.preprocess(input_dict, data_preprocessing_steps)

  # model inference
  preprocessed_image = preprocessed_input_dict[fields.InputDataFields.image]
  preprocessed_image_shape = tf.shape(preprocessed_image)
  predictions_dict = model.predict(tf.expand_dims(preprocessed_image, 0))
  recognitions = model.postprocess(predictions_dict)

  def _lexicon_search(lexicon, word):
    edit_distances = []
    for lex_word in lexicon:
      edit_distances.append(editdistance.eval(lex_word.lower(), word.lower()))
    edit_distances = np.asarray(edit_distances, dtype=np.int)
    argmin = np.argmin(edit_distances)
    return lexicon[argmin]

  if evaluate_with_lexicon:
    lexicon = input_dict[fields.InputDataFields.lexicon]
    recognition_text = tf.py_func(
      _lexicon_search,
      [lexicon, recognitions['text'][0]],
      tf.string,
      stateful=False,
    )
  else:
    recognition_text = recognitions['text'][0]

  tensor_dict = {
    'original_image': original_image,
    'original_image_shape': original_image_shape,
    'preprocessed_image_shape': preprocessed_image_shape,
    'filename': preprocessed_input_dict[fields.InputDataFields.filename],
    'groundtruth_text': input_dict[fields.InputDataFields.groundtruth_text],
    'recognition_text': recognition_text,
  }
  if 'control_points' in predictions_dict:
    tensor_dict.update({
      'control_points': predictions_dict['control_points'],
      'rectified_images': predictions_dict['rectified_images']
    })

  return tensor_dict
//  [python]
def _extract_prediction_tensors(model,
                                create_input_dict_fn,
                                data_preprocessing_steps,
                                ignore_groundtruth=False,
                                evaluate_with_lexicon=False):
  # input queue
  input_dict = create_input_dict_fn()
  prefetch_queue = prefetcher.prefetch(input_dict, capacity=500)
  input_dict = prefetch_queue.dequeue()
  original_image = tf.cast(input_dict[fields.InputDataFields.image], tf.float32)
  original_image_shape = tf.shape(original_image)
  input_dict[fields.InputDataFields.image] = original_image

  # data preprocessing
  preprocessed_input_dict = preprocessor.preprocess(input_dict, data_preprocessing_steps)

  # model inference
  preprocessed_image = preprocessed_input_dict[fields.InputDataFields.image]
  preprocessed_image_shape = tf.shape(preprocessed_image)
  predictions_dict = model.predict(tf.expand_dims(preprocessed_image, 0))
  recognitions = model.postprocess(predictions_dict)

  def _lexicon_search(lexicon, word):
    edit_distances = []
    for lex_word in lexicon:
      edit_distances.append(editdistance.eval(lex_word.lower(), word.lower()))
    edit_distances = np.asarray(edit_distances, dtype=np.int)
    argmin = np.argmin(edit_distances)
    return lexicon[argmin]

  if evaluate_with_lexicon:
    lexicon = input_dict[fields.InputDataFields.lexicon]
    recognition_text = tf.py_func(
      _lexicon_search,
      [lexicon, recognitions['text'][0]],
      tf.string,
      stateful=False,
    )
  else:
    recognition_text = recognitions['text'][0]

  tensor_dict = {
    'original_image': original_image,
    'original_image_shape': original_image_shape,
    'preprocessed_image_shape': preprocessed_image_shape,
    'filename': preprocessed_input_dict[fields.InputDataFields.filename],
    'groundtruth_text': input_dict[fields.InputDataFields.groundtruth_text],
    'recognition_text': recognition_text,
  }
  if 'control_points' in predictions_dict:
    tensor_dict.update({
      'control_points': predictions_dict['control_points'],
      'rectified_images': predictions_dict['rectified_images']
    })

  return tensor_dict
//  [python]
def testSubsampleThreeByThree(self):
    x = tf.reshape(tf.to_float(tf.range(9)), [1, 3, 3, 1])
    x = resnet_utils.subsample(x, 2)
    expected = tf.reshape(tf.constant([0, 2, 6, 8]), [1, 2, 2, 1])
    with self.test_session():
      self.assertAllClose(x.eval(), expected.eval())
//  [python]
def testSubsampleThreeByThree(self):
    x = tf.reshape(tf.cast(tf.range(9), tf.float32), [1, 3, 3, 1])
    x = resnet_utils.subsample(x, 2)
    expected = tf.reshape(tf.constant([0, 2, 6, 8]), [1, 2, 2, 1])
    with self.test_session():
      self.assertAllClose(x.eval(), expected.eval())
//  [python]
def create_test_input(batch_size, height, width, channels):
  """Create test input tensor.

  Args:
    batch_size: The number of images per batch or `None` if unknown.
    height: The height of each image or `None` if unknown.
    width: The width of each image or `None` if unknown.
    channels: The number of channels per image or `None` if unknown.

  Returns:
    Either a placeholder `Tensor` of dimension
      [batch_size, height, width, channels] if any of the inputs are `None` or a
    constant `Tensor` with the mesh grid values along the spatial dimensions.
  """
  if None in [batch_size, height, width, channels]:
    return tf.placeholder(tf.float32, (batch_size, height, width, channels))
  else:
    return tf.to_float(
        np.tile(
            np.reshape(
                np.reshape(np.arange(height), [height, 1]) +
                np.reshape(np.arange(width), [1, width]),
                [1, height, width, 1]),
            [batch_size, 1, 1, channels]))
//  [python]
def create_test_input(batch_size, height, width, channels):
  """Create test input tensor.

  Args:
    batch_size: The number of images per batch or `None` if unknown.
    height: The height of each image or `None` if unknown.
    width: The width of each image or `None` if unknown.
    channels: The number of channels per image or `None` if unknown.

  Returns:
    Either a placeholder `Tensor` of dimension
      [batch_size, height, width, channels] if any of the inputs are `None` or a
    constant `Tensor` with the mesh grid values along the spatial dimensions.
  """
  if None in [batch_size, height, width, channels]:
    return tf.placeholder(tf.float32, (batch_size, height, width, channels))
  else:
    return tf.cast(
        np.tile(
            np.reshape(
                np.reshape(np.arange(height), [height, 1]) +
                np.reshape(np.arange(width), [1, width]),
                [1, height, width, 1]),
            [batch_size, 1, 1, channels]), tf.float32)
//  [python]
def testSubsampleFourByFour(self):
    x = tf.reshape(tf.to_float(tf.range(16)), [1, 4, 4, 1])
    x = resnet_utils.subsample(x, 2)
    expected = tf.reshape(tf.constant([0, 2, 8, 10]), [1, 2, 2, 1])
    with self.test_session():
      self.assertAllClose(x.eval(), expected.eval())
//  [python]
def testSubsampleFourByFour(self):
    x = tf.reshape(tf.cast(tf.range(16), tf.float32), [1, 4, 4, 1])
    x = resnet_utils.subsample(x, 2)
    expected = tf.reshape(tf.constant([0, 2, 8, 10]), [1, 2, 2, 1])
    with self.test_session():
      self.assertAllClose(x.eval(), expected.eval())
//  [python]
def body(self, features):
    hparams = copy.copy(self._hparams)
    targets = features["targets"]
    if (hparams.likelihood == cia.DistributionType.DMOL and
        hparams.num_channels != 1):
      raise ValueError("When using DMOL for the likelihood, bottom function "
                       " must be identity and num_channels must be 1.")
    if (not tf.get_variable_scope().reuse and
        hparams.mode != tf.estimator.ModeKeys.PREDICT):
      tf.summary.image("targets", tf.to_float(targets), max_outputs=1)

    # Extra losses list if we want to use moe.
    losses = []
    # Prepare decoder inputs and bias.
    decoder_input, rows, cols = cia.prepare_decoder(targets, hparams)
    # Add class label to decoder input.
    if not hparams.unconditional:
      inputs = features["inputs"]
      decoder_input += tf.reshape(
          inputs,
          [common_layers.shape_list(targets)[0], 1, 1, hparams.hidden_size])
    decoder_output = cia.transformer_decoder_layers(
        decoder_input,
        None,
        hparams.num_decoder_layers or hparams.num_hidden_layers,
        hparams,
        attention_type=hparams.dec_attention_type,
        losses=losses,
        name="decoder")
    output = cia.create_output(decoder_output, rows, cols, targets, hparams)

    if losses:
      return output, {"extra_loss": tf.add_n(losses)}
    else:
      return output
//  [python]
def body(self, features):
    hparams = copy.copy(self._hparams)
    targets = features["targets"]
    if (hparams.likelihood == cia.DistributionType.DMOL and
        hparams.num_channels != 1):
      raise ValueError("When using DMOL for the likelihood, bottom function "
                       " must be identity and num_channels must be 1.")
    if (not tf.get_variable_scope().reuse and
        hparams.mode != tf.estimator.ModeKeys.PREDICT):
      tf.summary.image("targets", tf.cast(targets, tf.float32), max_outputs=1)

    # Extra losses list if we want to use moe.
    losses = []
    # Prepare decoder inputs and bias.
    decoder_input, rows, cols = cia.prepare_decoder(targets, hparams)
    # Add class label to decoder input.
    if not hparams.unconditional:
      inputs = features["inputs"]
      decoder_input += tf.reshape(
          inputs,
          [common_layers.shape_list(targets)[0], 1, 1, hparams.hidden_size])
    decoder_output = cia.transformer_decoder_layers(
        decoder_input,
        None,
        hparams.num_decoder_layers or hparams.num_hidden_layers,
        hparams,
        attention_type=hparams.dec_attention_type,
        losses=losses,
        name="decoder")
    output = cia.create_output(decoder_output, rows, cols, targets, hparams)

    if losses:
      return output, {"extra_loss": tf.add_n(losses)}
    else:
      return output
//  [python]
def extract_resnet_features(max_step_run=39000):
  """Not checked provide_resnet_noisy_data dataset might change."""
  g = tf.Graph()

  with g.as_default():
    tf_global_step = tf.train.get_or_create_global_step()

    # pylint: disable=line-too-long
    images, one_hot_labels, num_examples, num_of_classes, clean_labels, image_ids = cifar_data_provider.provide_resnet_noisy_data(
        FLAGS.dataset_name,
        'train',
        FLAGS.batch_size,
        dataset_dir=FLAGS.data_dir)

    hps = resnet_model.HParams(
        batch_size=FLAGS.batch_size,
        num_classes=num_of_classes,
        min_lrn_rate=0.0001,
        lrn_rate=0,
        num_residual_units=9,
        use_bottleneck=False,
        weight_decay_rate=0.0002,
        relu_leakiness=0.1,
        optimizer='mom')

    images.set_shape([FLAGS.batch_size, 32, 32, 3])

    # Define the model:
    resnet = resnet_model.ResNet(hps, images, one_hot_labels, mode='train')
    logits = resnet.build_model()

    # Specify the loss function:
    loss = tf.nn.softmax_cross_entropy_with_logits(
        labels=one_hot_labels, logits=logits)
    labels = tf.argmax(one_hot_labels, 1)

    loss = tf.reshape(loss, [-1, 1])

    epoch_step = tf.to_int32(
        tf.floor(tf.divide(tf_global_step, max_step_run) * 100))

    ckpt_model = FLAGS.checkpoint_dir

    num_batches = int(math.ceil(num_examples / float(FLAGS.batch_size)))
    init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
        ckpt_model,
        tf.contrib.framework.get_variables_to_restore(include=['.*']))
    outfile = open(FLAGS.output_csv_file, 'w')
    with tf.Session() as sess:
      with slim.queues.QueueRunners(sess):
        init_fn(sess)
        for _ in xrange(num_batches):
          image_ids_val, epoch_step_val, labels_val, loss_val, clean_labels_val = sess.run(
              [image_ids, epoch_step, labels, loss, clean_labels])
          clean_labels_val = np.squeeze(clean_labels_val)
          loss_val = np.squeeze(loss_val)
          image_ids_val = np.squeeze(image_ids_val)
          labels_val = np.squeeze(labels_val)
          for k in range(FLAGS.batch_size):
            outfile.write('{} {} {} {} {:5f}\n'.format(
                int(image_ids_val[k]), epoch_step_val, labels_val[k],
                clean_labels_val[k], loss_val[k]))

    outfile.flush()
    outfile.close()
//  [python]
def extract_resnet_features(max_step_run=39000):
  """Not checked provide_resnet_noisy_data dataset might change."""
  g = tf.Graph()

  with g.as_default():
    tf_global_step = tf.train.get_or_create_global_step()

    # pylint: disable=line-too-long
    images, one_hot_labels, num_examples, num_of_classes, clean_labels, image_ids = cifar_data_provider.provide_resnet_noisy_data(
        FLAGS.dataset_name,
        'train',
        FLAGS.batch_size,
        dataset_dir=FLAGS.data_dir)

    hps = resnet_model.HParams(
        batch_size=FLAGS.batch_size,
        num_classes=num_of_classes,
        min_lrn_rate=0.0001,
        lrn_rate=0,
        num_residual_units=9,
        use_bottleneck=False,
        weight_decay_rate=0.0002,
        relu_leakiness=0.1,
        optimizer='mom')

    images.set_shape([FLAGS.batch_size, 32, 32, 3])

    # Define the model:
    resnet = resnet_model.ResNet(hps, images, one_hot_labels, mode='train')
    logits = resnet.build_model()

    # Specify the loss function:
    loss = tf.nn.softmax_cross_entropy_with_logits(
        labels=one_hot_labels, logits=logits)
    labels = tf.argmax(one_hot_labels, 1)

    loss = tf.reshape(loss, [-1, 1])

    epoch_step = tf.cast(
        tf.floor(tf.divide(tf_global_step, max_step_run) * 100), tf.int32)

    ckpt_model = FLAGS.checkpoint_dir

    num_batches = int(math.ceil(num_examples / float(FLAGS.batch_size)))
    init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
        ckpt_model,
        tf.contrib.framework.get_variables_to_restore(include=['.*']))
    outfile = open(FLAGS.output_csv_file, 'w')
    with tf.Session() as sess:
      with slim.queues.QueueRunners(sess):
        init_fn(sess)
        for _ in xrange(num_batches):
          image_ids_val, epoch_step_val, labels_val, loss_val, clean_labels_val = sess.run(
              [image_ids, epoch_step, labels, loss, clean_labels])
          clean_labels_val = np.squeeze(clean_labels_val)
          loss_val = np.squeeze(loss_val)
          image_ids_val = np.squeeze(image_ids_val)
          labels_val = np.squeeze(labels_val)
          for k in range(FLAGS.batch_size):
            outfile.write('{} {} {} {} {:5f}\n'.format(
                int(image_ids_val[k]), epoch_step_val, labels_val[k],
                clean_labels_val[k], loss_val[k]))

    outfile.flush()
    outfile.close()
//  [python]
def compute_loss(self, unreduced_loss):
    """Computes scaled loss based on mask out size."""
    # construct mask to identify zero padding that was introduced to
    # make the batch rectangular
    batch_duration = tf.shape(self.pianorolls)[1]
    indices = tf.to_float(tf.range(batch_duration))
    pad_mask = tf.to_float(
        indices[None, :, None, None] < self.lengths[:, None, None, None])

    # construct mask and its complement, respecting pad mask
    mask = pad_mask * self.masks
    unmask = pad_mask * (1. - self.masks)

    # Compute numbers of variables
    # #timesteps * #variables per timestep
    variable_axis = 3 if self.hparams.use_softmax_loss else 2
    dd = (
        self.lengths[:, None, None, None] * tf.to_float(
            tf.shape(self.pianorolls)[variable_axis]))
    reduced_dd = tf.reduce_sum(dd)

    # Compute numbers of variables to be predicted/conditioned on
    mask_size = tf.reduce_sum(mask, axis=[1, variable_axis], keep_dims=True)
    unmask_size = tf.reduce_sum(unmask, axis=[1, variable_axis], keep_dims=True)

    unreduced_loss *= pad_mask
    if self.hparams.rescale_loss:
      unreduced_loss *= dd / mask_size

    # Compute average loss over entire set of variables
    self.loss_total = tf.reduce_sum(unreduced_loss) / reduced_dd

    # Compute separate losses for masked/unmasked variables
    # NOTE: indexing the pitch dimension with 0 because the mask is constant
    # across pitch. Except in the sigmoid case, but then the pitch dimension
    # will have been reduced over.
    self.reduced_mask_size = tf.reduce_sum(mask_size[:, :, 0, :])
    self.reduced_unmask_size = tf.reduce_sum(unmask_size[:, :, 0, :])

    assert_partition_op = tf.group(
        tf.assert_equal(tf.reduce_sum(mask * unmask), 0.),
        tf.assert_equal(self.reduced_mask_size + self.reduced_unmask_size,
                        reduced_dd))
    with tf.control_dependencies([assert_partition_op]):
      self.loss_mask = (
          tf.reduce_sum(mask * unreduced_loss) / self.reduced_mask_size)
      self.loss_unmask = (
          tf.reduce_sum(unmask * unreduced_loss) / self.reduced_unmask_size)

    # Check which loss to use as objective function.
    self.loss = (
        self.loss_mask if self.hparams.optimize_mask_only else self.loss_total)
//  [python]
def compute_loss(self, unreduced_loss):
    """Computes scaled loss based on mask out size."""
    # construct mask to identify zero padding that was introduced to
    # make the batch rectangular
    batch_duration = tf.shape(self.pianorolls)[1]
    indices = tf.cast(tf.range(batch_duration), tf.float32)
    pad_mask = tf.cast(
        indices[None, :, None, None] < self.lengths[:, None, None, None], tf.float32)

    # construct mask and its complement, respecting pad mask
    mask = pad_mask * self.masks
    unmask = pad_mask * (1. - self.masks)

    # Compute numbers of variables
    # #timesteps * #variables per timestep
    variable_axis = 3 if self.hparams.use_softmax_loss else 2
    dd = (
        self.lengths[:, None, None, None] * tf.cast(
            tf.shape(self.pianorolls)[variable_axis], tf.float32))
    reduced_dd = tf.reduce_sum(dd)

    # Compute numbers of variables to be predicted/conditioned on
    mask_size = tf.reduce_sum(mask, axis=[1, variable_axis], keep_dims=True)
    unmask_size = tf.reduce_sum(unmask, axis=[1, variable_axis], keep_dims=True)

    unreduced_loss *= pad_mask
    if self.hparams.rescale_loss:
      unreduced_loss *= dd / mask_size

    # Compute average loss over entire set of variables
    self.loss_total = tf.reduce_sum(unreduced_loss) / reduced_dd

    # Compute separate losses for masked/unmasked variables
    # NOTE: indexing the pitch dimension with 0 because the mask is constant
    # across pitch. Except in the sigmoid case, but then the pitch dimension
    # will have been reduced over.
    self.reduced_mask_size = tf.reduce_sum(mask_size[:, :, 0, :])
    self.reduced_unmask_size = tf.reduce_sum(unmask_size[:, :, 0, :])

    assert_partition_op = tf.group(
        tf.assert_equal(tf.reduce_sum(mask * unmask), 0.),
        tf.assert_equal(self.reduced_mask_size + self.reduced_unmask_size,
                        reduced_dd))
    with tf.control_dependencies([assert_partition_op]):
      self.loss_mask = (
          tf.reduce_sum(mask * unreduced_loss) / self.reduced_mask_size)
      self.loss_unmask = (
          tf.reduce_sum(unmask * unreduced_loss) / self.reduced_unmask_size)

    # Check which loss to use as objective function.
    self.loss = (
        self.loss_mask if self.hparams.optimize_mask_only else self.loss_total)
//  [python]
def logits_to_predict(attr_logits, name=None):
    """
    Args:
        :param name: add name for tensor if name is not None
        :param attr_logits:
    Returns:
        predict_label nx1 [-1,1,0,-1,-1] int64

    """
    specific_logits = attr_logits[:, 0]    # 1 means sure    0 means not sure
    attribute_logits = attr_logits[:, 1]   # 1 means Yes     0 means No

    prediction = tf.where(attribute_logits > 0.5, tf.ones_like(attribute_logits), tf.zeros_like(attribute_logits))
    prediction = tf.where(specific_logits < 0.5, -tf.ones_like(prediction), prediction)
    predict_label = tf.to_int32(prediction)
    if name:
        return tf.identity(predict_label, name='{}_predict'.format(name))
    else:
        return predict_label
//  [python]
def logits_to_predict(attr_logits, name=None):
    """
    Args:
        :param name: add name for tensor if name is not None
        :param attr_logits:
    Returns:
        predict_label nx1 [-1,1,0,-1,-1] int64

    """
    specific_logits = attr_logits[:, 0]    # 1 means sure    0 means not sure
    attribute_logits = attr_logits[:, 1]   # 1 means Yes     0 means No

    prediction = tf.where(attribute_logits > 0.5, tf.ones_like(attribute_logits), tf.zeros_like(attribute_logits))
    prediction = tf.where(specific_logits < 0.5, -tf.ones_like(prediction), prediction)
    predict_label = tf.cast(prediction, tf.int32)
    if name:
        return tf.identity(predict_label, name='{}_predict'.format(name))
    else:
        return predict_label
//  [python]
def logits_to_predict_v2(attr_logits, name=None):
    """
    this function is only contains two type
    Args:
        :param name: add name for tensor if name is not None
        :param attr_logits:
    Returns:
        predict_label nx1 [-1,1,0,-1,-1] int64

    """
    prediction = tf.argmax(attr_logits, axis=-1)
    predict_label = tf.to_int32(prediction)
    if name:
        return tf.identity(predict_label, name='{}_predict'.format(name))
    else:
        return predict_label
//  [python]
def logits_to_predict_v2(attr_logits, name=None):
    """
    this function is only contains two type
    Args:
        :param name: add name for tensor if name is not None
        :param attr_logits:
    Returns:
        predict_label nx1 [-1,1,0,-1,-1] int64

    """
    prediction = tf.argmax(attr_logits, axis=-1)
    predict_label = tf.cast(prediction, tf.int32)
    if name:
        return tf.identity(predict_label, name='{}_predict'.format(name))
    else:
        return predict_label
//  [python]
def testConv2DSameEven(self):
    n, n2 = 4, 2

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    w = create_test_input(1, 3, 3, 1)
    w = tf.reshape(w, [3, 3, 1, 1])

    tf.get_variable('Conv/weights', initializer=w)
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.conv2d(x, 1, [3, 3], stride=1, scope='Conv')
    y1_expected = tf.to_float([[14, 28, 43, 26],
                               [28, 48, 66, 37],
                               [43, 66, 84, 46],
                               [26, 37, 46, 22]])
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.to_float([[14, 43],
                               [43, 84]])
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = resnet_utils.conv2d_same(x, 1, 3, stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.conv2d(x, 1, [3, 3], stride=2, scope='Conv')
    y4_expected = tf.to_float([[48, 37],
                               [37, 22]])
    y4_expected = tf.reshape(y4_expected, [1, n2, n2, 1])

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def testConv2DSameEven(self):
    n, n2 = 4, 2

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    w = create_test_input(1, 3, 3, 1)
    w = tf.reshape(w, [3, 3, 1, 1])

    tf.get_variable('Conv/weights', initializer=w)
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.conv2d(x, 1, [3, 3], stride=1, scope='Conv')
    y1_expected = tf.cast([[14, 28, 43, 26],
                               [28, 48, 66, 37],
                               [43, 66, 84, 46],
                               [26, 37, 46, 22]], tf.float32)
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.cast([[14, 43],
                               [43, 84]], tf.float32)
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = resnet_utils.conv2d_same(x, 1, 3, stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.conv2d(x, 1, [3, 3], stride=2, scope='Conv')
    y4_expected = tf.cast([[48, 37],
                               [37, 22]], tf.float32)
    y4_expected = tf.reshape(y4_expected, [1, n2, n2, 1])

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def build(self, images):
    """Builds the embedder with the given speicifcation.

    Args:
      images: a tensor that contains the input images which has the shape of
        NxTxHxWxC where N is the batch size, T is the maximum length of the
        sequence, H and W are the height and width of the images and C is the
        number of channels.

    Returns:
      A tensor that is the embedding of the images.
    """

    shape = images.get_shape().as_list()
    images = tf.reshape(images,
                        [shape[0] * shape[1], shape[2], shape[3], shape[4]])

    with slim.arg_scope(
        [slim.conv2d, slim.fully_connected],
        activation_fn=tf.nn.relu,
        weights_regularizer=slim.l2_regularizer(self._params.weight_decay_rate),
        biases_initializer=tf.zeros_initializer()):
      with slim.arg_scope([slim.conv2d], padding='SAME'):
        # convert the image to one hot if needed.
        if self._params.to_one_hot:
          net = tf.one_hot(
              tf.squeeze(tf.to_int32(images), axis=[-1]),
              self._params.one_hot_length)
        else:
          net = images

        p = self._params
        # Adding conv layers with the specified configurations.
        for conv_id, kernel_stride_channel in enumerate(
            zip(p.conv_sizes, p.conv_strides, p.conv_channels)):
          kernel_size, stride, channels = kernel_stride_channel
          net = slim.conv2d(
              net,
              channels, [kernel_size, kernel_size],
              stride,
              scope='conv_{}'.format(conv_id + 1))

        net = slim.flatten(net)
        net = slim.fully_connected(net, self._params.embedding_size, scope='fc')

        output = tf.reshape(net, [shape[0], shape[1], -1])
        return output
//  [python]
def build(self, images):
    """Builds the embedder with the given speicifcation.

    Args:
      images: a tensor that contains the input images which has the shape of
        NxTxHxWxC where N is the batch size, T is the maximum length of the
        sequence, H and W are the height and width of the images and C is the
        number of channels.

    Returns:
      A tensor that is the embedding of the images.
    """

    shape = images.get_shape().as_list()
    images = tf.reshape(images,
                        [shape[0] * shape[1], shape[2], shape[3], shape[4]])

    with slim.arg_scope(
        [slim.conv2d, slim.fully_connected],
        activation_fn=tf.nn.relu,
        weights_regularizer=slim.l2_regularizer(self._params.weight_decay_rate),
        biases_initializer=tf.zeros_initializer()):
      with slim.arg_scope([slim.conv2d], padding='SAME'):
        # convert the image to one hot if needed.
        if self._params.to_one_hot:
          net = tf.one_hot(
              tf.squeeze(tf.cast(images, tf.int32), axis=[-1]),
              self._params.one_hot_length)
        else:
          net = images

        p = self._params
        # Adding conv layers with the specified configurations.
        for conv_id, kernel_stride_channel in enumerate(
            zip(p.conv_sizes, p.conv_strides, p.conv_channels)):
          kernel_size, stride, channels = kernel_stride_channel
          net = slim.conv2d(
              net,
              channels, [kernel_size, kernel_size],
              stride,
              scope='conv_{}'.format(conv_id + 1))

        net = slim.flatten(net)
        net = slim.fully_connected(net, self._params.embedding_size, scope='fc')

        output = tf.reshape(net, [shape[0], shape[1], -1])
        return output
//  [python]
def testConv2DSameOdd(self):
    n, n2 = 5, 3

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    w = create_test_input(1, 3, 3, 1)
    w = tf.reshape(w, [3, 3, 1, 1])

    tf.get_variable('Conv/weights', initializer=w)
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.conv2d(x, 1, [3, 3], stride=1, scope='Conv')
    y1_expected = tf.to_float([[14, 28, 43, 58, 34],
                               [28, 48, 66, 84, 46],
                               [43, 66, 84, 102, 55],
                               [58, 84, 102, 120, 64],
                               [34, 46, 55, 64, 30]])
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.to_float([[14, 43, 34],
                               [43, 84, 55],
                               [34, 55, 30]])
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = resnet_utils.conv2d_same(x, 1, 3, stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.conv2d(x, 1, [3, 3], stride=2, scope='Conv')
    y4_expected = y2_expected

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def testConv2DSameOdd(self):
    n, n2 = 5, 3

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    w = create_test_input(1, 3, 3, 1)
    w = tf.reshape(w, [3, 3, 1, 1])

    tf.get_variable('Conv/weights', initializer=w)
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.conv2d(x, 1, [3, 3], stride=1, scope='Conv')
    y1_expected = tf.cast([[14, 28, 43, 58, 34],
                               [28, 48, 66, 84, 46],
                               [43, 66, 84, 102, 55],
                               [58, 84, 102, 120, 64],
                               [34, 46, 55, 64, 30]], tf.float32)
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.cast([[14, 43, 34],
                               [43, 84, 55],
                               [34, 55, 30]], tf.float32)
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = resnet_utils.conv2d_same(x, 1, 3, stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.conv2d(x, 1, [3, 3], stride=2, scope='Conv')
    y4_expected = y2_expected

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def bit_to_int(x_bit, num_bits, base=2):
  """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.

  Args:
    x_bit: Tensor containing numbers in a particular base to be converted to
      int.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Integer representation of this number.
  """
  x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
  x_labels = [
      x_l[:, i] * tf.to_int32(base)**tf.to_int32(i) for i in range(num_bits)]
  res = sum(x_labels)
  return tf.to_int32(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]))
//  [python]
def bit_to_int(x_bit, num_bits, base=2):
  """Turn x_bit representing numbers bitwise (lower-endian) to int tensor.

  Args:
    x_bit: Tensor containing numbers in a particular base to be converted to
      int.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Integer representation of this number.
  """
  x_l = tf.stop_gradient(tf.cast(tf.reshape(x_bit, [-1, num_bits]), tf.int32))
  x_labels = [
      x_l[:, i] * tf.cast(base**i, tf.int32) for i in range(num_bits)]
  res = sum(x_labels)
  return tf.cast(tf.reshape(res, common_layers.shape_list(x_bit)[:-1]), tf.int32)
//  [python]
def int_to_bit_embed(x_int, num_bits, embedding_size, base=2):
  """Turn x_int into a bitwise (lower-endian) tensor and embed densly."""
  shape = common_layers.shape_list(x_int)
  inputs = int_to_bit(x_int, num_bits, base=base)
  inputs = tf.reshape(inputs, shape[:-1] + [shape[-1] * 8])
  inputs = 2.0 * tf.to_float(inputs) - 1.0  # Move from 0/1 to -1/1.
  return tf.layers.dense(inputs, embedding_size, name="int_to_bit_embed")
//  [python]
def int_to_bit_embed(x_int, num_bits, embedding_size, base=2):
  """Turn x_int into a bitwise (lower-endian) tensor and embed densly."""
  shape = common_layers.shape_list(x_int)
  inputs = int_to_bit(x_int, num_bits, base=base)
  inputs = tf.reshape(inputs, shape[:-1] + [shape[-1] * 8])
  inputs = 2.0 * tf.cast(inputs, tf.float32) - 1.0  # Move from 0/1 to -1/1.
  return tf.layers.dense(inputs, embedding_size, name="int_to_bit_embed")
//  [python]
def _build_model(self, images):
    """Builds the model."""

    # Convert images to floats and normalize them.
    images = tf.to_float(images)
    bs = images.get_shape().as_list()[0]
    images = [
        tf.image.per_image_standardization(tf.squeeze(i))
        for i in tf.split(images, bs)
    ]
    images = tf.concat([tf.expand_dims(i, axis=0) for i in images], axis=0)

    with tf.variable_scope('init'):
      x = self._conv('init_conv', images, 3, 3, 16, self._stride_arr(1))

    strides = [1, 2, 2]
    activate_before_residual = [True, False, False]
    if self._params.use_bottleneck:
      res_func = self._bottleneck_residual
      filters = [16, 64, 128, 256]
    else:
      res_func = self._residual
      filters = [16, 16, 32, 128]

    with tf.variable_scope('unit_1_0'):
      x = res_func(x, filters[0], filters[1], self._stride_arr(strides[0]),
                   activate_before_residual[0])
    for i in xrange(1, self._params.num_residual_units):
      with tf.variable_scope('unit_1_%d' % i):
        x = res_func(x, filters[1], filters[1], self._stride_arr(1), False)

    with tf.variable_scope('unit_2_0'):
      x = res_func(x, filters[1], filters[2], self._stride_arr(strides[1]),
                   activate_before_residual[1])
    for i in xrange(1, self._params.num_residual_units):
      with tf.variable_scope('unit_2_%d' % i):
        x = res_func(x, filters[2], filters[2], self._stride_arr(1), False)

    with tf.variable_scope('unit_3_0'):
      x = res_func(x, filters[2], filters[3], self._stride_arr(strides[2]),
                   activate_before_residual[2])
    for i in xrange(1, self._params.num_residual_units):
      with tf.variable_scope('unit_3_%d' % i):
        x = res_func(x, filters[3], filters[3], self._stride_arr(1), False)

    with tf.variable_scope('unit_last'):
      x = self._batch_norm('final_bn', x)
      x = self._relu(x, self._params.relu_leakiness)

    with tf.variable_scope('pool_logit'):
      x = self._global_avg_pooling(x)

    return x
//  [python]
def _build_model(self, images):
    """Builds the model."""

    # Convert images to floats and normalize them.
    images = tf.cast(images, tf.float32)
    bs = images.get_shape().as_list()[0]
    images = [
        tf.image.per_image_standardization(tf.squeeze(i))
        for i in tf.split(images, bs)
    ]
    images = tf.concat([tf.expand_dims(i, axis=0) for i in images], axis=0)

    with tf.variable_scope('init'):
      x = self._conv('init_conv', images, 3, 3, 16, self._stride_arr(1))

    strides = [1, 2, 2]
    activate_before_residual = [True, False, False]
    if self._params.use_bottleneck:
      res_func = self._bottleneck_residual
      filters = [16, 64, 128, 256]
    else:
      res_func = self._residual
      filters = [16, 16, 32, 128]

    with tf.variable_scope('unit_1_0'):
      x = res_func(x, filters[0], filters[1], self._stride_arr(strides[0]),
                   activate_before_residual[0])
    for i in xrange(1, self._params.num_residual_units):
      with tf.variable_scope('unit_1_%d' % i):
        x = res_func(x, filters[1], filters[1], self._stride_arr(1), False)

    with tf.variable_scope('unit_2_0'):
      x = res_func(x, filters[1], filters[2], self._stride_arr(strides[1]),
                   activate_before_residual[1])
    for i in xrange(1, self._params.num_residual_units):
      with tf.variable_scope('unit_2_%d' % i):
        x = res_func(x, filters[2], filters[2], self._stride_arr(1), False)

    with tf.variable_scope('unit_3_0'):
      x = res_func(x, filters[2], filters[3], self._stride_arr(strides[2]),
                   activate_before_residual[2])
    for i in xrange(1, self._params.num_residual_units):
      with tf.variable_scope('unit_3_%d' % i):
        x = res_func(x, filters[3], filters[3], self._stride_arr(1), False)

    with tf.variable_scope('unit_last'):
      x = self._batch_norm('final_bn', x)
      x = self._relu(x, self._params.relu_leakiness)

    with tf.variable_scope('pool_logit'):
      x = self._global_avg_pooling(x)

    return x
//  [python]
def int_to_bit(x_int, num_bits, base=2):
  """Turn x_int representing numbers into a bitwise (lower-endian) tensor.

  Args:
    x_int: Tensor containing integer to be converted into base notation.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Corresponding number expressed in base.
  """
  x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
  x_labels = [tf.floormod(
      tf.floordiv(tf.to_int32(x_l), tf.to_int32(base)**i), tf.to_int32(base))
              for i in range(num_bits)]
  res = tf.concat(x_labels, axis=-1)
  return tf.to_float(res)
//  [python]
def int_to_bit(x_int, num_bits, base=2):
  """Turn x_int representing numbers into a bitwise (lower-endian) tensor.

  Args:
    x_int: Tensor containing integer to be converted into base notation.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Corresponding number expressed in base.
  """
  x_l = tf.cast(tf.expand_dims(x_int, axis=-1), tf.int32)
  x_labels = [tf.floormod(
      tf.floordiv(tf.cast(x_l, tf.int32), tf.cast(base, tf.int32)**i), tf.cast(base, tf.int32))
              for i in range(num_bits)]
  res = tf.concat(x_labels, axis=-1)
  return tf.cast(res, tf.float32)
//  [python]
def vq_discrete_unbottleneck(x, hidden_size):
  """Simple undiscretization from vector quantized representation."""
  x_shape = common_layers.shape_list(x)
  x = tf.to_float(x)
  bottleneck_size = common_layers.shape_list(x)[-1]
  means, _, _ = get_vq_codebook(bottleneck_size, hidden_size)
  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
  return tf.reshape(result, x_shape[:-1] + [hidden_size])
//  [python]
def vq_discrete_unbottleneck(x, hidden_size):
  """Simple undiscretization from vector quantized representation."""
  x_shape = common_layers.shape_list(x)
  x = tf.cast(x, tf.float32)
  bottleneck_size = common_layers.shape_list(x)[-1]
  means, _, _ = get_vq_codebook(bottleneck_size, hidden_size)
  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
  return tf.reshape(result, x_shape[:-1] + [hidden_size])
//  [python]
def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                             discretize_warmup_steps, mode):
  """Simple discretization through tanh, flip bottleneck_noise many bits."""
  x = tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck")
  d0 = tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x))) - 1.0
  if mode == tf.estimator.ModeKeys.TRAIN:
    x += tf.truncated_normal(
        common_layers.shape_list(x), mean=0.0, stddev=0.2)
  x = tf.tanh(x)
  d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
  if mode == tf.estimator.ModeKeys.TRAIN:
    noise = tf.random_uniform(common_layers.shape_list(x))
    noise = 2.0 * tf.to_float(tf.less(bottleneck_noise, noise)) - 1.0
    d *= noise
  d = common_layers.mix(d, x, discretize_warmup_steps,
                        mode == tf.estimator.ModeKeys.TRAIN)
  return d, d0
//  [python]
def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                             discretize_warmup_steps, mode):
  """Simple discretization through tanh, flip bottleneck_noise many bits."""
  x = tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck")
  d0 = tf.stop_gradient(2.0 * tf.cast(tf.less(0.0, x), tf.float32)) - 1.0
  if mode == tf.estimator.ModeKeys.TRAIN:
    x += tf.truncated_normal(
        common_layers.shape_list(x), mean=0.0, stddev=0.2)
  x = tf.tanh(x)
  d = x + tf.stop_gradient(2.0 * tf.cast(tf.less(0.0, x), tf.float32) - 1.0 - x)
  if mode == tf.estimator.ModeKeys.TRAIN:
    noise = tf.random_uniform(common_layers.shape_list(x))
    noise = 2.0 * tf.cast(tf.less(bottleneck_noise, noise), tf.float32) - 1.0
    d *= noise
  d = common_layers.mix(d, x, discretize_warmup_steps,
                        mode == tf.estimator.ModeKeys.TRAIN)
  return d, d0
//  [python]
def jitter_scale(image, image_size, keypoints, center, scale_range=None,
                 scale_factor=None):
    with tf.name_scope(None, 'jitter_scale', [image, image_size, keypoints]):
        if scale_factor is None:
            scale_factor = tf.random_uniform(
                shape=(1,),
                minval=scale_range[0],
                maxval=scale_range[1],
                dtype=tf.float32
            )
        scale_factor = 2 ** scale_factor
        new_size = tf.to_int32(tf.to_float(image_size) * scale_factor)
        new_image = tf.image.resize_images(image, new_size)

        # This is [height, width] -> [y, x] -> [col, row]
        actual_factor = tf.to_float(
            tf.shape(new_image)[:2]) / tf.to_float(image_size)
        x = keypoints[0, :] * actual_factor[1]
        y = keypoints[1, :] * actual_factor[0]

        cx = tf.cast(center[0], actual_factor.dtype) * actual_factor[1]
        cy = tf.cast(center[1], actual_factor.dtype) * actual_factor[0]

        return new_image, tf.stack([x, y]), tf.cast(
            tf.stack([cx, cy]), tf.int32)
//  [python]
def jitter_scale(image, image_size, keypoints, center, scale_range=None,
                 scale_factor=None):
    with tf.name_scope(None, 'jitter_scale', [image, image_size, keypoints]):
        if scale_factor is None:
            scale_factor = tf.random_uniform(
                shape=(1,),
                minval=scale_range[0],
                maxval=scale_range[1],
                dtype=tf.float32
            )
        scale_factor = 2 ** scale_factor
        new_size = tf.cast(tf.to_float(image_size) * scale_factor, tf.int32)
        new_image = tf.image.resize_images(image, new_size)

        # This is [height, width] -> [y, x] -> [col, row]
        actual_factor = tf.to_float(
            tf.shape(new_image)[:2]) / tf.to_float(image_size)
        x = keypoints[0, :] * actual_factor[1]
        y = keypoints[1, :] * actual_factor[0]

        cx = tf.cast(center[0], actual_factor.dtype) * actual_factor[1]
        cy = tf.cast(center[1], actual_factor.dtype) * actual_factor[0]

        return new_image, tf.stack([x, y]), tf.cast(
            tf.stack([cx, cy]), tf.int32)