# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Travis Yates

"""Tests for object_detection.export_inference_graph."""
import os
import mock
import numpy as np
import tensorflow as tf
from object_detection import exporter
from object_detection.builders import model_builder
from object_detection.core import model
from object_detection.protos import pipeline_pb2


class FakeModel(model.DetectionModel):

  def preprocess(self, inputs):
    return (tf.identity(inputs) *
            tf.get_variable('dummy', shape=(),
                            initializer=tf.constant_initializer(2),
                            dtype=tf.float32))

  def predict(self, preprocessed_inputs):
    return {'image': tf.identity(preprocessed_inputs)}

  def postprocess(self, prediction_dict):
    with tf.control_dependencies(prediction_dict.values()):
      return {
          'detection_boxes': tf.constant([[0.0, 0.0, 0.5, 0.5],
                                          [0.5, 0.5, 0.8, 0.8]], tf.float32),
          'detection_scores': tf.constant([[0.7, 0.6]], tf.float32),
          'detection_classes': tf.constant([[0, 1]], tf.float32),
          'num_detections': tf.constant([2], tf.float32)
      }

  def restore_fn(self, checkpoint_path, from_detection_checkpoint):
    pass

  def loss(self, prediction_dict):
    pass


class ExportInferenceGraphTest(tf.test.TestCase):

  def _save_checkpoint_from_mock_model(self, checkpoint_path,
                                       use_moving_averages):
    g = tf.Graph()
    with g.as_default():
      mock_model = FakeModel(num_classes=1)
      mock_model.preprocess(tf.constant([1, 3, 4, 3], tf.float32))
      if use_moving_averages:
        tf.train.ExponentialMovingAverage(0.0).apply()
      saver = tf.train.Saver()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        saver.save(sess, checkpoint_path)

  def _load_inference_graph(self, inference_graph_path):
    od_graph = tf.Graph()
    with od_graph.as_default():
      od_graph_def = tf.GraphDef()
      with tf.gfile.GFile(inference_graph_path) as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')
    return od_graph

  def _create_tf_example(self, image_array):
    with self.test_session():
      encoded_image = tf.image.encode_jpeg(tf.constant(image_array)).eval()
    def _bytes_feature(value):
      return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
    example = tf.train.Example(features=tf.train.Features(feature={
        'image/encoded': _bytes_feature(encoded_image),
        'image/format': _bytes_feature('jpg'),
        'image/source_id': _bytes_feature('image_id')
    })).SerializeToString()
    return example

  def test_export_graph_with_image_tensor_input(self):
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(num_classes=1)
      inference_graph_path = os.path.join(self.get_temp_dir(),
                                          'exported_graph.pbtxt')

      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          checkpoint_path=None,
          inference_graph_path=inference_graph_path)

  def test_export_graph_with_tf_example_input(self):
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(num_classes=1)
      inference_graph_path = os.path.join(self.get_temp_dir(),
                                          'exported_graph.pbtxt')
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='tf_example',
          pipeline_config=pipeline_config,
          checkpoint_path=None,
          inference_graph_path=inference_graph_path)

  def test_export_frozen_graph(self):
    checkpoint_path = os.path.join(self.get_temp_dir(), 'model-ckpt')
    self._save_checkpoint_from_mock_model(checkpoint_path,
                                          use_moving_averages=False)
    inference_graph_path = os.path.join(self.get_temp_dir(),
                                        'exported_graph.pb')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(num_classes=1)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          checkpoint_path=checkpoint_path,
          inference_graph_path=inference_graph_path)

  def test_export_frozen_graph_with_moving_averages(self):
    checkpoint_path = os.path.join(self.get_temp_dir(), 'model-ckpt')
    self._save_checkpoint_from_mock_model(checkpoint_path,
                                          use_moving_averages=True)
    inference_graph_path = os.path.join(self.get_temp_dir(),
                                        'exported_graph.pb')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(num_classes=1)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = True
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          checkpoint_path=checkpoint_path,
          inference_graph_path=inference_graph_path)

  def test_export_and_run_inference_with_image_tensor(self):
    checkpoint_path = os.path.join(self.get_temp_dir(), 'model-ckpt')
    self._save_checkpoint_from_mock_model(checkpoint_path,
                                          use_moving_averages=False)
    inference_graph_path = os.path.join(self.get_temp_dir(),
                                        'exported_graph.pb')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(num_classes=1)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          checkpoint_path=checkpoint_path,
          inference_graph_path=inference_graph_path)

    inference_graph = self._load_inference_graph(inference_graph_path)
    with self.test_session(graph=inference_graph) as sess:
      image_tensor = inference_graph.get_tensor_by_name('image_tensor:0')
      boxes = inference_graph.get_tensor_by_name('detection_boxes:0')
      scores = inference_graph.get_tensor_by_name('detection_scores:0')
      classes = inference_graph.get_tensor_by_name('detection_classes:0')
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
      (boxes, scores, classes, num_detections) = sess.run(
          [boxes, scores, classes, num_detections],
          feed_dict={image_tensor: np.ones((1, 4, 4, 3)).astype(np.uint8)})
      self.assertAllClose(boxes, [[0.0, 0.0, 0.5, 0.5],
                                  [0.5, 0.5, 0.8, 0.8]])
      self.assertAllClose(scores, [[0.7, 0.6]])
      self.assertAllClose(classes, [[1, 2]])
      self.assertAllClose(num_detections, [2])

  def test_export_and_run_inference_with_tf_example(self):
    checkpoint_path = os.path.join(self.get_temp_dir(), 'model-ckpt')
    self._save_checkpoint_from_mock_model(checkpoint_path,
                                          use_moving_averages=False)
    inference_graph_path = os.path.join(self.get_temp_dir(),
                                        'exported_graph.pb')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(num_classes=1)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='tf_example',
          pipeline_config=pipeline_config,
          checkpoint_path=checkpoint_path,
          inference_graph_path=inference_graph_path)

    inference_graph = self._load_inference_graph(inference_graph_path)
    with self.test_session(graph=inference_graph) as sess:
      tf_example = inference_graph.get_tensor_by_name('tf_example:0')
      boxes = inference_graph.get_tensor_by_name('detection_boxes:0')
      scores = inference_graph.get_tensor_by_name('detection_scores:0')
      classes = inference_graph.get_tensor_by_name('detection_classes:0')
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
      (boxes, scores, classes, num_detections) = sess.run(
          [boxes, scores, classes, num_detections],
          feed_dict={tf_example: self._create_tf_example(
              np.ones((4, 4, 3)).astype(np.uint8))})
      self.assertAllClose(boxes, [[0.0, 0.0, 0.5, 0.5],
                                  [0.5, 0.5, 0.8, 0.8]])
      self.assertAllClose(scores, [[0.7, 0.6]])
      self.assertAllClose(classes, [[1, 2]])
      self.assertAllClose(num_detections, [2])


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

