jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def diet_expert(x, hidden_size, params):
  """A two-layer feed-forward network with relu activation on hidden layer.

  Uses diet variables.
  Recomputes hidden layer on backprop to save activation memory.

  Args:
    x: a Tensor with shape [batch, io_size]
    hidden_size: an integer
    params: a diet variable HParams object.

  Returns:
    a Tensor with shape [batch, io_size]
  """

  @fn_with_diet_vars(params)
  def diet_expert_internal(x):
    dim = x.get_shape().as_list()[-1]
    h = tf.layers.dense(x, hidden_size, activation=tf.nn.relu, use_bias=False)
    y = tf.layers.dense(h, dim, use_bias=False)
    y *= tf.rsqrt(tf.to_float(dim * hidden_size))
    return y

  return diet_expert_internal(x)
//  [python]
def diet_expert(x, hidden_size, params):
  """A two-layer feed-forward network with relu activation on hidden layer.

  Uses diet variables.
  Recomputes hidden layer on backprop to save activation memory.

  Args:
    x: a Tensor with shape [batch, io_size]
    hidden_size: an integer
    params: a diet variable HParams object.

  Returns:
    a Tensor with shape [batch, io_size]
  """

  @fn_with_diet_vars(params)
  def diet_expert_internal(x):
    dim = x.get_shape().as_list()[-1]
    h = tf.layers.dense(x, hidden_size, activation=tf.nn.relu, use_bias=False)
    y = tf.layers.dense(h, dim, use_bias=False)
    y *= tf.rsqrt(tf.cast(dim * hidden_size, tf.float32))
    return y

  return diet_expert_internal(x)
//  [python]
def sequence_loss_fn(self, chars_logits, chars_labels):
    """Loss function for char sequence.

    Depending on values of hyper parameters it applies label smoothing and can
    also ignore all null chars after the first one.

    Args:
      chars_logits: logits for predicted characters,
        shape=[batch_size, seq_length, num_char_classes];
      chars_labels: ground truth ids of characters,
        shape=[batch_size, seq_length];
      mparams: method hyper parameters.

    Returns:
      A Tensor with shape [batch_size] - the log-perplexity for each sequence.
    """
    mparams = self._mparams['sequence_loss_fn']
    with tf.variable_scope('sequence_loss_fn/SLF'):
      if mparams.label_smoothing > 0:
        smoothed_one_hot_labels = self.label_smoothing_regularization(
            chars_labels, mparams.label_smoothing)
        labels_list = tf.unstack(smoothed_one_hot_labels, axis=1)
      else:
        # NOTE: in case of sparse softmax we are not using one-hot
        # encoding.
        labels_list = tf.unstack(chars_labels, axis=1)

      batch_size, seq_length, _ = chars_logits.shape.as_list()
      if mparams.ignore_nulls:
        weights = tf.ones((batch_size, seq_length), dtype=tf.float32)
      else:
        # Suppose that reject character is the last in the charset.
        reject_char = tf.constant(
            self._params.num_char_classes - 1,
            shape=(batch_size, seq_length),
            dtype=tf.int64)
        known_char = tf.not_equal(chars_labels, reject_char)
        weights = tf.to_float(known_char)

      logits_list = tf.unstack(chars_logits, axis=1)
      weights_list = tf.unstack(weights, axis=1)
      loss = tf.contrib.legacy_seq2seq.sequence_loss(
          logits_list,
          labels_list,
          weights_list,
          softmax_loss_function=get_softmax_loss_fn(mparams.label_smoothing),
          average_across_timesteps=mparams.average_across_timesteps)
      tf.losses.add_loss(loss)
      return loss
//  [python]
def sequence_loss_fn(self, chars_logits, chars_labels):
    """Loss function for char sequence.

    Depending on values of hyper parameters it applies label smoothing and can
    also ignore all null chars after the first one.

    Args:
      chars_logits: logits for predicted characters,
        shape=[batch_size, seq_length, num_char_classes];
      chars_labels: ground truth ids of characters,
        shape=[batch_size, seq_length];
      mparams: method hyper parameters.

    Returns:
      A Tensor with shape [batch_size] - the log-perplexity for each sequence.
    """
    mparams = self._mparams['sequence_loss_fn']
    with tf.variable_scope('sequence_loss_fn/SLF'):
      if mparams.label_smoothing > 0:
        smoothed_one_hot_labels = self.label_smoothing_regularization(
            chars_labels, mparams.label_smoothing)
        labels_list = tf.unstack(smoothed_one_hot_labels, axis=1)
      else:
        # NOTE: in case of sparse softmax we are not using one-hot
        # encoding.
        labels_list = tf.unstack(chars_labels, axis=1)

      batch_size, seq_length, _ = chars_logits.shape.as_list()
      if mparams.ignore_nulls:
        weights = tf.ones((batch_size, seq_length), dtype=tf.float32)
      else:
        # Suppose that reject character is the last in the charset.
        reject_char = tf.constant(
            self._params.num_char_classes - 1,
            shape=(batch_size, seq_length),
            dtype=tf.int64)
        known_char = tf.not_equal(chars_labels, reject_char)
        weights = tf.cast(known_char, tf.float32)

      logits_list = tf.unstack(chars_logits, axis=1)
      weights_list = tf.unstack(weights, axis=1)
      loss = tf.contrib.legacy_seq2seq.sequence_loss(
          logits_list,
          labels_list,
          weights_list,
          softmax_loss_function=get_softmax_loss_fn(mparams.label_smoothing),
          average_across_timesteps=mparams.average_across_timesteps)
      tf.losses.add_loss(loss)
      return loss
