# 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()

[((9561, 9575), 'tensorflow.test.main', 'tf.test.main', ([], {}), '()\n', (9573, 9575), True, 'import tensorflow as tf\n'), ((1516, 1526), 'tensorflow.Graph', 'tf.Graph', ([], {}), '()\n', (1524, 1526), True, 'import tensorflow as tf\n'), ((2004, 2014), 'tensorflow.Graph', 'tf.Graph', ([], {}), '()\n', (2012, 2014), True, 'import tensorflow as tf\n'), ((467, 486), 'tensorflow.identity', 'tf.identity', (['inputs'], {}), '(inputs)\n', (478, 486), True, 'import tensorflow as tf\n'), ((715, 747), 'tensorflow.identity', 'tf.identity', (['preprocessed_inputs'], {}), '(preprocessed_inputs)\n', (726, 747), True, 'import tensorflow as tf\n'), ((1762, 1778), 'tensorflow.train.Saver', 'tf.train.Saver', ([], {}), '()\n', (1776, 1778), True, 'import tensorflow as tf\n'), ((1792, 1825), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), '()\n', (1823, 1825), True, 'import tensorflow as tf\n'), ((2068, 2081), 'tensorflow.GraphDef', 'tf.GraphDef', ([], {}), '()\n', (2079, 2081), True, 'import tensorflow as tf\n'), ((2898, 2954), 'mock.patch.object', 'mock.patch.object', (['model_builder', '"""build"""'], {'autospec': '(True)'}), "(model_builder, 'build', autospec=True)\n", (2915, 2954), False, 'import mock\n'), ((3194, 3232), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ([], {}), '()\n', (3230, 3232), False, 'from object_detection.protos import pipeline_pb2\n'), ((3301, 3467), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', ([], {'input_type': '"""image_tensor"""', 'pipeline_config': 'pipeline_config', 'checkpoint_path': 'None', 'inference_graph_path': 'inference_graph_path'}), "(input_type='image_tensor', pipeline_config=\n    pipeline_config, checkpoint_path=None, inference_graph_path=\n    inference_graph_path)\n", (3332, 3467), False, 'from object_detection import exporter\n'), ((3562, 3618), 'mock.patch.object', 'mock.patch.object', (['model_builder', '"""build"""'], {'autospec': '(True)'}), "(model_builder, 'build', autospec=True)\n", (3579, 3618), False, 'import mock\n'), ((3857, 3895), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ([], {}), '()\n', (3893, 3895), False, 'from object_detection.protos import pipeline_pb2\n'), ((3964, 4128), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', ([], {'input_type': '"""tf_example"""', 'pipeline_config': 'pipeline_config', 'checkpoint_path': 'None', 'inference_graph_path': 'inference_graph_path'}), "(input_type='tf_example', pipeline_config=\n    pipeline_config, checkpoint_path=None, inference_graph_path=\n    inference_graph_path)\n", (3995, 4128), False, 'from object_detection import exporter\n'), ((4528, 4584), 'mock.patch.object', 'mock.patch.object', (['model_builder', '"""build"""'], {'autospec': '(True)'}), "(model_builder, 'build', autospec=True)\n", (4545, 4584), False, 'import mock\n'), ((4694, 4732), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ([], {}), '()\n', (4730, 4732), False, 'from object_detection.protos import pipeline_pb2\n'), ((4801, 4978), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', ([], {'input_type': '"""image_tensor"""', 'pipeline_config': 'pipeline_config', 'checkpoint_path': 'checkpoint_path', 'inference_graph_path': 'inference_graph_path'}), "(input_type='image_tensor', pipeline_config=\n    pipeline_config, checkpoint_path=checkpoint_path, inference_graph_path=\n    inference_graph_path)\n", (4832, 4978), False, 'from object_detection import exporter\n'), ((5398, 5454), 'mock.patch.object', 'mock.patch.object', (['model_builder', '"""build"""'], {'autospec': '(True)'}), "(model_builder, 'build', autospec=True)\n", (5415, 5454), False, 'import mock\n'), ((5564, 5602), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ([], {}), '()\n', (5600, 5602), False, 'from object_detection.protos import pipeline_pb2\n'), ((5670, 5847), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', ([], {'input_type': '"""image_tensor"""', 'pipeline_config': 'pipeline_config', 'checkpoint_path': 'checkpoint_path', 'inference_graph_path': 'inference_graph_path'}), "(input_type='image_tensor', pipeline_config=\n    pipeline_config, checkpoint_path=checkpoint_path, inference_graph_path=\n    inference_graph_path)\n", (5701, 5847), False, 'from object_detection import exporter\n'), ((6270, 6326), 'mock.patch.object', 'mock.patch.object', (['model_builder', '"""build"""'], {'autospec': '(True)'}), "(model_builder, 'build', autospec=True)\n", (6287, 6326), False, 'import mock\n'), ((6436, 6474), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ([], {}), '()\n', (6472, 6474), False, 'from object_detection.protos import pipeline_pb2\n'), ((6543, 6720), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', ([], {'input_type': '"""image_tensor"""', 'pipeline_config': 'pipeline_config', 'checkpoint_path': 'checkpoint_path', 'inference_graph_path': 'inference_graph_path'}), "(input_type='image_tensor', pipeline_config=\n    pipeline_config, checkpoint_path=checkpoint_path, inference_graph_path=\n    inference_graph_path)\n", (6574, 6720), False, 'from object_detection import exporter\n'), ((8080, 8136), 'mock.patch.object', 'mock.patch.object', (['model_builder', '"""build"""'], {'autospec': '(True)'}), "(model_builder, 'build', autospec=True)\n", (8097, 8136), False, 'import mock\n'), ((8246, 8284), 'object_detection.protos.pipeline_pb2.TrainEvalPipelineConfig', 'pipeline_pb2.TrainEvalPipelineConfig', ([], {}), '()\n', (8282, 8284), False, 'from object_detection.protos import pipeline_pb2\n'), ((8353, 8528), 'object_detection.exporter.export_inference_graph', 'exporter.export_inference_graph', ([], {'input_type': '"""tf_example"""', 'pipeline_config': 'pipeline_config', 'checkpoint_path': 'checkpoint_path', 'inference_graph_path': 'inference_graph_path'}), "(input_type='tf_example', pipeline_config=\n    pipeline_config, checkpoint_path=checkpoint_path, inference_graph_path=\n    inference_graph_path)\n", (8384, 8528), False, 'from object_detection import exporter\n'), ((896, 965), 'tensorflow.constant', 'tf.constant', (['[[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]]', 'tf.float32'], {}), '([[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]], tf.float32)\n', (907, 965), True, 'import tensorflow as tf\n'), ((1039, 1076), 'tensorflow.constant', 'tf.constant', (['[[0.7, 0.6]]', 'tf.float32'], {}), '([[0.7, 0.6]], tf.float32)\n', (1050, 1076), True, 'import tensorflow as tf\n'), ((1109, 1142), 'tensorflow.constant', 'tf.constant', (['[[0, 1]]', 'tf.float32'], {}), '([[0, 1]], tf.float32)\n', (1120, 1142), True, 'import tensorflow as tf\n'), ((1172, 1200), 'tensorflow.constant', 'tf.constant', (['[2]', 'tf.float32'], {}), '([2], tf.float32)\n', (1183, 1200), True, 'import tensorflow as tf\n'), ((1624, 1661), 'tensorflow.constant', 'tf.constant', (['[1, 3, 4, 3]', 'tf.float32'], {}), '([1, 3, 4, 3], tf.float32)\n', (1635, 1661), True, 'import tensorflow as tf\n'), ((2093, 2129), 'tensorflow.gfile.GFile', 'tf.gfile.GFile', (['inference_graph_path'], {}), '(inference_graph_path)\n', (2107, 2129), True, 'import tensorflow as tf\n'), ((2239, 2281), 'tensorflow.import_graph_def', 'tf.import_graph_def', (['od_graph_def'], {'name': '""""""'}), "(od_graph_def, name='')\n", (2258, 2281), True, 'import tensorflow as tf\n'), ((576, 602), 'tensorflow.constant_initializer', 'tf.constant_initializer', (['(2)'], {}), '(2)\n', (599, 602), True, 'import tensorflow as tf\n'), ((2526, 2559), 'tensorflow.train.BytesList', 'tf.train.BytesList', ([], {'value': '[value]'}), '(value=[value])\n', (2544, 2559), True, 'import tensorflow as tf\n'), ((1701, 1739), 'tensorflow.train.ExponentialMovingAverage', 'tf.train.ExponentialMovingAverage', (['(0.0)'], {}), '(0.0)\n', (1734, 1739), True, 'import tensorflow as tf\n'), ((2421, 2445), 'tensorflow.constant', 'tf.constant', (['image_array'], {}), '(image_array)\n', (2432, 2445), True, 'import tensorflow as tf\n'), ((7397, 7418), 'numpy.ones', 'np.ones', (['(1, 4, 4, 3)'], {}), '((1, 4, 4, 3))\n', (7404, 7418), True, 'import numpy as np\n'), ((9238, 9256), 'numpy.ones', 'np.ones', (['(4, 4, 3)'], {}), '((4, 4, 3))\n', (9245, 9256), True, 'import numpy as np\n')]