[(214, 'tensorflow.test.main', 'tf.test.main', True, 'import tensorflow as tf\n'), (48, 'tensorflow.Graph', 'tf.Graph', True, 'import tensorflow as tf\n'), (61, 'tensorflow.Graph', 'tf.Graph', True, 'import tensorflow as tf\n'), (19, 'tensorflow.identity', 'tf.identity', True, 'import tensorflow as tf\n'), (25, 'tensorflow.identity', 'tf.identity', True, 'import tensorflow as tf\n'), (54, 'tensorflow.train.Saver', 'tf.train.Saver', True, 'import tensorflow as tf\n'), (55, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', True, 'import tensorflow as tf\n'), (63, 'tensorflow.GraphDef', 'tf.GraphDef', True, 'import tensorflow as tf\n'), (83, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (89, 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', False, 'from object_detection.protos import pipeline_pb2\n'), (91, 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', False, 'from object_detection import exporter\n'), (98, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (103, 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', False, 'from object_detection.protos import pipeline_pb2\n'), (105, 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', False, 'from object_detection import exporter\n'), (117, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (120, 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', False, 'from object_detection.protos import pipeline_pb2\n'), (122, 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', False, 'from object_detection import exporter\n'), (134, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (137, 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', False, 'from object_detection.protos import pipeline_pb2\n'), (139, 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', False, 'from object_detection import exporter\n'), (151, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (154, 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', False, 'from object_detection.protos import pipeline_pb2\n'), (156, 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', False, 'from object_detection import exporter\n'), (184, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (187, 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', False, 'from object_detection.protos import pipeline_pb2\n'), (189, 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', False, 'from object_detection import exporter\n'), (30, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (32, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (33, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (34, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (51, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (64, 'tensorflow.gfile.GFile', 'tf.gfile.GFile', True, 'import tensorflow as tf\n'), (67, 'tensorflow.import_graph_def', 'tf.import_graph_def', True, 'import tensorflow as tf\n'), (21, 'tensorflow.constant_initializer', 'tf.constant_initializer', True, 'import tensorflow as tf\n'), (74, 'tensorflow.train.BytesList', 'tf.train.BytesList', True, 'import tensorflow as tf\n'), (53, 'tensorflow.train.ExponentialMovingAverage', 'tf.train.ExponentialMovingAverage', True, 'import tensorflow as tf\n'), (72, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (171, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (205, 'numpy.ones', 'np.ones', True, 'import numpy as np\n')]