//  [python]
def transformer_pointer_prediction_layer(feature_name,
                                         encoder_output,
                                         x,
                                         encoder_decoder_attention_bias,
                                         hparams,
                                         features,
                                         loss_mask,
                                         layer_collection=None):
  """Layer that predicts the start or end token position.

  Args:
    feature_name: 'targets_start_token' or 'targets_end_token'
    encoder_output: [batch_size, input_length, hidden_size] tensor with encoder
      outputs
    x: [batch_size, target_length, 1, hidden_size] tensor with decoder outputs
    encoder_decoder_attention_bias: [batch_size, input_length, target_length]
      attention mask
    hparams: Hyper parameters
    features: Feature dictionary
    loss_mask: [batch_size, target_length] mask for loss computation.
    layer_collection: Layer collection

  Returns:
    (x, logits, loss)
  """
  if isinstance(encoder_output, list):
    pointer_encoder_output = encoder_output[1]
    encoder_output = sum(encoder_output)
  else:
    pointer_encoder_output = encoder_output
  with tf.variable_scope("%s_prediction" % feature_name):
    x = maybe_flatten4d3d(x)
    encoder_decoder_attention_bias = common_layers.flatten4d3d(
        encoder_decoder_attention_bias)
    q = common_attention.compute_attention_component(x, hparams.hidden_size)
    k = common_attention.compute_attention_component(encoder_output,
                                                     hparams.hidden_size)
    # Scaled dot-product attention
    scalar = tf.rsqrt(tf.to_float(common_layers.shape_list(q)[2]))
    logits = tf.matmul(q * scalar, k, transpose_b=True)

    logits += encoder_decoder_attention_bias

    labels = features["%s_raw" % feature_name]
    xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels)
    loss = tf.reduce_sum(xent * loss_mask)

    pointer_out = gather_2d(pointer_encoder_output, labels)
    y = common_layers.layer_preprocess(
        pointer_out, hparams, layer_collection=layer_collection)
    x = common_layers.layer_postprocess(x, y, hparams)
    return x, logits, loss
//  [python]
def transformer_pointer_prediction_layer(feature_name,
                                         encoder_output,
                                         x,
                                         encoder_decoder_attention_bias,
                                         hparams,
                                         features,
                                         loss_mask,
                                         layer_collection=None):
  """Layer that predicts the start or end token position.

  Args:
    feature_name: 'targets_start_token' or 'targets_end_token'
    encoder_output: [batch_size, input_length, hidden_size] tensor with encoder
      outputs
    x: [batch_size, target_length, 1, hidden_size] tensor with decoder outputs
    encoder_decoder_attention_bias: [batch_size, input_length, target_length]
      attention mask
    hparams: Hyper parameters
    features: Feature dictionary
    loss_mask: [batch_size, target_length] mask for loss computation.
    layer_collection: Layer collection

  Returns:
    (x, logits, loss)
  """
  if isinstance(encoder_output, list):
    pointer_encoder_output = encoder_output[1]
    encoder_output = sum(encoder_output)
  else:
    pointer_encoder_output = encoder_output
  with tf.variable_scope("%s_prediction" % feature_name):
    x = maybe_flatten4d3d(x)
    encoder_decoder_attention_bias = common_layers.flatten4d3d(
        encoder_decoder_attention_bias)
    q = common_attention.compute_attention_component(x, hparams.hidden_size)
    k = common_attention.compute_attention_component(encoder_output,
                                                     hparams.hidden_size)
    # Scaled dot-product attention
    scalar = tf.rsqrt(tf.cast(common_layers.shape_list(q)[2], tf.float32))
    logits = tf.matmul(q * scalar, k, transpose_b=True)

    logits += encoder_decoder_attention_bias

    labels = features["%s_raw" % feature_name]
    xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels)
    loss = tf.reduce_sum(xent * loss_mask)

    pointer_out = gather_2d(pointer_encoder_output, labels)
    y = common_layers.layer_preprocess(
        pointer_out, hparams, layer_collection=layer_collection)
    x = common_layers.layer_postprocess(x, y, hparams)
    return x, logits, loss
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
    """Blurs the given image with separable convolution.
    Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.
    Returns:
    A Tensor representing the blurred image.
    """
    radius = tf.to_int32(kernel_size / 2)
    kernel_size = radius * 2 + 1
    x = tf.to_float(tf.range(-radius, radius + 1))
    blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.to_float(sigma), 2.0)))
    blur_filter /= tf.reduce_sum(blur_filter)
    # One vertical and one horizontal filter.
    blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
    blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
    num_channels = tf.shape(image)[-1]
    blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
    blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
    expand_batch_dim = image.shape.ndims == 3
    if expand_batch_dim:
        # Tensorflow requires batched input to convolutions, which we can fake with
        # an extra dimension.
        image = tf.expand_dims(image, axis=0)
    blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
    blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
    if expand_batch_dim:
        blurred = tf.squeeze(blurred, axis=0)
    return blurred
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
    """Blurs the given image with separable convolution.
    Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.
    Returns:
    A Tensor representing the blurred image.
    """
    radius = tf.cast(kernel_size / 2, tf.int32)
    kernel_size = radius * 2 + 1
    x = tf.cast(tf.range(-radius, radius + 1), tf.float32)
    blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, tf.float32), 2.0)))
    blur_filter /= tf.reduce_sum(blur_filter)
    # One vertical and one horizontal filter.
    blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
    blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
    num_channels = tf.shape(image)[-1]
    blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
    blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
    expand_batch_dim = image.shape.ndims == 3
    if expand_batch_dim:
        # Tensorflow requires batched input to convolutions, which we can fake with
        # an extra dimension.
        image = tf.expand_dims(image, axis=0)
    blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
    blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
    if expand_batch_dim:
        blurred = tf.squeeze(blurred, axis=0)
    return blurred
