"vscode:/vscode.git/clone" did not exist on "ee53be219bf3f1ef485a725db6b7047b09d0eb63"
exporter_test.py 53 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Copyright 2017 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 object_detection.export_inference_graph."""
import os
import numpy as np
19
import six
20
import tensorflow as tf
21
from google.protobuf import text_format
22
23
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
24
from tensorflow.python.tools import strip_unused_lib
25
from object_detection import exporter
26
from object_detection.builders import graph_rewriter_builder
27
28
from object_detection.builders import model_builder
from object_detection.core import model
29
from object_detection.protos import graph_rewriter_pb2
30
from object_detection.protos import pipeline_pb2
31
from object_detection.utils import ops
pkulzc's avatar
pkulzc committed
32
from object_detection.utils import variables_helper
33

34
35
36
37
38
if six.PY2:
  import mock  # pylint: disable=g-import-not-at-top
else:
  from unittest import mock  # pylint: disable=g-import-not-at-top

Vivek Rathod's avatar
Vivek Rathod committed
39
40
slim = tf.contrib.slim

41
42
43

class FakeModel(model.DetectionModel):

pkulzc's avatar
pkulzc committed
44
45
  def __init__(self, add_detection_keypoints=False, add_detection_masks=False,
               add_detection_features=False):
46
    self._add_detection_keypoints = add_detection_keypoints
47
    self._add_detection_masks = add_detection_masks
pkulzc's avatar
pkulzc committed
48
    self._add_detection_features = add_detection_features
49

50
  def preprocess(self, inputs):
51
52
    true_image_shapes = []  # Doesn't matter for the fake model.
    return tf.identity(inputs), true_image_shapes
53

54
  def predict(self, preprocessed_inputs, true_image_shapes):
55
    return {'image': tf.layers.conv2d(preprocessed_inputs, 3, 1)}
56

57
  def postprocess(self, prediction_dict, true_image_shapes):
58
    with tf.control_dependencies(prediction_dict.values()):
59
      postprocessed_tensors = {
60
61
62
63
64
65
          'detection_boxes': tf.constant([[[0.0, 0.0, 0.5, 0.5],
                                           [0.5, 0.5, 0.8, 0.8]],
                                          [[0.5, 0.5, 1.0, 1.0],
                                           [0.0, 0.0, 0.0, 0.0]]], tf.float32),
          'detection_scores': tf.constant([[0.7, 0.6],
                                           [0.9, 0.0]], tf.float32),
66
67
68
          'detection_multiclass_scores': tf.constant([[[0.3, 0.7], [0.4, 0.6]],
                                                      [[0.1, 0.9], [0.0, 0.0]]],
                                                     tf.float32),
69
70
          'detection_classes': tf.constant([[0, 1],
                                            [1, 0]], tf.float32),
71
72
73
74
75
76
77
78
          'num_detections': tf.constant([2, 1], tf.float32),
          'raw_detection_boxes': tf.constant([[[0.0, 0.0, 0.5, 0.5],
                                               [0.5, 0.5, 0.8, 0.8]],
                                              [[0.5, 0.5, 1.0, 1.0],
                                               [0.0, 0.5, 0.0, 0.5]]],
                                             tf.float32),
          'raw_detection_scores': tf.constant([[0.7, 0.6],
                                               [0.9, 0.5]], tf.float32),
79
      }
80
81
82
      if self._add_detection_keypoints:
        postprocessed_tensors['detection_keypoints'] = tf.constant(
            np.arange(48).reshape([2, 2, 6, 2]), tf.float32)
83
84
      if self._add_detection_masks:
        postprocessed_tensors['detection_masks'] = tf.constant(
85
            np.arange(64).reshape([2, 2, 4, 4]), tf.float32)
pkulzc's avatar
pkulzc committed
86
87
88
89
90
      if self._add_detection_features:
        # let fake detection features have shape [4, 4, 10]
        postprocessed_tensors['detection_features'] = tf.constant(
            np.ones((2, 2, 4, 4, 10)), tf.float32)

91
    return postprocessed_tensors
92

93
  def restore_map(self, checkpoint_path, fine_tune_checkpoint_type):
94
95
    pass

96
  def loss(self, prediction_dict, true_image_shapes):
97
98
    pass

99
100
101
102
103
104
  def regularization_losses(self):
    pass

  def updates(self):
    pass

105
106
107

class ExportInferenceGraphTest(tf.test.TestCase):

108
109
110
111
  def _save_checkpoint_from_mock_model(self,
                                       checkpoint_path,
                                       use_moving_averages,
                                       enable_quantization=False):
112
113
    g = tf.Graph()
    with g.as_default():
114
      mock_model = FakeModel()
115
      preprocessed_inputs, true_image_shapes = mock_model.preprocess(
116
          tf.placeholder(tf.float32, shape=[None, None, None, 3]))
117
118
      predictions = mock_model.predict(preprocessed_inputs, true_image_shapes)
      mock_model.postprocess(predictions, true_image_shapes)
119
120
      if use_moving_averages:
        tf.train.ExponentialMovingAverage(0.0).apply()
121
122
123
124
125
126
127
      tf.train.get_or_create_global_step()
      if enable_quantization:
        graph_rewriter_config = graph_rewriter_pb2.GraphRewriter()
        graph_rewriter_config.quantization.delay = 500000
        graph_rewriter_fn = graph_rewriter_builder.build(
            graph_rewriter_config, is_training=False)
        graph_rewriter_fn()
128
129
130
131
132
133
      saver = tf.train.Saver()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        saver.save(sess, checkpoint_path)

134
  def _load_inference_graph(self, inference_graph_path, is_binary=True):
135
136
137
138
    od_graph = tf.Graph()
    with od_graph.as_default():
      od_graph_def = tf.GraphDef()
      with tf.gfile.GFile(inference_graph_path) as fid:
139
140
141
142
        if is_binary:
          od_graph_def.ParseFromString(fid.read())
        else:
          text_format.Parse(fid.read(), od_graph_def)
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
        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):
159
160
161
162
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
163
164
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
165
      mock_builder.return_value = FakeModel()
166
      output_directory = os.path.join(tmp_dir, 'output')
167
168
169
170
171
      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,
172
173
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
Vivek Rathod's avatar
Vivek Rathod committed
174
175
176
      self.assertTrue(os.path.exists(os.path.join(
          output_directory, 'saved_model', 'saved_model.pb')))

177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
  def test_write_inference_graph(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel()
      output_directory = os.path.join(tmp_dir, 'output')
      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,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory,
          write_inference_graph=True)
      self.assertTrue(os.path.exists(os.path.join(
          output_directory, 'inference_graph.pbtxt')))

Vivek Rathod's avatar
Vivek Rathod committed
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
  def test_export_graph_with_fixed_size_image_tensor_input(self):
    input_shape = [1, 320, 320, 3]

    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(
        trained_checkpoint_prefix, use_moving_averages=False)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel()
      output_directory = os.path.join(tmp_dir, 'output')
      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,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory,
          input_shape=input_shape)
      saved_model_path = os.path.join(output_directory, 'saved_model')
      self.assertTrue(
          os.path.exists(os.path.join(saved_model_path, 'saved_model.pb')))

    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        meta_graph = tf.saved_model.loader.load(
            sess, [tf.saved_model.tag_constants.SERVING], saved_model_path)
        signature = meta_graph.signature_def['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        image_tensor = od_graph.get_tensor_by_name(input_tensor_name)
        self.assertSequenceEqual(image_tensor.get_shape().as_list(),
                                 input_shape)
229
230

  def test_export_graph_with_tf_example_input(self):
231
232
233
234
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
235
236
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
237
      mock_builder.return_value = FakeModel()
238
      output_directory = os.path.join(tmp_dir, 'output')
239
240
241
242
243
      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,
244
245
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
Vivek Rathod's avatar
Vivek Rathod committed
246
247
      self.assertTrue(os.path.exists(os.path.join(
          output_directory, 'saved_model', 'saved_model.pb')))
248

249
  def test_export_graph_with_encoded_image_string_input(self):
250
251
252
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
253
254
255
                                          use_moving_averages=False)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
256
      mock_builder.return_value = FakeModel()
257
      output_directory = os.path.join(tmp_dir, 'output')
258
259
260
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
261
          input_type='encoded_image_string_tensor',
262
          pipeline_config=pipeline_config,
263
264
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
Vivek Rathod's avatar
Vivek Rathod committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
      self.assertTrue(os.path.exists(os.path.join(
          output_directory, 'saved_model', 'saved_model.pb')))

  def _get_variables_in_checkpoint(self, checkpoint_file):
    return set([
        var_name
        for var_name, _ in tf.train.list_variables(checkpoint_file)])

  def test_replace_variable_values_with_moving_averages(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    new_checkpoint_prefix = os.path.join(tmp_dir, 'new.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    graph = tf.Graph()
    with graph.as_default():
      fake_model = FakeModel()
282
      preprocessed_inputs, true_image_shapes = fake_model.preprocess(
Vivek Rathod's avatar
Vivek Rathod committed
283
          tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3]))
284
285
      predictions = fake_model.predict(preprocessed_inputs, true_image_shapes)
      fake_model.postprocess(predictions, true_image_shapes)
Vivek Rathod's avatar
Vivek Rathod committed
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
      exporter.replace_variable_values_with_moving_averages(
          graph, trained_checkpoint_prefix, new_checkpoint_prefix)

    expected_variables = set(['conv2d/bias', 'conv2d/kernel'])
    variables_in_old_ckpt = self._get_variables_in_checkpoint(
        trained_checkpoint_prefix)
    self.assertIn('conv2d/bias/ExponentialMovingAverage',
                  variables_in_old_ckpt)
    self.assertIn('conv2d/kernel/ExponentialMovingAverage',
                  variables_in_old_ckpt)
    variables_in_new_ckpt = self._get_variables_in_checkpoint(
        new_checkpoint_prefix)
    self.assertTrue(expected_variables.issubset(variables_in_new_ckpt))
    self.assertNotIn('conv2d/bias/ExponentialMovingAverage',
                     variables_in_new_ckpt)
    self.assertNotIn('conv2d/kernel/ExponentialMovingAverage',
                     variables_in_new_ckpt)
303

304
305
306
307
  def test_export_graph_with_moving_averages(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
308
                                          use_moving_averages=True)
309
    output_directory = os.path.join(tmp_dir, 'output')
310
311
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
312
      mock_builder.return_value = FakeModel()
313
314
315
316
317
      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,
318
319
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
Vivek Rathod's avatar
Vivek Rathod committed
320
321
322
323
324
325
      self.assertTrue(os.path.exists(os.path.join(
          output_directory, 'saved_model', 'saved_model.pb')))
    expected_variables = set(['conv2d/bias', 'conv2d/kernel', 'global_step'])
    actual_variables = set(
        [var_name for var_name, _ in tf.train.list_variables(output_directory)])
    self.assertTrue(expected_variables.issubset(actual_variables))
326

327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
  def test_export_model_with_quantization_nodes(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(
        trained_checkpoint_prefix,
        use_moving_averages=False,
        enable_quantization=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'inference_graph.pbtxt')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel()
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      text_format.Merge(
          """graph_rewriter {
               quantization {
                 delay: 50000
                 activation_bits: 8
                 weight_bits: 8
               }
             }""", pipeline_config)
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory,
          write_inference_graph=True)
    self._load_inference_graph(inference_graph_path, is_binary=False)
    has_quant_nodes = False
pkulzc's avatar
pkulzc committed
357
    for v in variables_helper.get_global_variables_safely():
358
359
360
361
362
      if v.op.name.endswith('act_quant/min'):
        has_quant_nodes = True
        break
    self.assertTrue(has_quant_nodes)

363
  def test_export_model_with_all_output_nodes(self):
364
365
366
367
368
369
370
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
371
372
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
373
      mock_builder.return_value = FakeModel(
pkulzc's avatar
pkulzc committed
374
375
          add_detection_keypoints=True, add_detection_masks=True,
          add_detection_features=True)
376
377
378
379
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
380
381
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
382
383
384
385
386
    inference_graph = self._load_inference_graph(inference_graph_path)
    with self.test_session(graph=inference_graph):
      inference_graph.get_tensor_by_name('image_tensor:0')
      inference_graph.get_tensor_by_name('detection_boxes:0')
      inference_graph.get_tensor_by_name('detection_scores:0')
387
      inference_graph.get_tensor_by_name('detection_multiclass_scores:0')
388
      inference_graph.get_tensor_by_name('detection_classes:0')
389
      inference_graph.get_tensor_by_name('detection_keypoints:0')
390
391
      inference_graph.get_tensor_by_name('detection_masks:0')
      inference_graph.get_tensor_by_name('num_detections:0')
pkulzc's avatar
pkulzc committed
392
      inference_graph.get_tensor_by_name('detection_features:0')
393
394

  def test_export_model_with_detection_only_nodes(self):
395
396
397
398
399
400
401
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
402
403
404
405
406
407
408
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(add_detection_masks=False)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
409
410
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
411
412
413
414
415
    inference_graph = self._load_inference_graph(inference_graph_path)
    with self.test_session(graph=inference_graph):
      inference_graph.get_tensor_by_name('image_tensor:0')
      inference_graph.get_tensor_by_name('detection_boxes:0')
      inference_graph.get_tensor_by_name('detection_scores:0')
416
      inference_graph.get_tensor_by_name('detection_multiclass_scores:0')
417
418
419
      inference_graph.get_tensor_by_name('detection_classes:0')
      inference_graph.get_tensor_by_name('num_detections:0')
      with self.assertRaises(KeyError):
420
        inference_graph.get_tensor_by_name('detection_keypoints:0')
421
422
        inference_graph.get_tensor_by_name('detection_masks:0')

pkulzc's avatar
pkulzc committed
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
  def test_export_model_with_detection_only_nodes_and_detection_features(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(add_detection_features=True)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
    inference_graph = self._load_inference_graph(inference_graph_path)
    with self.test_session(graph=inference_graph):
      inference_graph.get_tensor_by_name('image_tensor:0')
      inference_graph.get_tensor_by_name('detection_boxes:0')
      inference_graph.get_tensor_by_name('detection_scores:0')
      inference_graph.get_tensor_by_name('detection_multiclass_scores:0')
      inference_graph.get_tensor_by_name('detection_classes:0')
      inference_graph.get_tensor_by_name('num_detections:0')
      inference_graph.get_tensor_by_name('detection_features:0')
      with self.assertRaises(KeyError):
        inference_graph.get_tensor_by_name('detection_keypoints:0')
        inference_graph.get_tensor_by_name('detection_masks:0')

453
  def test_export_and_run_inference_with_image_tensor(self):
454
455
456
457
458
459
460
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
461
462
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
463
464
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
465
466
467
468
469
      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,
470
471
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
472
473
474
475
476
477
478

    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')
479
      keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0')
480
      masks = inference_graph.get_tensor_by_name('detection_masks:0')
481
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
482
483
484
485
      (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
       num_detections_np) = sess.run(
           [boxes, scores, classes, keypoints, masks, num_detections],
           feed_dict={image_tensor: np.ones((2, 4, 4, 3)).astype(np.uint8)})
486
487
488
489
490
491
492
493
      self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                      [0.5, 0.5, 0.8, 0.8]],
                                     [[0.5, 0.5, 1.0, 1.0],
                                      [0.0, 0.0, 0.0, 0.0]]])
      self.assertAllClose(scores_np, [[0.7, 0.6],
                                      [0.9, 0.0]])
      self.assertAllClose(classes_np, [[1, 2],
                                       [2, 1]])
494
      self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
495
496
      self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
      self.assertAllClose(num_detections_np, [2, 1])
497

498
499
500
501
502
503
504
505
506
507
508
509
510
  def _create_encoded_image_string(self, image_array_np, encoding_format):
    od_graph = tf.Graph()
    with od_graph.as_default():
      if encoding_format == 'jpg':
        encoded_string = tf.image.encode_jpeg(image_array_np)
      elif encoding_format == 'png':
        encoded_string = tf.image.encode_png(image_array_np)
      else:
        raise ValueError('Supports only the following formats: `jpg`, `png`')
    with self.test_session(graph=od_graph):
      return encoded_string.eval()

  def test_export_and_run_inference_with_encoded_image_string_tensor(self):
511
512
513
514
515
516
517
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
518
519
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
520
521
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
522
523
524
525
526
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='encoded_image_string_tensor',
          pipeline_config=pipeline_config,
527
528
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
529
530
531
532
533
534
535
536
537
538
539

    inference_graph = self._load_inference_graph(inference_graph_path)
    jpg_image_str = self._create_encoded_image_string(
        np.ones((4, 4, 3)).astype(np.uint8), 'jpg')
    png_image_str = self._create_encoded_image_string(
        np.ones((4, 4, 3)).astype(np.uint8), 'png')
    with self.test_session(graph=inference_graph) as sess:
      image_str_tensor = inference_graph.get_tensor_by_name(
          'encoded_image_string_tensor:0')
      boxes = inference_graph.get_tensor_by_name('detection_boxes:0')
      scores = inference_graph.get_tensor_by_name('detection_scores:0')
540
541
      multiclass_scores = inference_graph.get_tensor_by_name(
          'detection_multiclass_scores:0')
542
      classes = inference_graph.get_tensor_by_name('detection_classes:0')
543
      keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0')
544
545
546
      masks = inference_graph.get_tensor_by_name('detection_masks:0')
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
      for image_str in [jpg_image_str, png_image_str]:
547
        image_str_batch_np = np.hstack([image_str]* 2)
548
549
550
551
552
553
        (boxes_np, scores_np, multiclass_scores_np, classes_np, keypoints_np,
         masks_np, num_detections_np) = sess.run(
             [
                 boxes, scores, multiclass_scores, classes, keypoints, masks,
                 num_detections
             ],
554
555
556
557
558
559
560
             feed_dict={image_str_tensor: image_str_batch_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
561
562
        self.assertAllClose(multiclass_scores_np, [[[0.3, 0.7], [0.4, 0.6]],
                                                   [[0.1, 0.9], [0.0, 0.0]]])
563
564
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
565
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
566
567
568
569
570
571
572
573
574
575
576
577
578
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])

  def test_raise_runtime_error_on_images_with_different_sizes(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
579
580
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='encoded_image_string_tensor',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)

    inference_graph = self._load_inference_graph(inference_graph_path)
    large_image = self._create_encoded_image_string(
        np.ones((4, 4, 3)).astype(np.uint8), 'jpg')
    small_image = self._create_encoded_image_string(
        np.ones((2, 2, 3)).astype(np.uint8), 'jpg')

    image_str_batch_np = np.hstack([large_image, small_image])
    with self.test_session(graph=inference_graph) as sess:
      image_str_tensor = inference_graph.get_tensor_by_name(
          'encoded_image_string_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')
602
      keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0')
603
604
605
      masks = inference_graph.get_tensor_by_name('detection_masks:0')
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
606
                                   'TensorArray.*shape'):
607
608
609
        sess.run(
            [boxes, scores, classes, keypoints, masks, num_detections],
            feed_dict={image_str_tensor: image_str_batch_np})
610

611
  def test_export_and_run_inference_with_tf_example(self):
612
613
614
615
616
617
618
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
619
620
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
621
622
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
623
624
625
626
627
      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,
628
629
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
630
631

    inference_graph = self._load_inference_graph(inference_graph_path)
632
633
    tf_example_np = np.expand_dims(self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8)), axis=0)
634
635
636
637
638
    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')
639
      keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0')
640
      masks = inference_graph.get_tensor_by_name('detection_masks:0')
641
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
642
643
644
645
      (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
       num_detections_np) = sess.run(
           [boxes, scores, classes, keypoints, masks, num_detections],
           feed_dict={tf_example: tf_example_np})
646
647
648
649
650
651
652
653
      self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                      [0.5, 0.5, 0.8, 0.8]],
                                     [[0.5, 0.5, 1.0, 1.0],
                                      [0.0, 0.0, 0.0, 0.0]]])
      self.assertAllClose(scores_np, [[0.7, 0.6],
                                      [0.9, 0.0]])
      self.assertAllClose(classes_np, [[1, 2],
                                       [2, 1]])
654
      self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
655
656
      self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
      self.assertAllClose(num_detections_np, [2, 1])
657

658
659
660
661
662
663
664
665
666
667
668
  def test_write_frozen_graph(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    inference_graph_path = os.path.join(output_directory,
                                        'frozen_inference_graph.pb')
    tf.gfile.MakeDirs(output_directory)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
669
670
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
671
672
673
674
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      detection_model = model_builder.build(pipeline_config.model,
                                            is_training=False)
675
      outputs, _ = exporter.build_detection_graph(
676
677
678
679
680
681
682
683
          input_type='tf_example',
          detection_model=detection_model,
          input_shape=None,
          output_collection_name='inference_op',
          graph_hook_fn=None)
      output_node_names = ','.join(outputs.keys())
      saver = tf.train.Saver()
      input_saver_def = saver.as_saver_def()
684
      exporter.freeze_graph_with_def_protos(
685
686
687
688
689
690
          input_graph_def=tf.get_default_graph().as_graph_def(),
          input_saver_def=input_saver_def,
          input_checkpoint=trained_checkpoint_prefix,
          output_node_names=output_node_names,
          restore_op_name='save/restore_all',
          filename_tensor_name='save/Const:0',
691
          output_graph=inference_graph_path,
692
693
694
695
696
697
698
699
700
701
702
          clear_devices=True,
          initializer_nodes='')

    inference_graph = self._load_inference_graph(inference_graph_path)
    tf_example_np = np.expand_dims(self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8)), axis=0)
    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')
703
      keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0')
704
705
      masks = inference_graph.get_tensor_by_name('detection_masks:0')
      num_detections = inference_graph.get_tensor_by_name('num_detections:0')
706
707
708
709
      (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
       num_detections_np) = sess.run(
           [boxes, scores, classes, keypoints, masks, num_detections],
           feed_dict={tf_example: tf_example_np})
710
711
712
713
714
715
716
717
      self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                      [0.5, 0.5, 0.8, 0.8]],
                                     [[0.5, 0.5, 1.0, 1.0],
                                      [0.0, 0.0, 0.0, 0.0]]])
      self.assertAllClose(scores_np, [[0.7, 0.6],
                                      [0.9, 0.0]])
      self.assertAllClose(classes_np, [[1, 2],
                                       [2, 1]])
718
      self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
719
720
721
      self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
      self.assertAllClose(num_detections_np, [2, 1])

722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
  def test_export_graph_saves_pipeline_file(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=True)
    output_directory = os.path.join(tmp_dir, 'output')
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel()
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
      expected_pipeline_path = os.path.join(
          output_directory, 'pipeline.config')
      self.assertTrue(os.path.exists(expected_pipeline_path))

      written_pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      with tf.gfile.GFile(expected_pipeline_path, 'r') as f:
        proto_str = f.read()
        text_format.Merge(proto_str, written_pipeline_config)
        self.assertProtoEquals(pipeline_config, written_pipeline_config)

747
  def test_export_saved_model_and_run_inference(self):
748
749
750
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
751
                                          use_moving_averages=False)
752
753
    output_directory = os.path.join(tmp_dir, 'output')
    saved_model_path = os.path.join(output_directory, 'saved_model')
754
755
756

    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
757
758
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
759
760
761
762
763
      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,
764
765
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)
766

767
768
    tf_example_np = np.hstack([self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8))] * 2)
769
770
    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
Vivek Rathod's avatar
Vivek Rathod committed
771
        meta_graph = tf.saved_model.loader.load(
772
            sess, [tf.saved_model.tag_constants.SERVING], saved_model_path)
Vivek Rathod's avatar
Vivek Rathod committed
773
774
775
776
777
778
779
780
781

        signature = meta_graph.signature_def['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        tf_example = od_graph.get_tensor_by_name(input_tensor_name)

        boxes = od_graph.get_tensor_by_name(
            signature.outputs['detection_boxes'].name)
        scores = od_graph.get_tensor_by_name(
            signature.outputs['detection_scores'].name)
pkulzc's avatar
pkulzc committed
782
783
        multiclass_scores = od_graph.get_tensor_by_name(
            signature.outputs['detection_multiclass_scores'].name)
Vivek Rathod's avatar
Vivek Rathod committed
784
785
        classes = od_graph.get_tensor_by_name(
            signature.outputs['detection_classes'].name)
786
787
        keypoints = od_graph.get_tensor_by_name(
            signature.outputs['detection_keypoints'].name)
Vivek Rathod's avatar
Vivek Rathod committed
788
789
790
791
792
        masks = od_graph.get_tensor_by_name(
            signature.outputs['detection_masks'].name)
        num_detections = od_graph.get_tensor_by_name(
            signature.outputs['num_detections'].name)

pkulzc's avatar
pkulzc committed
793
794
795
796
        (boxes_np, scores_np, multiclass_scores_np, classes_np, keypoints_np,
         masks_np, num_detections_np) = sess.run(
             [boxes, scores, multiclass_scores, classes, keypoints, masks,
              num_detections],
797
798
799
800
801
802
803
             feed_dict={tf_example: tf_example_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
pkulzc's avatar
pkulzc committed
804
805
        self.assertAllClose(multiclass_scores_np, [[[0.3, 0.7], [0.4, 0.6]],
                                                   [[0.1, 0.9], [0.0, 0.0]]])
806
807
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
808
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
809
810
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])
811

812
813
814
815
816
817
818
819
820
821
  def test_write_saved_model(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
    output_directory = os.path.join(tmp_dir, 'output')
    saved_model_path = os.path.join(output_directory, 'saved_model')
    tf.gfile.MakeDirs(output_directory)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
822
823
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
824
825
826
827
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      detection_model = model_builder.build(pipeline_config.model,
                                            is_training=False)
828
      outputs, placeholder_tensor = exporter.build_detection_graph(
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
          input_type='tf_example',
          detection_model=detection_model,
          input_shape=None,
          output_collection_name='inference_op',
          graph_hook_fn=None)
      output_node_names = ','.join(outputs.keys())
      saver = tf.train.Saver()
      input_saver_def = saver.as_saver_def()
      frozen_graph_def = exporter.freeze_graph_with_def_protos(
          input_graph_def=tf.get_default_graph().as_graph_def(),
          input_saver_def=input_saver_def,
          input_checkpoint=trained_checkpoint_prefix,
          output_node_names=output_node_names,
          restore_op_name='save/restore_all',
          filename_tensor_name='save/Const:0',
844
          output_graph='',
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
          clear_devices=True,
          initializer_nodes='')
      exporter.write_saved_model(
          saved_model_path=saved_model_path,
          frozen_graph_def=frozen_graph_def,
          inputs=placeholder_tensor,
          outputs=outputs)

    tf_example_np = np.hstack([self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8))] * 2)
    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        meta_graph = tf.saved_model.loader.load(
            sess, [tf.saved_model.tag_constants.SERVING], saved_model_path)

        signature = meta_graph.signature_def['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        tf_example = od_graph.get_tensor_by_name(input_tensor_name)

        boxes = od_graph.get_tensor_by_name(
            signature.outputs['detection_boxes'].name)
        scores = od_graph.get_tensor_by_name(
            signature.outputs['detection_scores'].name)
        classes = od_graph.get_tensor_by_name(
            signature.outputs['detection_classes'].name)
870
871
        keypoints = od_graph.get_tensor_by_name(
            signature.outputs['detection_keypoints'].name)
872
873
874
875
876
        masks = od_graph.get_tensor_by_name(
            signature.outputs['detection_masks'].name)
        num_detections = od_graph.get_tensor_by_name(
            signature.outputs['num_detections'].name)

877
        (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
878
         num_detections_np) = sess.run(
879
             [boxes, scores, classes, keypoints, masks, num_detections],
880
881
882
883
884
885
886
887
888
             feed_dict={tf_example: tf_example_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
889
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
890
891
892
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])

893
894
895
896
897
898
899
900
901
902
903
  def test_export_checkpoint_and_run_inference(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
    output_directory = os.path.join(tmp_dir, 'output')
    model_path = os.path.join(output_directory, 'model.ckpt')
    meta_graph_path = model_path + '.meta'

    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
904
905
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
906
907
908
909
910
911
912
913
      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,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)

914
915
    tf_example_np = np.hstack([self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8))] * 2)
916
917
918
919
920
921
922
923
924
    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        new_saver = tf.train.import_meta_graph(meta_graph_path)
        new_saver.restore(sess, model_path)

        tf_example = od_graph.get_tensor_by_name('tf_example:0')
        boxes = od_graph.get_tensor_by_name('detection_boxes:0')
        scores = od_graph.get_tensor_by_name('detection_scores:0')
        classes = od_graph.get_tensor_by_name('detection_classes:0')
925
        keypoints = od_graph.get_tensor_by_name('detection_keypoints:0')
926
927
        masks = od_graph.get_tensor_by_name('detection_masks:0')
        num_detections = od_graph.get_tensor_by_name('num_detections:0')
928
        (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
929
         num_detections_np) = sess.run(
930
             [boxes, scores, classes, keypoints, masks, num_detections],
931
932
933
934
935
936
937
938
939
             feed_dict={tf_example: tf_example_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
940
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
941
942
943
944
945
946
947
948
949
950
951
952
953
954
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])

  def test_write_graph_and_checkpoint(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
    output_directory = os.path.join(tmp_dir, 'output')
    model_path = os.path.join(output_directory, 'model.ckpt')
    meta_graph_path = model_path + '.meta'
    tf.gfile.MakeDirs(output_directory)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
955
956
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
957
958
959
960
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      detection_model = model_builder.build(pipeline_config.model,
                                            is_training=False)
961
      exporter.build_detection_graph(
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
          input_type='tf_example',
          detection_model=detection_model,
          input_shape=None,
          output_collection_name='inference_op',
          graph_hook_fn=None)
      saver = tf.train.Saver()
      input_saver_def = saver.as_saver_def()
      exporter.write_graph_and_checkpoint(
          inference_graph_def=tf.get_default_graph().as_graph_def(),
          model_path=model_path,
          input_saver_def=input_saver_def,
          trained_checkpoint_prefix=trained_checkpoint_prefix)

    tf_example_np = np.hstack([self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8))] * 2)
    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        new_saver = tf.train.import_meta_graph(meta_graph_path)
        new_saver.restore(sess, model_path)

        tf_example = od_graph.get_tensor_by_name('tf_example:0')
        boxes = od_graph.get_tensor_by_name('detection_boxes:0')
        scores = od_graph.get_tensor_by_name('detection_scores:0')
985
986
        raw_boxes = od_graph.get_tensor_by_name('raw_detection_boxes:0')
        raw_scores = od_graph.get_tensor_by_name('raw_detection_scores:0')
987
        classes = od_graph.get_tensor_by_name('detection_classes:0')
988
        keypoints = od_graph.get_tensor_by_name('detection_keypoints:0')
989
990
        masks = od_graph.get_tensor_by_name('detection_masks:0')
        num_detections = od_graph.get_tensor_by_name('num_detections:0')
991
992
993
994
        (boxes_np, scores_np, raw_boxes_np, raw_scores_np, classes_np,
         keypoints_np, masks_np, num_detections_np) = sess.run(
             [boxes, scores, raw_boxes, raw_scores, classes, keypoints, masks,
              num_detections],
995
996
997
998
999
1000
1001
             feed_dict={tf_example: tf_example_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
1002
1003
1004
1005
1006
1007
        self.assertAllClose(raw_boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                            [0.5, 0.5, 0.8, 0.8]],
                                           [[0.5, 0.5, 1.0, 1.0],
                                            [0.0, 0.5, 0.0, 0.5]]])
        self.assertAllClose(raw_scores_np, [[0.7, 0.6],
                                            [0.9, 0.5]])
1008
1009
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
1010
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
1011
1012
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])
1013

1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
  def test_rewrite_nn_resize_op(self):
    g = tf.Graph()
    with g.as_default():
      x = array_ops.placeholder(dtypes.float32, shape=(8, 10, 10, 8))
      y = array_ops.placeholder(dtypes.float32, shape=(8, 20, 20, 8))
      s = ops.nearest_neighbor_upsampling(x, 2)
      t = s + y
      exporter.rewrite_nn_resize_op()

    resize_op_found = False
    for op in g.get_operations():
      if op.type == 'ResizeNearestNeighbor':
        resize_op_found = True
        self.assertEqual(op.inputs[0], x)
        self.assertEqual(op.outputs[0].consumers()[0], t.op)
        break

    self.assertTrue(resize_op_found)

  def test_rewrite_nn_resize_op_quantized(self):
    g = tf.Graph()
    with g.as_default():
      x = array_ops.placeholder(dtypes.float32, shape=(8, 10, 10, 8))
      x_conv = tf.contrib.slim.conv2d(x, 8, 1)
      y = array_ops.placeholder(dtypes.float32, shape=(8, 20, 20, 8))
      s = ops.nearest_neighbor_upsampling(x_conv, 2)
      t = s + y

      graph_rewriter_config = graph_rewriter_pb2.GraphRewriter()
      graph_rewriter_config.quantization.delay = 500000
      graph_rewriter_fn = graph_rewriter_builder.build(
          graph_rewriter_config, is_training=False)
      graph_rewriter_fn()

      exporter.rewrite_nn_resize_op(is_quantized=True)

    resize_op_found = False
    for op in g.get_operations():
      if op.type == 'ResizeNearestNeighbor':
        resize_op_found = True
        self.assertEqual(op.inputs[0].op.type, 'FakeQuantWithMinMaxVars')
        self.assertEqual(op.outputs[0].consumers()[0], t.op)
        break

    self.assertTrue(resize_op_found)

1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
  def test_rewrite_nn_resize_op_multiple_path(self):
    g = tf.Graph()
    with g.as_default():
      with tf.name_scope('nearest_upsampling'):
        x = array_ops.placeholder(dtypes.float32, shape=(8, 10, 10, 8))
        x_stack = tf.stack([tf.stack([x] * 2, axis=3)] * 2, axis=2)
        x_reshape = tf.reshape(x_stack, [8, 20, 20, 8])

      with tf.name_scope('nearest_upsampling'):
        x_2 = array_ops.placeholder(dtypes.float32, shape=(8, 10, 10, 8))
        x_stack_2 = tf.stack([tf.stack([x_2] * 2, axis=3)] * 2, axis=2)
        x_reshape_2 = tf.reshape(x_stack_2, [8, 20, 20, 8])

      t = x_reshape + x_reshape_2

      exporter.rewrite_nn_resize_op()

    graph_def = g.as_graph_def()
    graph_def = strip_unused_lib.strip_unused(
        graph_def,
        input_node_names=[
            'nearest_upsampling/Placeholder', 'nearest_upsampling_1/Placeholder'
        ],
        output_node_names=['add'],
        placeholder_type_enum=dtypes.float32.as_datatype_enum)

    counter_resize_op = 0
    t_input_ops = [op.name for op in t.op.inputs]
    for node in graph_def.node:
      # Make sure Stacks are replaced.
      self.assertNotEqual(node.op, 'Pack')
      if node.op == 'ResizeNearestNeighbor':
        counter_resize_op += 1
        self.assertIn(node.name + ':0', t_input_ops)
    self.assertEqual(counter_resize_op, 2)

1096

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