# Lint as: python3
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for model_export."""
import os

import numpy as np
from absl.testing import flagsaver
import tensorflow as tf
from tensorflow.compat.v1 import flags

import common_flags
import model_export

_CHECKPOINT = 'model.ckpt-399731'
_CHECKPOINT_URL = (
    'http://download.tensorflow.org/models/attention_ocr_2017_08_09.tar.gz')


def _clean_up():
  tf.io.gfile.rmtree(tf.compat.v1.test.get_temp_dir())


def _create_tf_example_string(image):
  """Create a serialized tf.Example proto for feeding the model."""
  example = tf.train.Example()
  example.features.feature['image/encoded'].float_list.value.extend(
      list(np.reshape(image, (-1))))
  return example.SerializeToString()


class AttentionOcrExportTest(tf.test.TestCase):
  """Tests for model_export.export_model."""

  def setUp(self):
    for suffix in ['.meta', '.index', '.data-00000-of-00001']:
      filename = _CHECKPOINT + suffix
      self.assertTrue(
          tf.io.gfile.exists(filename),
          msg='Missing checkpoint file %s. '
          'Please download and extract it from %s' %
          (filename, _CHECKPOINT_URL))
    flags.FLAGS.dataset_name = 'fsns'
    flags.FLAGS.checkpoint = _CHECKPOINT
    flags.FLAGS.dataset_dir = os.path.join(
        os.path.dirname(__file__), 'datasets/testdata/fsns')
    tf.test.TestCase.setUp(self)
    _clean_up()
    self.export_dir = os.path.join(
        tf.compat.v1.test.get_temp_dir(), 'exported_model')
    self.minimal_output_signature = {
        'predictions': 'AttentionOcr_v1/predicted_chars:0',
        'scores': 'AttentionOcr_v1/predicted_scores:0',
        'predicted_length': 'AttentionOcr_v1/predicted_length:0',
        'predicted_text': 'AttentionOcr_v1/predicted_text:0',
        'predicted_conf': 'AttentionOcr_v1/predicted_conf:0',
        'normalized_seq_conf': 'AttentionOcr_v1/normalized_seq_conf:0'
    }

  def create_input_feed(self, graph_def, serving):
    """Returns the input feed for the model.

    Creates random images, according to the size specified by dataset_name,
    format it in the correct way depending on whether the model was exported
    for serving, and return the correctly keyed feed_dict for inference.

    Args:
      graph_def: Graph definition of the loaded model.
      serving: Whether the model was exported for Serving.

    Returns:
      The feed_dict suitable for model inference.
    """
    # Creates a dataset based on FLAGS.dataset_name.
    self.dataset = common_flags.create_dataset('test')
    # Create some random images to test inference for any dataset.
    self.images = {
        'img1':
            np.random.uniform(low=64, high=192,
                              size=self.dataset.image_shape).astype('uint8'),
        'img2':
            np.random.uniform(low=32, high=224,
                              size=self.dataset.image_shape).astype('uint8'),
    }
    signature_def = graph_def.signature_def[
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    if serving:
      input_name = signature_def.inputs[
          tf.saved_model.CLASSIFY_INPUTS].name
      # Model for serving takes input: inputs['inputs'] = 'tf_example:0'
      feed_dict = {
          input_name: [
              _create_tf_example_string(self.images['img1']),
              _create_tf_example_string(self.images['img2'])
          ]
      }
    else:
      input_name = signature_def.inputs['images'].name
      # Model for direct use takes input: inputs['images'] = 'original_image:0'
      feed_dict = {
          input_name: np.stack([self.images['img1'], self.images['img2']])
      }
    return feed_dict

  def verify_export_load_and_inference(self, export_for_serving=False):
    """Verify exported model can be loaded and inference can run successfully.

    This function will load the exported model in self.export_dir, then create
    some fake images according to the specification of FLAGS.dataset_name.
    It then feeds the input through the model, and verify the minimal set of
    output signatures are present.
    Note: Model and dataset creation in the underlying library depends on the
          following commandline flags:
            FLAGS.dataset_name
    Args:
      export_for_serving: True if the model was exported for Serving. This
        affects how input is fed into the model.
    """
    tf.compat.v1.reset_default_graph()
    sess = tf.compat.v1.Session()
    graph_def = tf.compat.v1.saved_model.loader.load(
        sess=sess,
        tags=[tf.saved_model.SERVING],
        export_dir=self.export_dir)
    feed_dict = self.create_input_feed(graph_def, export_for_serving)
    results = sess.run(self.minimal_output_signature, feed_dict=feed_dict)

    out_shape = (2,)
    self.assertEqual(np.shape(results['predicted_conf']), out_shape)
    self.assertEqual(np.shape(results['predicted_text']), out_shape)
    self.assertEqual(np.shape(results['predicted_length']), out_shape)
    self.assertEqual(np.shape(results['normalized_seq_conf']), out_shape)
    out_shape = (2, self.dataset.max_sequence_length)
    self.assertEqual(np.shape(results['scores']), out_shape)
    self.assertEqual(np.shape(results['predictions']), out_shape)

  @flagsaver.flagsaver
  def test_fsns_export_for_serving_and_load_inference(self):
    model_export.export_model(self.export_dir, True)
    self.verify_export_load_and_inference(True)

  @flagsaver.flagsaver
  def test_fsns_export_and_load_inference(self):
    model_export.export_model(self.export_dir, False, batch_size=2)
    self.verify_export_load_and_inference(False)


if __name__ == '__main__':
  tf.test.main()