//  [python]
def update_variable(self, var, grad_var):
    """Update the variable and its slots."""
    params = self.params
    global_step = tf.to_float(self.global_step) + 1

    # compute learning rate
    lrate = params.learning_rate
    if params.learning_rate_decay_scheme == "noam":
      lrate *= tf.minimum(global_step * params.learning_rate_warmup_steps**-1.5,
                          global_step**-0.5)
    else:
      assert params.learning_rate_decay_scheme == "none"
      lrate *= tf.minimum(global_step / params.learning_rate_warmup_steps, 1.0)

    # compute adjustment due to second moment
    slots = params.slots[var.op.name]
    grad_squared = tf.square(grad_var)
    beta2_pow = tf.pow(params.beta2, global_step)
    if params.factored_second_moment_accumulator and len(var.shape) == 2:
      vr_update = tf.assign(slots["adam_vr"], slots["adam_vr"] * params.beta2 +
                            tf.reduce_mean(grad_squared, 1, keepdims=True) *
                            (1.0 - params.beta2))
      vc_update = tf.assign(slots["adam_vc"], slots["adam_vc"] * params.beta2 +
                            tf.reduce_mean(grad_squared, 0, keepdims=True) *
                            (1.0 - params.beta2))
      with tf.control_dependencies([vr_update, vc_update]):
        vr = tf.sqrt(slots["adam_vr"] / (1.0 - beta2_pow)) + params.epsilon
        vc = tf.sqrt(slots["adam_vc"] / (1.0 - beta2_pow)) + params.epsilon
        vc /= tf.reduce_mean(vc)
        denom = vr * vc
    else:
      v_update = tf.assign(slots["adam_v"],
                           slots["adam_v"] * params.beta2 + grad_squared *
                           (1.0 - params.beta2))
      with tf.control_dependencies([v_update]):
        denom = tf.sqrt(slots["adam_v"] / (1.0 - beta2_pow)) + params.epsilon

    # compute momentum if applicable
    if params.beta1 != 0.0:
      m_update = tf.assign(slots["adam_m"],
                           slots["adam_m"] * params.beta1 + grad_var *
                           (1.0 - params.beta1))
      with tf.control_dependencies([m_update]):
        grad_var = slots["adam_m"]

    # update var
    subtrahend = lrate * grad_var / denom
    new_val = _quantize(_dequantize(var, params) - subtrahend, params)
    return tf.assign(var, new_val)
//  [python]
def update_variable(self, var, grad_var):
    """Update the variable and its slots."""
    params = self.params
    global_step = tf.cast(self.global_step, tf.float32) + 1

    # compute learning rate
    lrate = params.learning_rate
    if params.learning_rate_decay_scheme == "noam":
      lrate *= tf.minimum(global_step * params.learning_rate_warmup_steps**-1.5,
                          global_step**-0.5)
    else:
      assert params.learning_rate_decay_scheme == "none"
      lrate *= tf.minimum(global_step / params.learning_rate_warmup_steps, 1.0)

    # compute adjustment due to second moment
    slots = params.slots[var.op.name]
    grad_squared = tf.square(grad_var)
    beta2_pow = tf.pow(params.beta2, global_step)
    if params.factored_second_moment_accumulator and len(var.shape) == 2:
      vr_update = tf.assign(slots["adam_vr"], slots["adam_vr"] * params.beta2 +
                            tf.reduce_mean(grad_squared, 1, keepdims=True) *
                            (1.0 - params.beta2))
      vc_update = tf.assign(slots["adam_vc"], slots["adam_vc"] * params.beta2 +
                            tf.reduce_mean(grad_squared, 0, keepdims=True) *
                            (1.0 - params.beta2))
      with tf.control_dependencies([vr_update, vc_update]):
        vr = tf.sqrt(slots["adam_vr"] / (1.0 - beta2_pow)) + params.epsilon
        vc = tf.sqrt(slots["adam_vc"] / (1.0 - beta2_pow)) + params.epsilon
        vc /= tf.reduce_mean(vc)
        denom = vr * vc
    else:
      v_update = tf.assign(slots["adam_v"],
                           slots["adam_v"] * params.beta2 + grad_squared *
                           (1.0 - params.beta2))
      with tf.control_dependencies([v_update]):
        denom = tf.sqrt(slots["adam_v"] / (1.0 - beta2_pow)) + params.epsilon

    # compute momentum if applicable
    if params.beta1 != 0.0:
      m_update = tf.assign(slots["adam_m"],
                           slots["adam_m"] * params.beta1 + grad_var *
                           (1.0 - params.beta1))
      with tf.control_dependencies([m_update]):
        grad_var = slots["adam_m"]

    # update var
    subtrahend = lrate * grad_var / denom
    new_val = _quantize(_dequantize(var, params) - subtrahend, params)
    return tf.assign(var, new_val)
