# 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, 2, 214, 16), 'tensorflow.test.main', 'tf.test.main', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((48, 8, 48, 18), 'tensorflow.Graph', 'tf.Graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((61, 15, 61, 25), 'tensorflow.Graph', 'tf.Graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((19, 12, 19, 31), 'tensorflow.identity', 'tf.identity', ({(19, 24, 19, 30): 'inputs'}, {}), '(inputs)', True, 'import tensorflow as tf\n'), ((25, 21, 25, 53), 'tensorflow.identity', 'tf.identity', ({(25, 33, 25, 52): 'preprocessed_inputs'}, {}), '(preprocessed_inputs)', True, 'import tensorflow as tf\n'), ((54, 14, 54, 30), 'tensorflow.train.Saver', 'tf.train.Saver', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((55, 13, 55, 46), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((63, 21, 63, 34), 'tensorflow.GraphDef', 'tf.GraphDef', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((83, 9, 84, 46), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((89, 24, 89, 62), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ({}, {}), '()', False, 'from object_detection.protos import pipeline_pb2\n'), ((91, 6, 95, 52), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', (), '', False, 'from object_detection import exporter\n'), ((98, 9, 99, 46), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((103, 24, 103, 62), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ({}, {}), '()', False, 'from object_detection.protos import pipeline_pb2\n'), ((105, 6, 109, 52), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', (), '', False, 'from object_detection import exporter\n'), ((117, 9, 118, 46), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((120, 24, 120, 62), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ({}, {}), '()', False, 'from object_detection.protos import pipeline_pb2\n'), ((122, 6, 126, 52), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', (), '', False, 'from object_detection import exporter\n'), ((134, 9, 135, 46), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((137, 24, 137, 62), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ({}, {}), '()', False, 'from object_detection.protos import pipeline_pb2\n'), ((139, 6, 143, 52), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', (), '', False, 'from object_detection import exporter\n'), ((151, 9, 152, 46), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((154, 24, 154, 62), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ({}, {}), '()', False, 'from object_detection.protos import pipeline_pb2\n'), ((156, 6, 160, 52), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', (), '', False, 'from object_detection import exporter\n'), ((184, 9, 185, 46), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((187, 24, 187, 62), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ({}, {}), '()', False, 'from object_detection.protos import pipeline_pb2\n'), ((189, 6, 193, 52), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', (), '', False, 'from object_detection import exporter\n'), ((30, 29, 31, 76), 'tensorflow.constant', 'tf.constant', ({(30, 41, 31, 63): '[[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]]', (31, 65, 31, 75): 'tf.float32'}, {}), '([[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]], tf.float32)', True, 'import tensorflow as tf\n'), ((32, 30, 32, 67), 'tensorflow.constant', 'tf.constant', ({(32, 42, 32, 54): '[[0.7, 0.6]]', (32, 56, 32, 66): 'tf.float32'}, {}), '([[0.7, 0.6]], tf.float32)', True, 'import tensorflow as tf\n'), ((33, 31, 33, 64), 'tensorflow.constant', 'tf.constant', ({(33, 43, 33, 51): '[[0, 1]]', (33, 53, 33, 63): 'tf.float32'}, {}), '([[0, 1]], tf.float32)', True, 'import tensorflow as tf\n'), ((34, 28, 34, 56), 'tensorflow.constant', 'tf.constant', ({(34, 40, 34, 43): '[2]', (34, 45, 34, 55): 'tf.float32'}, {}), '([2], tf.float32)', True, 'import tensorflow as tf\n'), ((51, 28, 51, 65), 'tensorflow.constant', 'tf.constant', ({(51, 40, 51, 52): '[1, 3, 4, 3]', (51, 54, 51, 64): 'tf.float32'}, {}), '([1, 3, 4, 3], tf.float32)', True, 'import tensorflow as tf\n'), ((64, 11, 64, 47), 'tensorflow.gfile.GFile', 'tf.gfile.GFile', ({(64, 26, 64, 46): 'inference_graph_path'}, {}), '(inference_graph_path)', True, 'import tensorflow as tf\n'), ((67, 8, 67, 50), 'tensorflow.import_graph_def', 'tf.import_graph_def', (), '', True, 'import tensorflow as tf\n'), ((21, 40, 21, 66), 'tensorflow.constant_initializer', 'tf.constant_initializer', ({(21, 64, 21, 65): '(2)'}, {}), '(2)', True, 'import tensorflow as tf\n'), ((74, 41, 74, 74), 'tensorflow.train.BytesList', 'tf.train.BytesList', (), '', True, 'import tensorflow as tf\n'), ((53, 8, 53, 46), 'tensorflow.train.ExponentialMovingAverage', 'tf.train.ExponentialMovingAverage', ({(53, 42, 53, 45): '(0.0)'}, {}), '(0.0)', True, 'import tensorflow as tf\n'), ((72, 43, 72, 67), 'tensorflow.constant', 'tf.constant', ({(72, 55, 72, 66): 'image_array'}, {}), '(image_array)', True, 'import tensorflow as tf\n'), ((171, 35, 171, 56), 'numpy.ones', 'np.ones', ({(171, 43, 171, 55): '(1, 4, 4, 3)'}, {}), '((1, 4, 4, 3))', True, 'import numpy as np\n'), ((205, 14, 205, 32), 'numpy.ones', 'np.ones', ({(205, 22, 205, 31): '(4, 4, 3)'}, {}), '((4, 4, 3))', True, 'import numpy as np\n')]