//  [python]
def input_generator(zs_file, batch_size):

    time_placeholder = tf.placeholder(dtype=tf.float32, shape=())
    attribute_placeholder = tf.placeholder(dtype=tf.float32, shape=(100,))

    def l2_norm(x):
        euclidean_norm = tf.sqrt(tf.reduce_sum(tf.square(x)))
        return euclidean_norm

    def dot_product(x, y):
        return tf.reduce_sum(tf.mul(x, y))

    def slerp(initial, final, progress):
        omega = tf.acos(dot_product(initial / l2_norm(initial), final / l2_norm(final)))
        so = tf.sin(omega)
        return tf.sin((1.0-progress)*omega) / so * initial + tf.sin(progress*omega)/so * final

    with open(zs_file, 'rb') as f:
        zs = pickle.load(f)
        img_count = len(zs)
        zs = tf.constant(zs, dtype=tf.float32)

    tensors = []

    epoch = tf.to_int32(time_placeholder)
    indices = tf.range(batch_size)
    indices_init = (indices * batch_size + epoch) % img_count
    indices_final = (indices_init + 1) % img_count

    for i in xrange(batch_size):
        z_init = zs[indices_init[i]]
        z_final = zs[indices_final[i]]

        progress = tf.mod(time_placeholder, 1)

        # progress = tf.Print(progress, [progress])

        z = slerp(z_init, z_final, progress)

        tensors.append(z)

    batch = tf.pack(tensors) + attribute_placeholder

    return batch, time_placeholder, attribute_placeholder
//  [python]
def input_generator(zs_file, batch_size):

    time_placeholder = tf.placeholder(dtype=tf.float32, shape=())
    attribute_placeholder = tf.placeholder(dtype=tf.float32, shape=(100,))

    def l2_norm(x):
        euclidean_norm = tf.sqrt(tf.reduce_sum(tf.square(x)))
        return euclidean_norm

    def dot_product(x, y):
        return tf.reduce_sum(tf.mul(x, y))

    def slerp(initial, final, progress):
        omega = tf.acos(dot_product(initial / l2_norm(initial), final / l2_norm(final)))
        so = tf.sin(omega)
        return tf.sin((1.0-progress)*omega) / so * initial + tf.sin(progress*omega)/so * final

    with open(zs_file, 'rb') as f:
        zs = pickle.load(f)
        img_count = len(zs)
        zs = tf.constant(zs, dtype=tf.float32)

    tensors = []

    epoch = tf.cast(time_placeholder, tf.int32)
    indices = tf.range(batch_size)
    indices_init = (indices * batch_size + epoch) % img_count
    indices_final = (indices_init + 1) % img_count

    for i in xrange(batch_size):
        z_init = zs[indices_init[i]]
        z_final = zs[indices_final[i]]

        progress = tf.mod(time_placeholder, 1)

        # progress = tf.Print(progress, [progress])

        z = slerp(z_init, z_final, progress)

        tensors.append(z)

    batch = tf.pack(tensors) + attribute_placeholder

    return batch, time_placeholder, attribute_placeholder
//  [python]
def _sequence_correct(labels, predictions):
  """Computes a per-example sequence accuracy."""
  target_decode_steps = decode_utils.decode_steps_from_labels(
      labels, trim_start_symbol=True)
  predicted_decode_steps = decode_utils.decode_steps_from_predictions(
      predictions)

  decode_utils.assert_shapes_match(target_decode_steps, predicted_decode_steps)

  equal_tokens = decode_utils.compare_decode_steps(target_decode_steps,
                                                   predicted_decode_steps)
  target_len = labels["target_len"] - 1
  loss_mask = tf.sequence_mask(
      lengths=tf.to_int32(target_len),
      maxlen=tf.to_int32(tf.shape(equal_tokens)[1]))
  equal_tokens = tf.logical_or(equal_tokens, tf.logical_not(loss_mask))
  all_equal = tf.cast(tf.reduce_all(equal_tokens, 1), tf.float32)
  return all_equal
//  [python]
def _sequence_correct(labels, predictions):
  """Computes a per-example sequence accuracy."""
  target_decode_steps = decode_utils.decode_steps_from_labels(
      labels, trim_start_symbol=True)
  predicted_decode_steps = decode_utils.decode_steps_from_predictions(
      predictions)

  decode_utils.assert_shapes_match(target_decode_steps, predicted_decode_steps)

  equal_tokens = decode_utils.compare_decode_steps(target_decode_steps,
                                                   predicted_decode_steps)
  target_len = labels["target_len"] - 1
  loss_mask = tf.sequence_mask(
      lengths=tf.cast(target_len, tf.int32),
      maxlen=tf.cast(tf.shape(equal_tokens)[1], tf.int32))
  equal_tokens = tf.logical_or(equal_tokens, tf.logical_not(loss_mask))
  all_equal = tf.cast(tf.reduce_all(equal_tokens, 1), tf.float32)
  return all_equal
//  [python]
def main(_):
  checkpoint_dir = os.path.join(FLAGS.exp_dir, 'log')
  # eval_dir = os.path.join(FLAGS.exp_dir, 'log/eval')
  model_config, _, _ = get_configs_from_exp_dir()

  model = model_builder.build(model_config, is_training=False)

  input_image_str_tensor = tf.placeholder(
    dtype=tf.string,
    shape=[])
  input_image_tensor = tf.image.decode_jpeg(
    input_image_str_tensor,
    channels=3,
  )
  resized_image_tensor = tf.image.resize_images(
    tf.to_float(input_image_tensor),
    [64, 256])

  predictions_dict = model.predict(tf.expand_dims(resized_image_tensor, 0))
  recognitions = model.postprocess(predictions_dict)
  recognition_text = recognitions['text'][0]
  control_points = predictions_dict['control_points'],
  rectified_images = predictions_dict['rectified_images']

  saver = tf.train.Saver(tf.global_variables())
  checkpoint = os.path.join(FLAGS.exp_dir, 'log/model.ckpt')

  fetches = {
    'original_image': input_image_tensor,
    'recognition_text': recognition_text,
    'control_points': predictions_dict['control_points'],
    'rectified_images': predictions_dict['rectified_images'],
  }

  with open(FLAGS.input_image, 'rb') as f:
    input_image_str = f.read()

  with tf.Session() as sess:
    sess.run([
      tf.global_variables_initializer(),
      tf.local_variables_initializer(),
      tf.tables_initializer()])
    saver.restore(sess, checkpoint)
    sess_outputs = sess.run(fetches, feed_dict={input_image_str_tensor: input_image_str})

  print('Recognized text: {}'.format(sess_outputs['recognition_text'].decode('utf-8')))
  
  rectified_image = sess_outputs['rectified_images'][0]
  rectified_image_pil = Image.fromarray((128 * (rectified_image + 1.0)).astype(np.uint8))
  input_image_dir = os.path.dirname(FLAGS.input_image)
  rectified_image_save_path = os.path.join(input_image_dir, 'rectified_image.jpg')
  rectified_image_pil.save(rectified_image_save_path)
  print('Rectified image saved to {}'.format(rectified_image_save_path))
//  [python]
def main(_):
  checkpoint_dir = os.path.join(FLAGS.exp_dir, 'log')
  # eval_dir = os.path.join(FLAGS.exp_dir, 'log/eval')
  model_config, _, _ = get_configs_from_exp_dir()

  model = model_builder.build(model_config, is_training=False)

  input_image_str_tensor = tf.placeholder(
    dtype=tf.string,
    shape=[])
  input_image_tensor = tf.image.decode_jpeg(
    input_image_str_tensor,
    channels=3,
  )
  resized_image_tensor = tf.image.resize_images(
    tf.cast(input_image_tensor, tf.float32),
    [64, 256])

  predictions_dict = model.predict(tf.expand_dims(resized_image_tensor, 0))
  recognitions = model.postprocess(predictions_dict)
  recognition_text = recognitions['text'][0]
  control_points = predictions_dict['control_points'],
  rectified_images = predictions_dict['rectified_images']

  saver = tf.train.Saver(tf.global_variables())
  checkpoint = os.path.join(FLAGS.exp_dir, 'log/model.ckpt')

  fetches = {
    'original_image': input_image_tensor,
    'recognition_text': recognition_text,
    'control_points': predictions_dict['control_points'],
    'rectified_images': predictions_dict['rectified_images'],
  }

  with open(FLAGS.input_image, 'rb') as f:
    input_image_str = f.read()

  with tf.Session() as sess:
    sess.run([
      tf.global_variables_initializer(),
      tf.local_variables_initializer(),
      tf.tables_initializer()])
    saver.restore(sess, checkpoint)
    sess_outputs = sess.run(fetches, feed_dict={input_image_str_tensor: input_image_str})

  print('Recognized text: {}'.format(sess_outputs['recognition_text'].decode('utf-8')))
  
  rectified_image = sess_outputs['rectified_images'][0]
  rectified_image_pil = Image.fromarray((128 * (rectified_image + 1.0)).astype(np.uint8))
  input_image_dir = os.path.dirname(FLAGS.input_image)
  rectified_image_save_path = os.path.join(input_image_dir, 'rectified_image.jpg')
  rectified_image_pil.save(rectified_image_save_path)
  print('Rectified image saved to {}'.format(rectified_image_save_path))
//  [python]
def dataset(directory, images_file, labels_file):
  """Download and parse MNIST dataset."""

  images_file = download(directory, images_file)
  labels_file = download(directory, labels_file)

  check_image_file_header(images_file)
  check_labels_file_header(labels_file)

  def decode_image(image):
    # Normalize from [0, 255] to [0.0, 1.0]
    image = tf.decode_raw(image, tf.uint8)
    image = tf.cast(image, tf.float32)
    image = tf.reshape(image, [784])
    return image / 255.0

  def decode_label(label):
    label = tf.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
    label = tf.reshape(label, [])  # label is a scalar
    return tf.to_int32(label)

  images = tf.data.FixedLengthRecordDataset(
      images_file, 28 * 28, header_bytes=16).map(decode_image)
  labels = tf.data.FixedLengthRecordDataset(
      labels_file, 1, header_bytes=8).map(decode_label)
  return tf.data.Dataset.zip((images, labels))
//  [python]
def dataset(directory, images_file, labels_file):
  """Download and parse MNIST dataset."""

  images_file = download(directory, images_file)
  labels_file = download(directory, labels_file)

  check_image_file_header(images_file)
  check_labels_file_header(labels_file)

  def decode_image(image):
    # Normalize from [0, 255] to [0.0, 1.0]
    image = tf.decode_raw(image, tf.uint8)
    image = tf.cast(image, tf.float32)
    image = tf.reshape(image, [784])
    return image / 255.0

  def decode_label(label):
    label = tf.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
    label = tf.reshape(label, [])  # label is a scalar
    return tf.cast(label, tf.int32)

  images = tf.data.FixedLengthRecordDataset(
      images_file, 28 * 28, header_bytes=16).map(decode_image)
  labels = tf.data.FixedLengthRecordDataset(
      labels_file, 1, header_bytes=8).map(decode_label)
  return tf.data.Dataset.zip((images, labels))
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
    """Blurs the given image with separable convolution.
    Args:
      image: Tensor of shape [height, width, channels] and dtype float to blur.
      kernel_size: Integer Tensor for the size of the blur kernel. This is should
        be an odd number. If it is an even number, the actual kernel size will be
        size + 1.
      sigma: Sigma value for gaussian operator.
      padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.
    Returns:
      A Tensor representing the blurred image.
    """
    radius = tf.to_int32(kernel_size / 2)
    kernel_size = radius * 2 + 1
    x = tf.to_float(tf.range(-radius, radius + 1))
    blur_filter = tf.exp(
        -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.to_float(sigma), 2.0)))
    blur_filter /= tf.reduce_sum(blur_filter)
    # One vertical and one horizontal filter.
    blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
    blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
    num_channels = tf.shape(image)[-1]
    blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
    blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
    expand_batch_dim = image.shape.ndims == 3
    if expand_batch_dim:
        # Tensorflow requires batched input to convolutions, which we can fake with
        # an extra dimension.
        image = tf.expand_dims(image, axis=0)
    blurred = tf.nn.depthwise_conv2d(
        image, blur_h, strides=[1, 1, 1, 1], padding=padding)
    blurred = tf.nn.depthwise_conv2d(
        blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
    if expand_batch_dim:
        blurred = tf.squeeze(blurred, axis=0)
    return blurred
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
    """Blurs the given image with separable convolution.
    Args:
      image: Tensor of shape [height, width, channels] and dtype float to blur.
      kernel_size: Integer Tensor for the size of the blur kernel. This is should
        be an odd number. If it is an even number, the actual kernel size will be
        size + 1.
      sigma: Sigma value for gaussian operator.
      padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.
    Returns:
      A Tensor representing the blurred image.
    """
    radius = tf.cast(kernel_size / 2, tf.int32)
    kernel_size = radius * 2 + 1
    x = tf.cast(tf.range(-radius, radius + 1), tf.float32)
    blur_filter = tf.exp(
        -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, tf.float32), 2.0)))
    blur_filter /= tf.reduce_sum(blur_filter)
    # One vertical and one horizontal filter.
    blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
    blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
    num_channels = tf.shape(image)[-1]
    blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
    blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
    expand_batch_dim = image.shape.ndims == 3
    if expand_batch_dim:
        # Tensorflow requires batched input to convolutions, which we can fake with
        # an extra dimension.
        image = tf.expand_dims(image, axis=0)
    blurred = tf.nn.depthwise_conv2d(
        image, blur_h, strides=[1, 1, 1, 1], padding=padding)
    blurred = tf.nn.depthwise_conv2d(
        blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
    if expand_batch_dim:
        blurred = tf.squeeze(blurred, axis=0)
    return blurred
//  [python]
def _build_train_op(self):
        """build train_op based on either REINFORCE or PPO
        """
        self.advantage = tf.placeholder(shape=(None, 1), dtype=tf.float32, name="advantage")
        self.reward = tf.placeholder(shape=(None, 1), dtype=tf.float32, name="reward")

        normalize = tf.to_float(self.num_layers * (self.num_layers - 1) / 2)
        self.skip_rate = tf.to_float(self.skip_count) / normalize

        self.input_arc_onehot = self.convert_arc_to_onehot(self)
        self.old_probs = [tf.placeholder(shape=self.onehot_probs[i].shape, dtype=tf.float32, name="old_prob_%i" % i) for
                          i in range(len(self.onehot_probs))]
        if self.skip_weight is not None:
            self.loss += self.skip_weight * tf.reduce_mean(self.onehot_skip_penaltys)
        if self.use_ppo_loss:
            self.loss += proximal_policy_optimization_loss(
                curr_prediction=self.onehot_probs,
                curr_onehot=self.input_arc_onehot,
                old_prediction=self.old_probs,
                old_onehotpred=self.input_arc_onehot,
                rewards=self.reward,
                advantage=self.advantage,
                clip_val=0.2)
        else:
            self.loss += tf.reshape(tf.tensordot(self.onehot_log_prob, self.advantage, axes=1), [])

        self.kl_div, self.ent = get_kl_divergence_n_entropy(curr_prediction=self.onehot_probs,
                                                            old_prediction=self.old_probs,
                                                            curr_onehot=self.input_arc_onehot,
                                                            old_onehotpred=self.input_arc_onehot)
        self.train_step = tf.Variable(
            0, dtype=tf.int32, trainable=False, name="train_step")
        tf_variables = [var
                        for var in tf.trainable_variables() if var.name.startswith(self.name)]

        self.train_op, self.lr, self.grad_norm, self.optimizer = get_keras_train_ops(
            loss=self.loss,
            tf_variables=tf_variables,
            optim_algo=self.optim_algo
        )
//  [python]
def _build_train_op(self):
        """build train_op based on either REINFORCE or PPO
        """
        self.advantage = tf.placeholder(shape=(None, 1), dtype=tf.float32, name="advantage")
        self.reward = tf.placeholder(shape=(None, 1), dtype=tf.float32, name="reward")

        normalize = tf.cast(self.num_layers * (self.num_layers - 1) / 2, tf.float32)
        self.skip_rate = tf.cast(self.skip_count, tf.float32) / normalize

        self.input_arc_onehot = self.convert_arc_to_onehot(self)
        self.old_probs = [tf.placeholder(shape=self.onehot_probs[i].shape, dtype=tf.float32, name="old_prob_%i" % i) for
                          i in range(len(self.onehot_probs))]
        if self.skip_weight is not None:
            self.loss += self.skip_weight * tf.reduce_mean(self.onehot_skip_penaltys)
        if self.use_ppo_loss:
            self.loss += proximal_policy_optimization_loss(
                curr_prediction=self.onehot_probs,
                curr_onehot=self.input_arc_onehot,
                old_prediction=self.old_probs,
                old_onehotpred=self.input_arc_onehot,
                rewards=self.reward,
                advantage=self.advantage,
                clip_val=0.2)
        else:
            self.loss += tf.reshape(tf.tensordot(self.onehot_log_prob, self.advantage, axes=1), [])

        self.kl_div, self.ent = get_kl_divergence_n_entropy(curr_prediction=self.onehot_probs,
                                                            old_prediction=self.old_probs,
                                                            curr_onehot=self.input_arc_onehot,
                                                            old_onehotpred=self.input_arc_onehot)
        self.train_step = tf.Variable(
            0, dtype=tf.int32, trainable=False, name="train_step")
        tf_variables = [var
                        for var in tf.trainable_variables() if var.name.startswith(self.name)]

        self.train_op, self.lr, self.grad_norm, self.optimizer = get_keras_train_ops(
            loss=self.loss,
            tf_variables=tf_variables,
            optim_algo=self.optim_algo
        )
//  [python]
def embedding_to_padding(emb):
  """Calculates the padding mask based on which embeddings are all zero.

  We have hacked symbol_modality to return all-zero embeddings for padding.

  Args:
    emb: a Tensor with shape [..., depth].

  Returns:
    a float Tensor with shape [...]. Each element is 1 if its corresponding
    embedding vector is all zero, and is 0 otherwise.
  """
  emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1)
  return tf.to_float(tf.equal(emb_sum, 0.0))
//  [python]
def embedding_to_padding(emb):
  """Calculates the padding mask based on which embeddings are all zero.

  We have hacked symbol_modality to return all-zero embeddings for padding.

  Args:
    emb: a Tensor with shape [..., depth].

  Returns:
    a float Tensor with shape [...]. Each element is 1 if its corresponding
    embedding vector is all zero, and is 0 otherwise.
  """
  emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1)
  return tf.cast(tf.equal(emb_sum, 0.0), tf.float32)
//  [python]
def get_timing_signal_1d(length,
                         channels,
                         min_timescale=1.0,
                         max_timescale=1.0e4,
                         start_index=0):
  """Gets a bunch of sinusoids of different frequencies.

  Each channel of the input Tensor is incremented by a sinusoid of a different
  frequency and phase.

  This allows attention to learn to use absolute and relative positions.
  Timing signals should be added to some precursors of both the query and the
  memory inputs to attention.

  The use of relative position is possible because sin(x+y) and cos(x+y) can be
  expressed in terms of y, sin(x) and cos(x).

  In particular, we use a geometric sequence of timescales starting with
  min_timescale and ending with max_timescale.  The number of different
  timescales is equal to channels / 2. For each timescale, we
  generate the two sinusoidal signals sin(timestep/timescale) and
  cos(timestep/timescale).  All of these sinusoids are concatenated in
  the channels dimension.

  Args:
    length: scalar, length of timing signal sequence.
    channels: scalar, size of timing embeddings to create. The number of
        different timescales is equal to channels / 2.
    min_timescale: a float
    max_timescale: a float
    start_index: index of first position

  Returns:
    a Tensor of timing signals [1, length, channels]
  """
  position = tf.to_float(tf.range(length) + start_index)
  num_timescales = channels // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      tf.maximum(tf.to_float(num_timescales) - 1, 1))
  inv_timescales = min_timescale * tf.exp(
      tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
  scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
  # Please note that this slightly differs from the published paper.
  # See a discussion here: https://github.com/tensorflow/tensor2tensor/pull/177
  signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
  signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
  signal = tf.reshape(signal, [1, length, channels])
  return signal
//  [python]
def get_timing_signal_1d(length,
                         channels,
                         min_timescale=1.0,
                         max_timescale=1.0e4,
                         start_index=0):
  """Gets a bunch of sinusoids of different frequencies.

  Each channel of the input Tensor is incremented by a sinusoid of a different
  frequency and phase.

  This allows attention to learn to use absolute and relative positions.
  Timing signals should be added to some precursors of both the query and the
  memory inputs to attention.

  The use of relative position is possible because sin(x+y) and cos(x+y) can be
  expressed in terms of y, sin(x) and cos(x).

  In particular, we use a geometric sequence of timescales starting with
  min_timescale and ending with max_timescale.  The number of different
  timescales is equal to channels / 2. For each timescale, we
  generate the two sinusoidal signals sin(timestep/timescale) and
  cos(timestep/timescale).  All of these sinusoids are concatenated in
  the channels dimension.

  Args:
    length: scalar, length of timing signal sequence.
    channels: scalar, size of timing embeddings to create. The number of
        different timescales is equal to channels / 2.
    min_timescale: a float
    max_timescale: a float
    start_index: index of first position

  Returns:
    a Tensor of timing signals [1, length, channels]
  """
  position = tf.cast(tf.range(length) + start_index, tf.float32)
  num_timescales = channels // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      tf.maximum(tf.cast(num_timescales - 1, tf.float32), 1))
  inv_timescales = min_timescale * tf.exp(
      tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)
  scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
  # Please note that this slightly differs from the published paper.
  # See a discussion here: https://github.com/tensorflow/tensor2tensor/pull/177
  signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
  signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
  signal = tf.reshape(signal, [1, length, channels])
  return signal
//  [python]
def padding_to_length(padding):
  """Calculate the length of mask based on padding.

  Args:
    padding: a Tensor with shape [..., length].
  Returns:
    a Tensor with shape [...].
  """
  non_padding = 1.0 - padding
  return tf.to_int32(tf.reduce_sum(non_padding, axis=-1))
//  [python]
def padding_to_length(padding):
  """Calculate the length of mask based on padding.

  Args:
    padding: a Tensor with shape [..., length].
  Returns:
    a Tensor with shape [...].
  """
  non_padding = 1.0 - padding
  return tf.cast(tf.reduce_sum(non_padding, axis=-1), tf.int32)
//  [python]
def add_positional_embedding(x, max_length, name=None, positions=None):
  """Adds positional embedding.

  Args:
    x: Tensor with shape [batch, length, depth].
    max_length: int representing static maximum size of any dimension.
    name: str representing name of the embedding tf.Variable.
    positions: Tensor with shape [batch, length].

  Returns:
    Tensor of same shape as x.
  """
  with tf.name_scope("add_positional_embedding"):
    _, length, depth = common_layers.shape_list(x)
    var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
    if positions is None:
      pad_length = tf.maximum(0, length - max_length)
      sliced = tf.cond(
          tf.less(length, max_length),
          lambda: tf.slice(var, [0, 0], [length, -1]),
          lambda: tf.pad(var, [[0, pad_length], [0, 0]]))
      return x + tf.expand_dims(sliced, 0)
    else:
      return x + tf.gather(var, tf.to_int32(positions))
//  [python]
def add_positional_embedding(x, max_length, name=None, positions=None):
  """Adds positional embedding.

  Args:
    x: Tensor with shape [batch, length, depth].
    max_length: int representing static maximum size of any dimension.
    name: str representing name of the embedding tf.Variable.
    positions: Tensor with shape [batch, length].

  Returns:
    Tensor of same shape as x.
  """
  with tf.name_scope("add_positional_embedding"):
    _, length, depth = common_layers.shape_list(x)
    var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
    if positions is None:
      pad_length = tf.maximum(0, length - max_length)
      sliced = tf.cond(
          tf.less(length, max_length),
          lambda: tf.slice(var, [0, 0], [length, -1]),
          lambda: tf.pad(var, [[0, pad_length], [0, 0]]))
      return x + tf.expand_dims(sliced, 0)
    else:
      return x + tf.gather(var, tf.cast(positions, tf.int32))
//  [python]
def add_timing_signals_given_positions(x,
                                       positions,
                                       min_timescale=1.0,
                                       max_timescale=1.0e4):
  """Adds sinusoids of diff frequencies to a Tensor, with timing positions given.

  Args:
    x: a Tensor with shape [batch, length, channels]
    positions: a list of positions, each of which can either be a Tensor of
      shape [batch, length] or None for a default of (0..length]
    min_timescale: a float
    max_timescale: a float

  Returns:
    a Tensor the same shape as x.
  """
  shape = common_layers.shape_list(x)
  batch = shape[0]
  length = shape[1]
  channels = shape[2]
  num_dims = len(positions)
  num_timescales = channels // (num_dims * 2)
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (tf.to_float(num_timescales) - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
  for dim, position in enumerate(positions):
    if position is None:
      # Create a [batch, length] Tensor of incrementing positions 0..length-1.
      position = tf.tile(
          tf.transpose(tf.expand_dims(tf.range(0, length), axis=1)), [batch, 1])
    scaled_time = (
        tf.expand_dims(tf.to_float(position), 2) *
        tf.expand_dims(tf.expand_dims(inv_timescales, 0), 0))
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
    prepad = dim * 2 * num_timescales
    postpad = channels - (dim + 1) * 2 * num_timescales
    signal = tf.pad(signal, [[0, 0], [0, 0], [prepad, postpad]])
    signal = common_layers.cast_like(signal, x)
    x += signal
  return x
//  [python]
def add_timing_signals_given_positions(x,
                                       positions,
                                       min_timescale=1.0,
                                       max_timescale=1.0e4):
  """Adds sinusoids of diff frequencies to a Tensor, with timing positions given.

  Args:
    x: a Tensor with shape [batch, length, channels]
    positions: a list of positions, each of which can either be a Tensor of
      shape [batch, length] or None for a default of (0..length]
    min_timescale: a float
    max_timescale: a float

  Returns:
    a Tensor the same shape as x.
  """
  shape = common_layers.shape_list(x)
  batch = shape[0]
  length = shape[1]
  channels = shape[2]
  num_dims = len(positions)
  num_timescales = channels // (num_dims * 2)
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (tf.cast(num_timescales, tf.float32) - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)
  for dim, position in enumerate(positions):
    if position is None:
      # Create a [batch, length] Tensor of incrementing positions 0..length-1.
      position = tf.tile(
          tf.transpose(tf.expand_dims(tf.range(0, length), axis=1)), [batch, 1])
    scaled_time = (
        tf.expand_dims(tf.cast(position, tf.float32), 2) *
        tf.expand_dims(tf.expand_dims(inv_timescales, 0), 0))
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
    prepad = dim * 2 * num_timescales
    postpad = channels - (dim + 1) * 2 * num_timescales
    signal = tf.pad(signal, [[0, 0], [0, 0], [prepad, postpad]])
    signal = common_layers.cast_like(signal, x)
    x += signal
  return x
//  [python]
def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
  """Gets edge vectors for the edge types in the adjacency matrix.

  Args:
    adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
    num_edge_types: Number of different edge types
    depth: Number of channels
    name: a string
  Returns:
    A [batch, num_nodes, num_nodes, depth] vector of tensors
  """
  with tf.variable_scope(name, default_name="edge_vectors"):
    att_adj_vectors_shape = [num_edge_types, depth]
    adjacency_matrix_shape = common_layers.shape_list(adjacency_matrix)
    adj_vectors = (
        tf.get_variable(
            "adj_vectors",
            att_adj_vectors_shape,
            initializer=tf.random_normal_initializer(0, depth**-0.5)) *
        (depth**0.5))
    # Avoiding gathers so that it works on TPUs
    # adjacency_matrix_one_hot has shape
    # [batch, num_nodes, num_nodes, num_edge_types]

    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)

    att_adj_vectors = tf.matmul(
        tf.reshape(tf.to_float(adjacency_matrix_one_hot), [-1, num_edge_types]),
        adj_vectors)
    return tf.reshape(att_adj_vectors,
                      [adjacency_matrix_shape[0], adjacency_matrix_shape[1],
                       adjacency_matrix_shape[2], depth])