Merge branch 'master' of github.com:tensorflow/models

lfads/synth_data/synthetic_data_utils.py

@@ -132,11 +132,10 @@ def spikify_data(data_e, rng, dt=1.0, max_firing_rate=100):
     dt: how often the data are sampled
     max_firing_rate: the firing rate that is associated with a value of 1.0
   Returns:
-    spikified_data_e: a list of length b of the data represented as spikes,
+    spikified_e: a list of length b of the data represented as spikes,
     sampled from the underlying poisson process.
     """
 
-  spikifies_data_e = []
   E = len(data_e)
   spikes_e = []
   for e in range(E):
@@ -152,6 +151,31 @@ def spikify_data(data_e, rng, dt=1.0, max_firing_rate=100):
   return spikes_e
 
 
+def gaussify_data(data_e, rng, dt=1.0, max_firing_rate=100):
+  """ Apply gaussian noise to a continuous dataset whose values are between
+  0.0 and 1.0
+
+  Args:
+    data_e: nexamples length list of NxT trials
+    dt: how often the data are sampled
+    max_firing_rate: the firing rate that is associated with a value of 1.0
+  Returns:
+    gauss_e: a list of length b of the data with noise.
+    """
+
+  E = len(data_e)
+  mfr = max_firing_rate
+  gauss_e = []
+  for e in range(E):
+    data = data_e[e]
+    N,T = data.shape
+    noisy_data = data * mfr + np.random.randn(N,T) * (5.0*mfr) * np.sqrt(dt)
+    gauss_e.append(noisy_data)
+
+  return gauss_e
+
+
+
 def get_train_n_valid_inds(num_trials, train_fraction, nspikifications):
   """Split the numbers between 0 and num_trials-1 into two portions for
   training and validation, based on the train fraction.
@@ -295,6 +319,8 @@ def add_alignment_projections(datasets, npcs, ntime=None, nsamples=None):
     W_chxp, _, _, _ = \
         np.linalg.lstsq(all_data_zm_chxtc.T, all_data_pca_pxtc.T)
     dataset['alignment_matrix_cxf'] = W_chxp
+    alignment_bias_cx1 = all_data_mean_nx1[cidx_s:cidx_f]
+    dataset['alignment_bias_c'] = np.squeeze(alignment_bias_cx1, axis=1)
 
   do_debug_plot = False
   if do_debug_plot:

lfads/utils.py

@@ -82,9 +82,9 @@ def linear(x, out_size, do_bias=True, alpha=1.0, identity_if_possible=False,
     return tf.matmul(x, W)
 
 
-def init_linear(in_size, out_size, do_bias=True, mat_init_value=None, alpha=1.0,
-                identity_if_possible=False, normalized=False,
-                name=None, collections=None):
+def init_linear(in_size, out_size, do_bias=True, mat_init_value=None,
+                bias_init_value=None, alpha=1.0, identity_if_possible=False,
+                normalized=False, name=None, collections=None):
   """Linear (affine) transformation, y = x W + b, for a variety of
   configurations.
 
@@ -110,6 +110,9 @@ def init_linear(in_size, out_size, do_bias=True, mat_init_value=None, alpha=1.0,
   if mat_init_value is not None and mat_init_value.shape != (in_size, out_size):
     raise ValueError(
         'Provided mat_init_value must have shape [%d, %d].'%(in_size, out_size))
+  if bias_init_value is not None and bias_init_value.shape != (1,out_size):
+    raise ValueError(
+        'Provided bias_init_value must have shape [1,%d].'%(out_size,))
 
   if mat_init_value is None:
     stddev = alpha/np.sqrt(float(in_size))
@@ -143,16 +146,20 @@ def init_linear(in_size, out_size, do_bias=True, mat_init_value=None, alpha=1.0,
       w = tf.get_variable(wname, [in_size, out_size], initializer=mat_init,
                           collections=w_collections)
 
+  b = None
   if do_bias:
     b_collections = [tf.GraphKeys.GLOBAL_VARIABLES]
     if collections:
       b_collections += collections
     bname = (name + "/b") if name else "/b"
-    b = tf.get_variable(bname, [1, out_size],
-                        initializer=tf.zeros_initializer(),
-                        collections=b_collections)
-  else:
-    b = None
+    if bias_init_value is None:
+      b = tf.get_variable(bname, [1, out_size],
+                          initializer=tf.zeros_initializer(),
+                          collections=b_collections)
+    else:
+      b = tf.Variable(bias_init_value, name=bname,
+                      collections=b_collections)
+
 
   return (w, b)
 

object_detection/README.md

@@ -54,6 +54,17 @@ Extras:
       Exporting a trained model for inference</a><br>
 * <a href='g3doc/defining_your_own_model.md'>
       Defining your own model architecture</a><br>
+* <a href='g3doc/using_your_own_dataset.md'>
+      Bringing in your own dataset</a><br>
+
+## Getting Help
+
+Please report bugs to the tensorflow/models/ Github
+[issue tracker](https://github.com/tensorflow/models/issues), prefixing the
+issue name with "object_detection". To get help with issues you may encounter
+using the Tensorflow Object Detection API, create a new question on
+[StackOverflow](https://stackoverflow.com/) with the tags "tensorflow" and
+"object-detection".
 
 ## Release information
 

object_detection/core/BUILD

@@ -270,6 +270,7 @@ py_library(
     deps = [
         "//tensorflow",
         "//tensorflow_models/object_detection/utils:ops",
+        "//tensorflow_models/object_detection/utils:shape_utils",
         "//tensorflow_models/object_detection/utils:static_shape",
     ],
 )

object_detection/core/box_predictor.py

@@ -29,6 +29,7 @@ few box predictor architectures are shared across many models.
 from abc import abstractmethod
 import tensorflow as tf
 from object_detection.utils import ops
+from object_detection.utils import shape_utils
 from object_detection.utils import static_shape
 
 slim = tf.contrib.slim
@@ -316,6 +317,8 @@ class MaskRCNNBoxPredictor(BoxPredictor):
     self._predict_instance_masks = predict_instance_masks
     self._mask_prediction_conv_depth = mask_prediction_conv_depth
     self._predict_keypoints = predict_keypoints
+    if self._predict_instance_masks:
+      raise ValueError('Mask prediction is unimplemented.')
     if self._predict_keypoints:
       raise ValueError('Keypoint prediction is unimplemented.')
     if ((self._predict_instance_masks or self._predict_keypoints) and
@@ -524,23 +527,21 @@ class ConvolutionalBoxPredictor(BoxPredictor):
           class_predictions_with_background = tf.sigmoid(
               class_predictions_with_background)
 
-    batch_size = static_shape.get_batch_size(image_features.get_shape())
-    if batch_size is None:
-      features_height = static_shape.get_height(image_features.get_shape())
-      features_width = static_shape.get_width(image_features.get_shape())
-      flattened_predictions_size = (features_height * features_width *
-                                    num_predictions_per_location)
-      box_encodings = tf.reshape(
-          box_encodings,
-          [-1, flattened_predictions_size, 1, self._box_code_size])
-      class_predictions_with_background = tf.reshape(
-          class_predictions_with_background,
-          [-1, flattened_predictions_size, num_class_slots])
-    else:
-      box_encodings = tf.reshape(
-          box_encodings, [batch_size, -1, 1, self._box_code_size])
-      class_predictions_with_background = tf.reshape(
-          class_predictions_with_background, [batch_size, -1, num_class_slots])
+    combined_feature_map_shape = shape_utils.combined_static_and_dynamic_shape(
+        image_features)
+    box_encodings = tf.reshape(
+        box_encodings, tf.stack([combined_feature_map_shape[0],
+                                 combined_feature_map_shape[1] *
+                                 combined_feature_map_shape[2] *
+                                 num_predictions_per_location,
+                                 1, self._box_code_size]))
+    class_predictions_with_background = tf.reshape(
+        class_predictions_with_background,
+        tf.stack([combined_feature_map_shape[0],
+                  combined_feature_map_shape[1] *
+                  combined_feature_map_shape[2] *
+                  num_predictions_per_location,
+                  num_class_slots]))
     return {BOX_ENCODINGS: box_encodings,
             CLASS_PREDICTIONS_WITH_BACKGROUND:
             class_predictions_with_background}

object_detection/core/model.py

@@ -228,25 +228,24 @@ class DetectionModel(object):
           fields.BoxListFields.keypoints] = groundtruth_keypoints_list
 
   @abstractmethod
-  def restore_fn(self, checkpoint_path, from_detection_checkpoint=True):
-    """Return callable for loading a foreign checkpoint into tensorflow graph.
+  def restore_map(self, from_detection_checkpoint=True):
+    """Returns a map of variables to load from a foreign checkpoint.
 
-    Loads variables from a different tensorflow graph (typically feature
-    extractor variables). This enables the model to initialize based on weights
-    from another task. For example, the feature extractor variables from a
+    Returns a map of variable names to load from a checkpoint to variables in
+    the model graph. This enables the model to initialize based on weights from
+    another task. For example, the feature extractor variables from a
     classification model can be used to bootstrap training of an object
     detector. When loading from an object detection model, the checkpoint model
     should have the same parameters as this detection model with exception of
     the num_classes parameter.
 
     Args:
-      checkpoint_path: path to checkpoint to restore.
       from_detection_checkpoint: whether to restore from a full detection
         checkpoint (with compatible variable names) or to restore from a
         classification checkpoint for initialization prior to training.
 
     Returns:
-      a callable which takes a tf.Session as input and loads a checkpoint when
-        run.
+      A dict mapping variable names (to load from a checkpoint) to variables in
+      the model graph.
     """
     pass

object_detection/core/post_processing.py

@@ -174,7 +174,8 @@ def batch_multiclass_non_max_suppression(boxes,
                                          change_coordinate_frame=False,
                                          num_valid_boxes=None,
                                          masks=None,
-                                         scope=None):
+                                         scope=None,
+                                         parallel_iterations=32):
   """Multi-class version of non maximum suppression that operates on a batch.
 
   This op is similar to `multiclass_non_max_suppression` but operates on a batch
@@ -208,26 +209,28 @@ def batch_multiclass_non_max_suppression(boxes,
       float32 tensor containing box masks. `q` can be either number of classes
       or 1 depending on whether a separate mask is predicted per class.
     scope: tf scope name.
+    parallel_iterations: (optional) number of batch items to process in
+      parallel.
 
   Returns:
-    A dictionary containing the following entries:
-    'detection_boxes': A [batch_size, max_detections, 4] float32 tensor
+    'nmsed_boxes': A [batch_size, max_detections, 4] float32 tensor
       containing the non-max suppressed boxes.
-    'detection_scores': A [bath_size, max_detections] float32 tensor containing
+    'nmsed_scores': A [batch_size, max_detections] float32 tensor containing
       the scores for the boxes.
-    'detection_classes': A [batch_size, max_detections] float32 tensor
+    'nmsed_classes': A [batch_size, max_detections] float32 tensor
       containing the class for boxes.
-    'num_detections': A [batchsize] float32 tensor indicating the number of
+    'nmsed_masks': (optional) a
+      [batch_size, max_detections, mask_height, mask_width] float32 tensor
+      containing masks for each selected box. This is set to None if input
+      `masks` is None.
+    'num_detections': A [batch_size] int32 tensor indicating the number of
       valid detections per batch item. Only the top num_detections[i] entries in
       nms_boxes[i], nms_scores[i] and nms_class[i] are valid. the rest of the
       entries are zero paddings.
-    'detection_masks': (optional) a
-      [batch_size, max_detections, mask_height, mask_width] float32 tensor
-      containing masks for each selected box.
 
   Raises:
-    ValueError: if iou_thresh is not in [0, 1] or if input boxlist does not have
-      a valid scores field.
+    ValueError: if `q` in boxes.shape is not 1 or not equal to number of
+      classes as inferred from scores.shape.
   """
   q = boxes.shape[2].value
   num_classes = scores.shape[2].value
@@ -235,36 +238,45 @@ def batch_multiclass_non_max_suppression(boxes,
     raise ValueError('third dimension of boxes must be either 1 or equal '
                      'to the third dimension of scores')
 
+  original_masks = masks
   with tf.name_scope(scope, 'BatchMultiClassNonMaxSuppression'):
-    per_image_boxes_list = tf.unstack(boxes)
-    per_image_scores_list = tf.unstack(scores)
-    num_valid_boxes_list = len(per_image_boxes_list) * [None]
-    per_image_masks_list = len(per_image_boxes_list) * [None]
-    if num_valid_boxes is not None:
-      num_valid_boxes_list = tf.unstack(num_valid_boxes)
-    if masks is not None:
-      per_image_masks_list = tf.unstack(masks)
+    boxes_shape = boxes.shape
+    batch_size = boxes_shape[0].value
+    num_anchors = boxes_shape[1].value
+
+    if batch_size is None:
+      batch_size = tf.shape(boxes)[0]
+    if num_anchors is None:
+      num_anchors = tf.shape(boxes)[1]
+
+    # If num valid boxes aren't provided, create one and mark all boxes as
+    # valid.
+    if num_valid_boxes is None:
+      num_valid_boxes = tf.ones([batch_size], dtype=tf.int32) * num_anchors
 
-    detection_boxes_list = []
-    detection_scores_list = []
-    detection_classes_list = []
-    num_detections_list = []
-    detection_masks_list = []
-    for (per_image_boxes, per_image_scores, per_image_masks, num_valid_boxes
-        ) in zip(per_image_boxes_list, per_image_scores_list,
-                 per_image_masks_list, num_valid_boxes_list):
-      if num_valid_boxes is not None:
-        per_image_boxes = tf.reshape(
-            tf.slice(per_image_boxes, 3*[0],
-                     tf.stack([num_valid_boxes, -1, -1])), [-1, q, 4])
-        per_image_scores = tf.reshape(
-            tf.slice(per_image_scores, [0, 0],
-                     tf.stack([num_valid_boxes, -1])), [-1, num_classes])
-        if masks is not None:
-          per_image_masks = tf.reshape(
-              tf.slice(per_image_masks, 4*[0],
-                       tf.stack([num_valid_boxes, -1, -1, -1])),
-              [-1, q, masks.shape[3].value, masks.shape[4].value])
+    # If masks aren't provided, create dummy masks so we can only have one copy
+    # of single_image_nms_fn and discard the dummy masks after map_fn.
+    if masks is None:
+      masks_shape = tf.stack([batch_size, num_anchors, 1, 0, 0])
+      masks = tf.zeros(masks_shape)
+
+    def single_image_nms_fn(args):
+      """Runs NMS on a single image and returns padded output."""
+      (per_image_boxes, per_image_scores, per_image_masks,
+       per_image_num_valid_boxes) = args
+      per_image_boxes = tf.reshape(
+          tf.slice(per_image_boxes, 3 * [0],
+                   tf.stack([per_image_num_valid_boxes, -1, -1])), [-1, q, 4])
+      per_image_scores = tf.reshape(
+          tf.slice(per_image_scores, [0, 0],
+                   tf.stack([per_image_num_valid_boxes, -1])),
+          [-1, num_classes])
+
+      per_image_masks = tf.reshape(
+          tf.slice(per_image_masks, 4 * [0],
+                   tf.stack([per_image_num_valid_boxes, -1, -1, -1])),
+          [-1, q, per_image_masks.shape[2].value,
+           per_image_masks.shape[3].value])
       nmsed_boxlist = multiclass_non_max_suppression(
           per_image_boxes,
           per_image_scores,
@@ -275,24 +287,26 @@ def batch_multiclass_non_max_suppression(boxes,
           masks=per_image_masks,
           clip_window=clip_window,
           change_coordinate_frame=change_coordinate_frame)
-      num_detections_list.append(tf.to_float(nmsed_boxlist.num_boxes()))
       padded_boxlist = box_list_ops.pad_or_clip_box_list(nmsed_boxlist,
                                                          max_total_size)
-      detection_boxes_list.append(padded_boxlist.get())
-      detection_scores_list.append(
-          padded_boxlist.get_field(fields.BoxListFields.scores))
-      detection_classes_list.append(
-          padded_boxlist.get_field(fields.BoxListFields.classes))
-      if masks is not None:
-        detection_masks_list.append(
-            padded_boxlist.get_field(fields.BoxListFields.masks))
+      num_detections = nmsed_boxlist.num_boxes()
+      nmsed_boxes = padded_boxlist.get()
+      nmsed_scores = padded_boxlist.get_field(fields.BoxListFields.scores)
+      nmsed_classes = padded_boxlist.get_field(fields.BoxListFields.classes)
+      nmsed_masks = padded_boxlist.get_field(fields.BoxListFields.masks)
+      return [nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+              num_detections]
 
-    nms_dict = {
-        'detection_boxes': tf.stack(detection_boxes_list),
-        'detection_scores': tf.stack(detection_scores_list),
-        'detection_classes': tf.stack(detection_classes_list),
-        'num_detections': tf.stack(num_detections_list)
-    }
-    if masks is not None:
-      nms_dict['detection_masks'] = tf.stack(detection_masks_list)
-    return nms_dict
+    (batch_nmsed_boxes, batch_nmsed_scores,
+     batch_nmsed_classes, batch_nmsed_masks,
+     batch_num_detections) = tf.map_fn(
+         single_image_nms_fn,
+         elems=[boxes, scores, masks, num_valid_boxes],
+         dtype=[tf.float32, tf.float32, tf.float32, tf.float32, tf.int32],
+         parallel_iterations=parallel_iterations)
+
+    if original_masks is None:
+      batch_nmsed_masks = None
+
+    return (batch_nmsed_boxes, batch_nmsed_scores, batch_nmsed_classes,
+            batch_nmsed_masks, batch_num_detections)

object_detection/core/post_processing_test.py

@@ -496,15 +496,21 @@ class MulticlassNonMaxSuppressionTest(tf.test.TestCase):
     exp_nms_scores = [[.95, .9, .85, .3]]
     exp_nms_classes = [[0, 0, 1, 0]]
 
-    nms_dict = post_processing.batch_multiclass_non_max_suppression(
-        boxes, scores, score_thresh, iou_thresh,
-        max_size_per_class=max_output_size, max_total_size=max_output_size)
+    (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+     num_detections) = post_processing.batch_multiclass_non_max_suppression(
+         boxes, scores, score_thresh, iou_thresh,
+         max_size_per_class=max_output_size, max_total_size=max_output_size)
+
+    self.assertIsNone(nmsed_masks)
+
     with self.test_session() as sess:
-      nms_output = sess.run(nms_dict)
-      self.assertAllClose(nms_output['detection_boxes'], exp_nms_corners)
-      self.assertAllClose(nms_output['detection_scores'], exp_nms_scores)
-      self.assertAllClose(nms_output['detection_classes'], exp_nms_classes)
-      self.assertEqual(nms_output['num_detections'], [4])
+      (nmsed_boxes, nmsed_scores, nmsed_classes,
+       num_detections) = sess.run([nmsed_boxes, nmsed_scores, nmsed_classes,
+                                   num_detections])
+      self.assertAllClose(nmsed_boxes, exp_nms_corners)
+      self.assertAllClose(nmsed_scores, exp_nms_scores)
+      self.assertAllClose(nmsed_classes, exp_nms_classes)
+      self.assertEqual(num_detections, [4])
 
   def test_batch_multiclass_nms_with_batch_size_2(self):
     boxes = tf.constant([[[[0, 0, 1, 1], [0, 0, 4, 5]],
@@ -524,28 +530,42 @@ class MulticlassNonMaxSuppressionTest(tf.test.TestCase):
     iou_thresh = .5
     max_output_size = 4
 
-    exp_nms_corners = [[[0, 10, 1, 11],
-                        [0, 0, 1, 1],
-                        [0, 0, 0, 0],
-                        [0, 0, 0, 0]],
-                       [[0, 999, 2, 1004],
-                        [0, 10.1, 1, 11.1],
-                        [0, 100, 1, 101],
-                        [0, 0, 0, 0]]]
-    exp_nms_scores = [[.95, .9, 0, 0],
-                      [.85, .5, .3, 0]]
-    exp_nms_classes = [[0, 0, 0, 0],
-                       [1, 0, 0, 0]]
+    exp_nms_corners = np.array([[[0, 10, 1, 11],
+                                 [0, 0, 1, 1],
+                                 [0, 0, 0, 0],
+                                 [0, 0, 0, 0]],
+                                [[0, 999, 2, 1004],
+                                 [0, 10.1, 1, 11.1],
+                                 [0, 100, 1, 101],
+                                 [0, 0, 0, 0]]])
+    exp_nms_scores = np.array([[.95, .9, 0, 0],
+                               [.85, .5, .3, 0]])
+    exp_nms_classes = np.array([[0, 0, 0, 0],
+                                [1, 0, 0, 0]])
+
+    (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+     num_detections) = post_processing.batch_multiclass_non_max_suppression(
+         boxes, scores, score_thresh, iou_thresh,
+         max_size_per_class=max_output_size, max_total_size=max_output_size)
+
+    self.assertIsNone(nmsed_masks)
+    # Check static shapes
+    self.assertAllEqual(nmsed_boxes.shape.as_list(),
+                        exp_nms_corners.shape)
+    self.assertAllEqual(nmsed_scores.shape.as_list(),
+                        exp_nms_scores.shape)
+    self.assertAllEqual(nmsed_classes.shape.as_list(),
+                        exp_nms_classes.shape)
+    self.assertEqual(num_detections.shape.as_list(), [2])
 
-    nms_dict = post_processing.batch_multiclass_non_max_suppression(
-        boxes, scores, score_thresh, iou_thresh,
-        max_size_per_class=max_output_size, max_total_size=max_output_size)
     with self.test_session() as sess:
-      nms_output = sess.run(nms_dict)
-      self.assertAllClose(nms_output['detection_boxes'], exp_nms_corners)
-      self.assertAllClose(nms_output['detection_scores'], exp_nms_scores)
-      self.assertAllClose(nms_output['detection_classes'], exp_nms_classes)
-      self.assertAllClose(nms_output['num_detections'], [2, 3])
+      (nmsed_boxes, nmsed_scores, nmsed_classes,
+       num_detections) = sess.run([nmsed_boxes, nmsed_scores, nmsed_classes,
+                                   num_detections])
+      self.assertAllClose(nmsed_boxes, exp_nms_corners)
+      self.assertAllClose(nmsed_scores, exp_nms_scores)
+      self.assertAllClose(nmsed_classes, exp_nms_classes)
+      self.assertAllClose(num_detections, [2, 3])
 
   def test_batch_multiclass_nms_with_masks(self):
     boxes = tf.constant([[[[0, 0, 1, 1], [0, 0, 4, 5]],
@@ -574,38 +594,126 @@ class MulticlassNonMaxSuppressionTest(tf.test.TestCase):
     iou_thresh = .5
     max_output_size = 4
 
-    exp_nms_corners = [[[0, 10, 1, 11],
-                        [0, 0, 1, 1],
-                        [0, 0, 0, 0],
-                        [0, 0, 0, 0]],
-                       [[0, 999, 2, 1004],
-                        [0, 10.1, 1, 11.1],
-                        [0, 100, 1, 101],
-                        [0, 0, 0, 0]]]
-    exp_nms_scores = [[.95, .9, 0, 0],
-                      [.85, .5, .3, 0]]
-    exp_nms_classes = [[0, 0, 0, 0],
-                       [1, 0, 0, 0]]
-    exp_nms_masks = [[[[6, 7], [8, 9]],
-                      [[0, 1], [2, 3]],
-                      [[0, 0], [0, 0]],
-                      [[0, 0], [0, 0]]],
-                     [[[13, 14], [15, 16]],
-                      [[8, 9], [10, 11]],
-                      [[10, 11], [12, 13]],
-                      [[0, 0], [0, 0]]]]
+    exp_nms_corners = np.array([[[0, 10, 1, 11],
+                                 [0, 0, 1, 1],
+                                 [0, 0, 0, 0],
+                                 [0, 0, 0, 0]],
+                                [[0, 999, 2, 1004],
+                                 [0, 10.1, 1, 11.1],
+                                 [0, 100, 1, 101],
+                                 [0, 0, 0, 0]]])
+    exp_nms_scores = np.array([[.95, .9, 0, 0],
+                               [.85, .5, .3, 0]])
+    exp_nms_classes = np.array([[0, 0, 0, 0],
+                                [1, 0, 0, 0]])
+    exp_nms_masks = np.array([[[[6, 7], [8, 9]],
+                               [[0, 1], [2, 3]],
+                               [[0, 0], [0, 0]],
+                               [[0, 0], [0, 0]]],
+                              [[[13, 14], [15, 16]],
+                               [[8, 9], [10, 11]],
+                               [[10, 11], [12, 13]],
+                               [[0, 0], [0, 0]]]])
+
+    (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+     num_detections) = post_processing.batch_multiclass_non_max_suppression(
+         boxes, scores, score_thresh, iou_thresh,
+         max_size_per_class=max_output_size, max_total_size=max_output_size,
+         masks=masks)
+
+    # Check static shapes
+    self.assertAllEqual(nmsed_boxes.shape.as_list(), exp_nms_corners.shape)
+    self.assertAllEqual(nmsed_scores.shape.as_list(), exp_nms_scores.shape)
+    self.assertAllEqual(nmsed_classes.shape.as_list(), exp_nms_classes.shape)
+    self.assertAllEqual(nmsed_masks.shape.as_list(), exp_nms_masks.shape)
+    self.assertEqual(num_detections.shape.as_list(), [2])
+
+    with self.test_session() as sess:
+      (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+       num_detections) = sess.run([nmsed_boxes, nmsed_scores, nmsed_classes,
+                                   nmsed_masks, num_detections])
+
+      self.assertAllClose(nmsed_boxes, exp_nms_corners)
+      self.assertAllClose(nmsed_scores, exp_nms_scores)
+      self.assertAllClose(nmsed_classes, exp_nms_classes)
+      self.assertAllClose(num_detections, [2, 3])
+      self.assertAllClose(nmsed_masks, exp_nms_masks)
+
+  def test_batch_multiclass_nms_with_dynamic_batch_size(self):
+    boxes_placeholder = tf.placeholder(tf.float32, shape=(None, None, 2, 4))
+    scores_placeholder = tf.placeholder(tf.float32, shape=(None, None, 2))
+    masks_placeholder = tf.placeholder(tf.float32, shape=(None, None, 2, 2, 2))
+
+    boxes = np.array([[[[0, 0, 1, 1], [0, 0, 4, 5]],
+                       [[0, 0.1, 1, 1.1], [0, 0.1, 2, 1.1]],
+                       [[0, -0.1, 1, 0.9], [0, -0.1, 1, 0.9]],
+                       [[0, 10, 1, 11], [0, 10, 1, 11]]],
+                      [[[0, 10.1, 1, 11.1], [0, 10.1, 1, 11.1]],
+                       [[0, 100, 1, 101], [0, 100, 1, 101]],
+                       [[0, 1000, 1, 1002], [0, 999, 2, 1004]],
+                       [[0, 1000, 1, 1002.1], [0, 999, 2, 1002.7]]]])
+    scores = np.array([[[.9, 0.01], [.75, 0.05],
+                        [.6, 0.01], [.95, 0]],
+                       [[.5, 0.01], [.3, 0.01],
+                        [.01, .85], [.01, .5]]])
+    masks = np.array([[[[[0, 1], [2, 3]], [[1, 2], [3, 4]]],
+                       [[[2, 3], [4, 5]], [[3, 4], [5, 6]]],
+                       [[[4, 5], [6, 7]], [[5, 6], [7, 8]]],
+                       [[[6, 7], [8, 9]], [[7, 8], [9, 10]]]],
+                      [[[[8, 9], [10, 11]], [[9, 10], [11, 12]]],
+                       [[[10, 11], [12, 13]], [[11, 12], [13, 14]]],
+                       [[[12, 13], [14, 15]], [[13, 14], [15, 16]]],
+                       [[[14, 15], [16, 17]], [[15, 16], [17, 18]]]]])
+    score_thresh = 0.1
+    iou_thresh = .5
+    max_output_size = 4
+
+    exp_nms_corners = np.array([[[0, 10, 1, 11],
+                                 [0, 0, 1, 1],
+                                 [0, 0, 0, 0],
+                                 [0, 0, 0, 0]],
+                                [[0, 999, 2, 1004],
+                                 [0, 10.1, 1, 11.1],
+                                 [0, 100, 1, 101],
+                                 [0, 0, 0, 0]]])
+    exp_nms_scores = np.array([[.95, .9, 0, 0],
+                               [.85, .5, .3, 0]])
+    exp_nms_classes = np.array([[0, 0, 0, 0],
+                                [1, 0, 0, 0]])
+    exp_nms_masks = np.array([[[[6, 7], [8, 9]],
+                               [[0, 1], [2, 3]],
+                               [[0, 0], [0, 0]],
+                               [[0, 0], [0, 0]]],
+                              [[[13, 14], [15, 16]],
+                               [[8, 9], [10, 11]],
+                               [[10, 11], [12, 13]],
+                               [[0, 0], [0, 0]]]])
+
+    (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+     num_detections) = post_processing.batch_multiclass_non_max_suppression(
+         boxes_placeholder, scores_placeholder, score_thresh, iou_thresh,
+         max_size_per_class=max_output_size, max_total_size=max_output_size,
+         masks=masks_placeholder)
+
+    # Check static shapes
+    self.assertAllEqual(nmsed_boxes.shape.as_list(), [None, 4, 4])
+    self.assertAllEqual(nmsed_scores.shape.as_list(), [None, 4])
+    self.assertAllEqual(nmsed_classes.shape.as_list(), [None, 4])
+    self.assertAllEqual(nmsed_masks.shape.as_list(), [None, 4, 2, 2])
+    self.assertEqual(num_detections.shape.as_list(), [None])
 
-    nms_dict = post_processing.batch_multiclass_non_max_suppression(
-        boxes, scores, score_thresh, iou_thresh,
-        max_size_per_class=max_output_size, max_total_size=max_output_size,
-        masks=masks)
     with self.test_session() as sess:
-      nms_output = sess.run(nms_dict)
-      self.assertAllClose(nms_output['detection_boxes'], exp_nms_corners)
-      self.assertAllClose(nms_output['detection_scores'], exp_nms_scores)
-      self.assertAllClose(nms_output['detection_classes'], exp_nms_classes)
-      self.assertAllClose(nms_output['num_detections'], [2, 3])
-      self.assertAllClose(nms_output['detection_masks'], exp_nms_masks)
+      (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+       num_detections) = sess.run([nmsed_boxes, nmsed_scores, nmsed_classes,
+                                   nmsed_masks, num_detections],
+                                  feed_dict={boxes_placeholder: boxes,
+                                             scores_placeholder: scores,
+                                             masks_placeholder: masks})
+      self.assertAllClose(nmsed_boxes, exp_nms_corners)
+      self.assertAllClose(nmsed_scores, exp_nms_scores)
+      self.assertAllClose(nmsed_classes, exp_nms_classes)
+      self.assertAllClose(num_detections, [2, 3])
+      self.assertAllClose(nmsed_masks, exp_nms_masks)
 
   def test_batch_multiclass_nms_with_masks_and_num_valid_boxes(self):
     boxes = tf.constant([[[[0, 0, 1, 1], [0, 0, 4, 5]],
@@ -656,17 +764,21 @@ class MulticlassNonMaxSuppressionTest(tf.test.TestCase):
                       [[0, 0], [0, 0]],
                       [[0, 0], [0, 0]]]]
 
-    nms_dict = post_processing.batch_multiclass_non_max_suppression(
-        boxes, scores, score_thresh, iou_thresh,
-        max_size_per_class=max_output_size, max_total_size=max_output_size,
-        num_valid_boxes=num_valid_boxes, masks=masks)
+    (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+     num_detections) = post_processing.batch_multiclass_non_max_suppression(
+         boxes, scores, score_thresh, iou_thresh,
+         max_size_per_class=max_output_size, max_total_size=max_output_size,
+         num_valid_boxes=num_valid_boxes, masks=masks)
+
     with self.test_session() as sess:
-      nms_output = sess.run(nms_dict)
-      self.assertAllClose(nms_output['detection_boxes'], exp_nms_corners)
-      self.assertAllClose(nms_output['detection_scores'], exp_nms_scores)
-      self.assertAllClose(nms_output['detection_classes'], exp_nms_classes)
-      self.assertAllClose(nms_output['num_detections'], [1, 1])
-      self.assertAllClose(nms_output['detection_masks'], exp_nms_masks)
+      (nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_masks,
+       num_detections) = sess.run([nmsed_boxes, nmsed_scores, nmsed_classes,
+                                   nmsed_masks, num_detections])
+      self.assertAllClose(nmsed_boxes, exp_nms_corners)
+      self.assertAllClose(nmsed_scores, exp_nms_scores)
+      self.assertAllClose(nmsed_classes, exp_nms_classes)
+      self.assertAllClose(num_detections, [1, 1])
+      self.assertAllClose(nmsed_masks, exp_nms_masks)
 
 
 if __name__ == '__main__':

object_detection/core/preprocessor.py

@@ -1255,6 +1255,82 @@ def random_resize_method(image, target_size):
   return resized_image
 
 
+def _compute_new_static_size(image,
+                             min_dimension,
+                             max_dimension):
+  """Compute new static shape for resize_to_range method."""
+  image_shape = image.get_shape().as_list()
+  orig_height = image_shape[0]
+  orig_width = image_shape[1]
+  orig_min_dim = min(orig_height, orig_width)
+  # Calculates the larger of the possible sizes
+  large_scale_factor = min_dimension / float(orig_min_dim)
+  # Scaling orig_(height|width) by large_scale_factor will make the smaller
+  # dimension equal to min_dimension, save for floating point rounding errors.
+  # For reasonably-sized images, taking the nearest integer will reliably
+  # eliminate this error.
+  large_height = int(round(orig_height * large_scale_factor))
+  large_width = int(round(orig_width * large_scale_factor))
+  large_size = [large_height, large_width]
+  if max_dimension:
+    # Calculates the smaller of the possible sizes, use that if the larger
+    # is too big.
+    orig_max_dim = max(orig_height, orig_width)
+    small_scale_factor = max_dimension / float(orig_max_dim)
+    # Scaling orig_(height|width) by small_scale_factor will make the larger
+    # dimension equal to max_dimension, save for floating point rounding
+    # errors. For reasonably-sized images, taking the nearest integer will
+    # reliably eliminate this error.
+    small_height = int(round(orig_height * small_scale_factor))
+    small_width = int(round(orig_width * small_scale_factor))
+    small_size = [small_height, small_width]
+    new_size = large_size
+    if max(large_size) > max_dimension:
+      new_size = small_size
+  else:
+    new_size = large_size
+  return tf.constant(new_size)
+
+
+def _compute_new_dynamic_size(image,
+                              min_dimension,
+                              max_dimension):
+  """Compute new dynamic shape for resize_to_range method."""
+  image_shape = tf.shape(image)
+  orig_height = tf.to_float(image_shape[0])
+  orig_width = tf.to_float(image_shape[1])
+  orig_min_dim = tf.minimum(orig_height, orig_width)
+  # Calculates the larger of the possible sizes
+  min_dimension = tf.constant(min_dimension, dtype=tf.float32)
+  large_scale_factor = min_dimension / orig_min_dim
+  # Scaling orig_(height|width) by large_scale_factor will make the smaller
+  # dimension equal to min_dimension, save for floating point rounding errors.
+  # For reasonably-sized images, taking the nearest integer will reliably
+  # eliminate this error.
+  large_height = tf.to_int32(tf.round(orig_height * large_scale_factor))
+  large_width = tf.to_int32(tf.round(orig_width * large_scale_factor))
+  large_size = tf.stack([large_height, large_width])
+  if max_dimension:
+    # Calculates the smaller of the possible sizes, use that if the larger
+    # is too big.
+    orig_max_dim = tf.maximum(orig_height, orig_width)
+    max_dimension = tf.constant(max_dimension, dtype=tf.float32)
+    small_scale_factor = max_dimension / orig_max_dim
+    # Scaling orig_(height|width) by small_scale_factor will make the larger
+    # dimension equal to max_dimension, save for floating point rounding
+    # errors. For reasonably-sized images, taking the nearest integer will
+    # reliably eliminate this error.
+    small_height = tf.to_int32(tf.round(orig_height * small_scale_factor))
+    small_width = tf.to_int32(tf.round(orig_width * small_scale_factor))
+    small_size = tf.stack([small_height, small_width])
+    new_size = tf.cond(
+        tf.to_float(tf.reduce_max(large_size)) > max_dimension,
+        lambda: small_size, lambda: large_size)
+  else:
+    new_size = large_size
+  return new_size
+
+
 def resize_to_range(image,
                     masks=None,
                     min_dimension=None,
@@ -1295,64 +1371,22 @@ def resize_to_range(image,
     raise ValueError('Image should be 3D tensor')
 
   with tf.name_scope('ResizeToRange', values=[image, min_dimension]):
-    image_shape = tf.shape(image)
-    orig_height = tf.to_float(image_shape[0])
-    orig_width = tf.to_float(image_shape[1])
-    orig_min_dim = tf.minimum(orig_height, orig_width)
-
-    # Calculates the larger of the possible sizes
-    min_dimension = tf.constant(min_dimension, dtype=tf.float32)
-    large_scale_factor = min_dimension / orig_min_dim
-    # Scaling orig_(height|width) by large_scale_factor will make the smaller
-    # dimension equal to min_dimension, save for floating point rounding errors.
-    # For reasonably-sized images, taking the nearest integer will reliably
-    # eliminate this error.
-    large_height = tf.to_int32(tf.round(orig_height * large_scale_factor))
-    large_width = tf.to_int32(tf.round(orig_width * large_scale_factor))
-    large_size = tf.stack([large_height, large_width])
-
-    if max_dimension:
-      # Calculates the smaller of the possible sizes, use that if the larger
-      # is too big.
-      orig_max_dim = tf.maximum(orig_height, orig_width)
-      max_dimension = tf.constant(max_dimension, dtype=tf.float32)
-      small_scale_factor = max_dimension / orig_max_dim
-      # Scaling orig_(height|width) by small_scale_factor will make the larger
-      # dimension equal to max_dimension, save for floating point rounding
-      # errors. For reasonably-sized images, taking the nearest integer will
-      # reliably eliminate this error.
-      small_height = tf.to_int32(tf.round(orig_height * small_scale_factor))
-      small_width = tf.to_int32(tf.round(orig_width * small_scale_factor))
-      small_size = tf.stack([small_height, small_width])
-
-      new_size = tf.cond(
-          tf.to_float(tf.reduce_max(large_size)) > max_dimension,
-          lambda: small_size, lambda: large_size)
+    if image.get_shape().is_fully_defined():
+      new_size = _compute_new_static_size(image, min_dimension,
+                                          max_dimension)
     else:
-      new_size = large_size
-
+      new_size = _compute_new_dynamic_size(image, min_dimension,
+                                           max_dimension)
     new_image = tf.image.resize_images(image, new_size,
                                        align_corners=align_corners)
 
     result = new_image
     if masks is not None:
-      num_instances = tf.shape(masks)[0]
-
-      def resize_masks_branch():
-        new_masks = tf.expand_dims(masks, 3)
-        new_masks = tf.image.resize_nearest_neighbor(
-            new_masks, new_size, align_corners=align_corners)
-        new_masks = tf.squeeze(new_masks, axis=3)
-        return new_masks
-
-      def reshape_masks_branch():
-        new_masks = tf.reshape(masks, [0, new_size[0], new_size[1]])
-        return new_masks
-
-      masks = tf.cond(num_instances > 0,
-                      resize_masks_branch,
-                      reshape_masks_branch)
-      result = [new_image, masks]
+      new_masks = tf.expand_dims(masks, 3)
+      new_masks = tf.image.resize_nearest_neighbor(new_masks, new_size,
+                                                   align_corners=align_corners)
+      new_masks = tf.squeeze(new_masks, 3)
+      result = [new_image, new_masks]
 
     return result
 

object_detection/core/preprocessor_test.py

@@ -1395,7 +1395,7 @@ class PreprocessorTest(tf.test.TestCase):
       self.assertAllEqual(expected_images_shape_,
                           resized_images_shape_)
 
-  def testResizeToRange(self):
+  def testResizeToRangePreservesStaticSpatialShape(self):
     """Tests image resizing, checking output sizes."""
     in_shape_list = [[60, 40, 3], [15, 30, 3], [15, 50, 3]]
     min_dim = 50
@@ -1406,13 +1406,27 @@ class PreprocessorTest(tf.test.TestCase):
       in_image = tf.random_uniform(in_shape)
       out_image = preprocessor.resize_to_range(
           in_image, min_dimension=min_dim, max_dimension=max_dim)
-      out_image_shape = tf.shape(out_image)
+      self.assertAllEqual(out_image.get_shape().as_list(), expected_shape)
+
+  def testResizeToRangeWithDynamicSpatialShape(self):
+    """Tests image resizing, checking output sizes."""
+    in_shape_list = [[60, 40, 3], [15, 30, 3], [15, 50, 3]]
+    min_dim = 50
+    max_dim = 100
+    expected_shape_list = [[75, 50, 3], [50, 100, 3], [30, 100, 3]]
 
+    for in_shape, expected_shape in zip(in_shape_list, expected_shape_list):
+      in_image = tf.placeholder(tf.float32, shape=(None, None, 3))
+      out_image = preprocessor.resize_to_range(
+          in_image, min_dimension=min_dim, max_dimension=max_dim)
+      out_image_shape = tf.shape(out_image)
       with self.test_session() as sess:
-        out_image_shape = sess.run(out_image_shape)
+        out_image_shape = sess.run(out_image_shape,
+                                   feed_dict={in_image:
+                                              np.random.randn(*in_shape)})
         self.assertAllEqual(out_image_shape, expected_shape)
 
-  def testResizeToRangeWithMasks(self):
+  def testResizeToRangeWithMasksPreservesStaticSpatialShape(self):
     """Tests image resizing, checking output sizes."""
     in_image_shape_list = [[60, 40, 3], [15, 30, 3]]
     in_masks_shape_list = [[15, 60, 40], [10, 15, 30]]
@@ -1430,30 +1444,25 @@ class PreprocessorTest(tf.test.TestCase):
       in_masks = tf.random_uniform(in_masks_shape)
       out_image, out_masks = preprocessor.resize_to_range(
           in_image, in_masks, min_dimension=min_dim, max_dimension=max_dim)
-      out_image_shape = tf.shape(out_image)
-      out_masks_shape = tf.shape(out_masks)
-
-      with self.test_session() as sess:
-        out_image_shape, out_masks_shape = sess.run(
-            [out_image_shape, out_masks_shape])
-        self.assertAllEqual(out_image_shape, expected_image_shape)
-        self.assertAllEqual(out_masks_shape, expected_mask_shape)
+      self.assertAllEqual(out_masks.get_shape().as_list(), expected_mask_shape)
+      self.assertAllEqual(out_image.get_shape().as_list(), expected_image_shape)
 
-  def testResizeToRangeWithNoInstanceMask(self):
+  def testResizeToRangeWithMasksAndDynamicSpatialShape(self):
     """Tests image resizing, checking output sizes."""
     in_image_shape_list = [[60, 40, 3], [15, 30, 3]]
-    in_masks_shape_list = [[0, 60, 40], [0, 15, 30]]
+    in_masks_shape_list = [[15, 60, 40], [10, 15, 30]]
     min_dim = 50
     max_dim = 100
     expected_image_shape_list = [[75, 50, 3], [50, 100, 3]]
-    expected_masks_shape_list = [[0, 75, 50], [0, 50, 100]]
+    expected_masks_shape_list = [[15, 75, 50], [10, 50, 100]]
 
     for (in_image_shape, expected_image_shape, in_masks_shape,
          expected_mask_shape) in zip(in_image_shape_list,
                                      expected_image_shape_list,
                                      in_masks_shape_list,
                                      expected_masks_shape_list):
-      in_image = tf.random_uniform(in_image_shape)
+      in_image = tf.placeholder(tf.float32, shape=(None, None, 3))
+      in_masks = tf.placeholder(tf.float32, shape=(None, None, None))
       in_masks = tf.random_uniform(in_masks_shape)
       out_image, out_masks = preprocessor.resize_to_range(
           in_image, in_masks, min_dimension=min_dim, max_dimension=max_dim)
@@ -1462,38 +1471,15 @@ class PreprocessorTest(tf.test.TestCase):
 
       with self.test_session() as sess:
         out_image_shape, out_masks_shape = sess.run(
-            [out_image_shape, out_masks_shape])
-        self.assertAllEqual(out_image_shape, expected_image_shape)
-        self.assertAllEqual(out_masks_shape, expected_mask_shape)
-
-  def testResizeImageWithMasks(self):
-    """Tests image resizing, checking output sizes."""
-    in_image_shape_list = [[60, 40, 3], [15, 30, 3]]
-    in_masks_shape_list = [[15, 60, 40], [10, 15, 30]]
-    height = 50
-    width = 100
-    expected_image_shape_list = [[50, 100, 3], [50, 100, 3]]
-    expected_masks_shape_list = [[15, 50, 100], [10, 50, 100]]
-
-    for (in_image_shape, expected_image_shape, in_masks_shape,
-         expected_mask_shape) in zip(in_image_shape_list,
-                                     expected_image_shape_list,
-                                     in_masks_shape_list,
-                                     expected_masks_shape_list):
-      in_image = tf.random_uniform(in_image_shape)
-      in_masks = tf.random_uniform(in_masks_shape)
-      out_image, out_masks = preprocessor.resize_image(
-          in_image, in_masks, new_height=height, new_width=width)
-      out_image_shape = tf.shape(out_image)
-      out_masks_shape = tf.shape(out_masks)
-
-      with self.test_session() as sess:
-        out_image_shape, out_masks_shape = sess.run(
-            [out_image_shape, out_masks_shape])
+            [out_image_shape, out_masks_shape],
+            feed_dict={
+                in_image: np.random.randn(*in_image_shape),
+                in_masks: np.random.randn(*in_masks_shape)
+            })
         self.assertAllEqual(out_image_shape, expected_image_shape)
         self.assertAllEqual(out_masks_shape, expected_mask_shape)
 
-  def testResizeImageWithNoInstanceMask(self):
+  def testResizeToRangeWithInstanceMasksTensorOfSizeZero(self):
     """Tests image resizing, checking output sizes."""
     in_image_shape_list = [[60, 40, 3], [15, 30, 3]]
     in_masks_shape_list = [[0, 60, 40], [0, 15, 30]]

object_detection/data/pascal_label_map.pbtxt

-item {
-  id: 0
-  name: 'none_of_the_above'
-}
-
 item {
   id: 1
   name: 'aeroplane'

object_detection/data/pet_label_map.pbtxt

-item {
-  id: 0
-  name: 'none_of_the_above'
-}
-
 item {
   id: 1
   name: 'Abyssinian'

object_detection/export_inference_graph.py

@@ -16,16 +16,19 @@
 r"""Tool to export an object detection model for inference.
 
 Prepares an object detection tensorflow graph for inference using model
-configuration and an optional trained checkpoint. Outputs either an inference
-graph or a SavedModel (https://tensorflow.github.io/serving/serving_basic.html).
+configuration and an optional trained checkpoint. Outputs inference
+graph, associated checkpoint files, a frozen inference graph and a
+SavedModel (https://tensorflow.github.io/serving/serving_basic.html).
 
 The inference graph contains one of three input nodes depending on the user
 specified option.
-  * `image_tensor`: Accepts a uint8 4-D tensor of shape [1, None, None, 3]
-  * `encoded_image_string_tensor`: Accepts a scalar string tensor of encoded PNG
-    or JPEG image.
-  * `tf_example`: Accepts a serialized TFExample proto. The batch size in this
-    case is always 1.
+  * `image_tensor`: Accepts a uint8 4-D tensor of shape [None, None, None, 3]
+  * `encoded_image_string_tensor`: Accepts a 1-D string tensor of shape [None]
+    containing encoded PNG or JPEG images. Image resolutions are expected to be
+    the same if more than 1 image is provided.
+  * `tf_example`: Accepts a 1-D string tensor of shape [None] containing
+    serialized TFExample protos. Image resolutions are expected to be the same
+    if more than 1 image is provided.
 
 and the following output nodes returned by the model.postprocess(..):
   * `num_detections`: Outputs float32 tensors of the form [batch]
@@ -41,23 +44,27 @@ and the following output nodes returned by the model.postprocess(..):
       masks for each box if its present in the dictionary of postprocessed
       tensors returned by the model.
 
-Note that currently `batch` is always 1, but we will support `batch` > 1 in
-the future.
-
-Optionally, one can freeze the graph by converting the weights in the provided
-checkpoint as graph constants thereby eliminating the need to use a checkpoint
-file during inference.
-
-Note that this tool uses `use_moving_averages` from eval_config to decide
-which weights to freeze.
+Notes:
+ * This tool uses `use_moving_averages` from eval_config to decide which
+   weights to freeze.
 
 Example Usage:
 --------------
 python export_inference_graph \
     --input_type image_tensor \
     --pipeline_config_path path/to/ssd_inception_v2.config \
-    --checkpoint_path path/to/model-ckpt \
-    --inference_graph_path path/to/inference_graph.pb
+    --trained_checkpoint_prefix path/to/model.ckpt \
+    --output_directory path/to/exported_model_directory
+
+The expected output would be in the directory
+path/to/exported_model_directory (which is created if it does not exist)
+with contents:
+ - graph.pbtxt
+ - model.ckpt.data-00000-of-00001
+ - model.ckpt.info
+ - model.ckpt.meta
+ - frozen_inference_graph.pb
+ + saved_model (a directory)
 """
 import tensorflow as tf
 from google.protobuf import text_format
@@ -70,31 +77,29 @@ flags = tf.app.flags
 flags.DEFINE_string('input_type', 'image_tensor', 'Type of input node. Can be '
                     'one of [`image_tensor`, `encoded_image_string_tensor`, '
                     '`tf_example`]')
-flags.DEFINE_string('pipeline_config_path', '',
+flags.DEFINE_string('pipeline_config_path', None,
                     'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
                     'file.')
-flags.DEFINE_string('checkpoint_path', '', 'Optional path to checkpoint file. '
-                    'If provided, bakes the weights from the checkpoint into '
-                    'the graph.')
-flags.DEFINE_string('inference_graph_path', '', 'Path to write the output '
-                    'inference graph.')
-flags.DEFINE_bool('export_as_saved_model', False, 'Whether the exported graph '
-                  'should be saved as a SavedModel')
+flags.DEFINE_string('trained_checkpoint_prefix', None,
+                    'Path to trained checkpoint, typically of the form '
+                    'path/to/model.ckpt')
+flags.DEFINE_string('output_directory', None, 'Path to write outputs.')
 
 FLAGS = flags.FLAGS
 
 
 def main(_):
-  assert FLAGS.pipeline_config_path, 'TrainEvalPipelineConfig missing.'
-  assert FLAGS.inference_graph_path, 'Inference graph path missing.'
-  assert FLAGS.input_type, 'Input type missing.'
+  assert FLAGS.pipeline_config_path, '`pipeline_config_path` is missing'
+  assert FLAGS.trained_checkpoint_prefix, (
+         '`trained_checkpoint_prefix` is missing')
+  assert FLAGS.output_directory, '`output_directory` is missing'
+
   pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
   with tf.gfile.GFile(FLAGS.pipeline_config_path, 'r') as f:
     text_format.Merge(f.read(), pipeline_config)
-  exporter.export_inference_graph(FLAGS.input_type, pipeline_config,
-                                  FLAGS.checkpoint_path,
-                                  FLAGS.inference_graph_path,
-                                  FLAGS.export_as_saved_model)
+  exporter.export_inference_graph(
+      FLAGS.input_type, pipeline_config, FLAGS.trained_checkpoint_prefix,
+      FLAGS.output_directory)
 
 
 if __name__ == '__main__':

object_detection/exporter.py

@@ -17,6 +17,7 @@
 import logging
 import os
 import tensorflow as tf
+from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import pywrap_tensorflow
 from tensorflow.python.client import session
 from tensorflow.python.framework import graph_util
@@ -42,6 +43,7 @@ def freeze_graph_with_def_protos(
     filename_tensor_name,
     clear_devices,
     initializer_nodes,
+    optimize_graph=False,
     variable_names_blacklist=''):
   """Converts all variables in a graph and checkpoint into constants."""
   del restore_op_name, filename_tensor_name  # Unused by updated loading code.
@@ -61,86 +63,106 @@ def freeze_graph_with_def_protos(
     for node in input_graph_def.node:
       node.device = ''
 
-  _ = importer.import_graph_def(input_graph_def, name='')
-
-  with session.Session() as sess:
-    if input_saver_def:
-      saver = saver_lib.Saver(saver_def=input_saver_def)
-      saver.restore(sess, input_checkpoint)
+  with tf.Graph().as_default():
+    tf.import_graph_def(input_graph_def, name='')
+
+    if optimize_graph:
+      logging.info('Graph Rewriter optimizations enabled')
+      rewrite_options = rewriter_config_pb2.RewriterConfig(
+          optimize_tensor_layout=True)
+      rewrite_options.optimizers.append('pruning')
+      rewrite_options.optimizers.append('constfold')
+      rewrite_options.optimizers.append('layout')
+      graph_options = tf.GraphOptions(
+          rewrite_options=rewrite_options, infer_shapes=True)
     else:
-      var_list = {}
-      reader = pywrap_tensorflow.NewCheckpointReader(input_checkpoint)
-      var_to_shape_map = reader.get_variable_to_shape_map()
-      for key in var_to_shape_map:
-        try:
-          tensor = sess.graph.get_tensor_by_name(key + ':0')
-        except KeyError:
-          # This tensor doesn't exist in the graph (for example it's
-          # 'global_step' or a similar housekeeping element) so skip it.
-          continue
-        var_list[key] = tensor
-      saver = saver_lib.Saver(var_list=var_list)
-      saver.restore(sess, input_checkpoint)
-      if initializer_nodes:
-        sess.run(initializer_nodes)
-
-    variable_names_blacklist = (variable_names_blacklist.split(',') if
-                                variable_names_blacklist else None)
-    output_graph_def = graph_util.convert_variables_to_constants(
-        sess,
-        input_graph_def,
-        output_node_names.split(','),
-        variable_names_blacklist=variable_names_blacklist)
+      logging.info('Graph Rewriter optimizations disabled')
+      graph_options = tf.GraphOptions()
+    config = tf.ConfigProto(graph_options=graph_options)
+    with session.Session(config=config) as sess:
+      if input_saver_def:
+        saver = saver_lib.Saver(saver_def=input_saver_def)
+        saver.restore(sess, input_checkpoint)
+      else:
+        var_list = {}
+        reader = pywrap_tensorflow.NewCheckpointReader(input_checkpoint)
+        var_to_shape_map = reader.get_variable_to_shape_map()
+        for key in var_to_shape_map:
+          try:
+            tensor = sess.graph.get_tensor_by_name(key + ':0')
+          except KeyError:
+            # This tensor doesn't exist in the graph (for example it's
+            # 'global_step' or a similar housekeeping element) so skip it.
+            continue
+          var_list[key] = tensor
+        saver = saver_lib.Saver(var_list=var_list)
+        saver.restore(sess, input_checkpoint)
+        if initializer_nodes:
+          sess.run(initializer_nodes)
+
+      variable_names_blacklist = (variable_names_blacklist.split(',') if
+                                  variable_names_blacklist else None)
+      output_graph_def = graph_util.convert_variables_to_constants(
+          sess,
+          input_graph_def,
+          output_node_names.split(','),
+          variable_names_blacklist=variable_names_blacklist)
 
   return output_graph_def
 
 
-def get_frozen_graph_def(inference_graph_def, use_moving_averages,
-                         input_checkpoint, output_node_names):
-  """Freezes all variables in a graph definition."""
-  saver = None
-  if use_moving_averages:
-    variable_averages = tf.train.ExponentialMovingAverage(0.0)
-    variables_to_restore = variable_averages.variables_to_restore()
-    saver = tf.train.Saver(variables_to_restore)
-  else:
-    saver = tf.train.Saver()
 
-  frozen_graph_def = freeze_graph_with_def_protos(
-      input_graph_def=inference_graph_def,
-      input_saver_def=saver.as_saver_def(),
-      input_checkpoint=input_checkpoint,
-      output_node_names=output_node_names,
-      restore_op_name='save/restore_all',
-      filename_tensor_name='save/Const:0',
-      clear_devices=True,
-      initializer_nodes='')
-  return frozen_graph_def
+def _image_tensor_input_placeholder():
+  """Returns placeholder and input node that accepts a batch of uint8 images."""
+  input_tensor = tf.placeholder(dtype=tf.uint8,
+                                shape=(None, None, None, 3),
+                                name='image_tensor')
+  return input_tensor, input_tensor
 
 
-# TODO: Support batch tf example inputs.
 def _tf_example_input_placeholder():
-  tf_example_placeholder = tf.placeholder(
-      tf.string, shape=[], name='tf_example')
-  tensor_dict = tf_example_decoder.TfExampleDecoder().decode(
-      tf_example_placeholder)
-  image = tensor_dict[fields.InputDataFields.image]
-  return tf.expand_dims(image, axis=0)
+  """Returns input that accepts a batch of strings with tf examples.
 
-
-def _image_tensor_input_placeholder():
-  return tf.placeholder(dtype=tf.uint8,
-                        shape=(1, None, None, 3),
-                        name='image_tensor')
+  Returns:
+    a tuple of placeholder and input nodes that output decoded images.
+  """
+  batch_tf_example_placeholder = tf.placeholder(
+      tf.string, shape=[None], name='tf_example')
+  def decode(tf_example_string_tensor):
+    tensor_dict = tf_example_decoder.TfExampleDecoder().decode(
+        tf_example_string_tensor)
+    image_tensor = tensor_dict[fields.InputDataFields.image]
+    return image_tensor
+  return (batch_tf_example_placeholder,
+          tf.map_fn(decode,
+                    elems=batch_tf_example_placeholder,
+                    dtype=tf.uint8,
+                    parallel_iterations=32,
+                    back_prop=False))
 
 
 def _encoded_image_string_tensor_input_placeholder():
-  image_str = tf.placeholder(dtype=tf.string,
-                             shape=[],
-                             name='encoded_image_string_tensor')
-  image_tensor = tf.image.decode_image(image_str, channels=3)
-  image_tensor.set_shape((None, None, 3))
-  return tf.expand_dims(image_tensor, axis=0)
+  """Returns input that accepts a batch of PNG or JPEG strings.
+
+  Returns:
+    a tuple of placeholder and input nodes that output decoded images.
+  """
+  batch_image_str_placeholder = tf.placeholder(
+      dtype=tf.string,
+      shape=[None],
+      name='encoded_image_string_tensor')
+  def decode(encoded_image_string_tensor):
+    image_tensor = tf.image.decode_image(encoded_image_string_tensor,
+                                         channels=3)
+    image_tensor.set_shape((None, None, 3))
+    return image_tensor
+  return (batch_image_str_placeholder,
+          tf.map_fn(
+              decode,
+              elems=batch_image_str_placeholder,
+              dtype=tf.uint8,
+              parallel_iterations=32,
+              back_prop=False))
 
 
 input_placeholder_fn_map = {
@@ -151,7 +173,8 @@ input_placeholder_fn_map = {
 }
 
 
-def _add_output_tensor_nodes(postprocessed_tensors):
+def _add_output_tensor_nodes(postprocessed_tensors,
+                             output_collection_name='inference_op'):
   """Adds output nodes for detection boxes and scores.
 
   Adds the following nodes for output tensors -
@@ -174,6 +197,7 @@ def _add_output_tensor_nodes(postprocessed_tensors):
       'detection_masks': [batch, max_detections, mask_height, mask_width]
         (optional).
       'num_detections': [batch]
+    output_collection_name: Name of collection to add output tensors to.
 
   Returns:
     A tensor dict containing the added output tensor nodes.
@@ -191,53 +215,29 @@ def _add_output_tensor_nodes(postprocessed_tensors):
   outputs['num_detections'] = tf.identity(num_detections, name='num_detections')
   if masks is not None:
     outputs['detection_masks'] = tf.identity(masks, name='detection_masks')
+  for output_key in outputs:
+    tf.add_to_collection(output_collection_name, outputs[output_key])
+  if masks is not None:
+    tf.add_to_collection(output_collection_name, outputs['detection_masks'])
   return outputs
 
 
-def _write_inference_graph(inference_graph_path,
-                           checkpoint_path=None,
-                           use_moving_averages=False,
-                           output_node_names=(
-                               'num_detections,detection_scores,'
-                               'detection_boxes,detection_classes')):
-  """Writes inference graph to disk with the option to bake in weights.
-
-  If checkpoint_path is not None bakes the weights into the graph thereby
-  eliminating the need of checkpoint files during inference. If the model
-  was trained with moving averages, setting use_moving_averages to true
-  restores the moving averages, otherwise the original set of variables
-  is restored.
+def _write_frozen_graph(frozen_graph_path, frozen_graph_def):
+  """Writes frozen graph to disk.
 
   Args:
-    inference_graph_path: Path to write inference graph.
-    checkpoint_path: Optional path to the checkpoint file.
-    use_moving_averages: Whether to export the original or the moving averages
-      of the trainable variables from the checkpoint.
-    output_node_names: Output tensor names, defaults are: num_detections,
-      detection_scores, detection_boxes, detection_classes.
+    frozen_graph_path: Path to write inference graph.
+    frozen_graph_def: tf.GraphDef holding frozen graph.
   """
-  inference_graph_def = tf.get_default_graph().as_graph_def()
-  if checkpoint_path:
-    output_graph_def = get_frozen_graph_def(
-        inference_graph_def=inference_graph_def,
-        use_moving_averages=use_moving_averages,
-        input_checkpoint=checkpoint_path,
-        output_node_names=output_node_names,
-    )
-
-    with gfile.GFile(inference_graph_path, 'wb') as f:
-      f.write(output_graph_def.SerializeToString())
-    logging.info('%d ops in the final graph.', len(output_graph_def.node))
-
-    return
-  tf.train.write_graph(inference_graph_def,
-                       os.path.dirname(inference_graph_path),
-                       os.path.basename(inference_graph_path),
-                       as_text=False)
-
-
-def _write_saved_model(inference_graph_path, inputs, outputs,
-                       checkpoint_path=None, use_moving_averages=False):
+  with gfile.GFile(frozen_graph_path, 'wb') as f:
+    f.write(frozen_graph_def.SerializeToString())
+  logging.info('%d ops in the final graph.', len(frozen_graph_def.node))
+
+
+def _write_saved_model(saved_model_path,
+                       frozen_graph_def,
+                       inputs,
+                       outputs):
   """Writes SavedModel to disk.
 
   If checkpoint_path is not None bakes the weights into the graph thereby
@@ -247,30 +247,17 @@ def _write_saved_model(inference_graph_path, inputs, outputs,
   is restored.
 
   Args:
-    inference_graph_path: Path to write inference graph.
+    saved_model_path: Path to write SavedModel.
+    frozen_graph_def: tf.GraphDef holding frozen graph.
     inputs: The input image tensor to use for detection.
     outputs: A tensor dictionary containing the outputs of a DetectionModel.
-    checkpoint_path: Optional path to the checkpoint file.
-    use_moving_averages: Whether to export the original or the moving averages
-      of the trainable variables from the checkpoint.
   """
-  inference_graph_def = tf.get_default_graph().as_graph_def()
-  checkpoint_graph_def = None
-  if checkpoint_path:
-    output_node_names = ','.join(outputs.keys())
-    checkpoint_graph_def = get_frozen_graph_def(
-        inference_graph_def=inference_graph_def,
-        use_moving_averages=use_moving_averages,
-        input_checkpoint=checkpoint_path,
-        output_node_names=output_node_names
-    )
-
   with tf.Graph().as_default():
     with session.Session() as sess:
 
-      tf.import_graph_def(checkpoint_graph_def)
+      tf.import_graph_def(frozen_graph_def, name='')
 
-      builder = tf.saved_model.builder.SavedModelBuilder(inference_graph_path)
+      builder = tf.saved_model.builder.SavedModelBuilder(saved_model_path)
 
       tensor_info_inputs = {
           'inputs': tf.saved_model.utils.build_tensor_info(inputs)}
@@ -294,46 +281,96 @@ def _write_saved_model(inference_graph_path, inputs, outputs,
       builder.save()
 
 
+def _write_graph_and_checkpoint(inference_graph_def,
+                                model_path,
+                                input_saver_def,
+                                trained_checkpoint_prefix):
+  for node in inference_graph_def.node:
+    node.device = ''
+  with tf.Graph().as_default():
+    tf.import_graph_def(inference_graph_def, name='')
+    with session.Session() as sess:
+      saver = saver_lib.Saver(saver_def=input_saver_def,
+                              save_relative_paths=True)
+      saver.restore(sess, trained_checkpoint_prefix)
+      saver.save(sess, model_path)
+
+
 def _export_inference_graph(input_type,
                             detection_model,
                             use_moving_averages,
-                            checkpoint_path,
-                            inference_graph_path,
-                            export_as_saved_model=False):
+                            trained_checkpoint_prefix,
+                            output_directory,
+                            optimize_graph=False,
+                            output_collection_name='inference_op'):
   """Export helper."""
+  tf.gfile.MakeDirs(output_directory)
+  frozen_graph_path = os.path.join(output_directory,
+                                   'frozen_inference_graph.pb')
+  saved_model_path = os.path.join(output_directory, 'saved_model')
+  model_path = os.path.join(output_directory, 'model.ckpt')
+
   if input_type not in input_placeholder_fn_map:
     raise ValueError('Unknown input type: {}'.format(input_type))
-  inputs = tf.to_float(input_placeholder_fn_map[input_type]())
+  placeholder_tensor, input_tensors = input_placeholder_fn_map[input_type]()
+  inputs = tf.to_float(input_tensors)
   preprocessed_inputs = detection_model.preprocess(inputs)
   output_tensors = detection_model.predict(preprocessed_inputs)
   postprocessed_tensors = detection_model.postprocess(output_tensors)
-  outputs = _add_output_tensor_nodes(postprocessed_tensors)
-  out_node_names = list(outputs.keys())
-  if export_as_saved_model:
-    _write_saved_model(inference_graph_path, inputs, outputs, checkpoint_path,
-                       use_moving_averages)
+  outputs = _add_output_tensor_nodes(postprocessed_tensors,
+                                     output_collection_name)
+
+  saver = None
+  if use_moving_averages:
+    variable_averages = tf.train.ExponentialMovingAverage(0.0)
+    variables_to_restore = variable_averages.variables_to_restore()
+    saver = tf.train.Saver(variables_to_restore)
   else:
-    _write_inference_graph(inference_graph_path, checkpoint_path,
-                           use_moving_averages,
-                           output_node_names=','.join(out_node_names))
+    saver = tf.train.Saver()
+  input_saver_def = saver.as_saver_def()
+
+  _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)
+
+  frozen_graph_def = 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=','.join(outputs.keys()),
+      restore_op_name='save/restore_all',
+      filename_tensor_name='save/Const:0',
+      clear_devices=True,
+      optimize_graph=optimize_graph,
+      initializer_nodes='')
+  _write_frozen_graph(frozen_graph_path, frozen_graph_def)
+  _write_saved_model(saved_model_path, frozen_graph_def, placeholder_tensor,
+                     outputs)
 
 
-def export_inference_graph(input_type, pipeline_config, checkpoint_path,
-                           inference_graph_path, export_as_saved_model=False):
+def export_inference_graph(input_type,
+                           pipeline_config,
+                           trained_checkpoint_prefix,
+                           output_directory,
+                           optimize_graph=False,
+                           output_collection_name='inference_op'):
   """Exports inference graph for the model specified in the pipeline config.
 
   Args:
     input_type: Type of input for the graph. Can be one of [`image_tensor`,
       `tf_example`].
     pipeline_config: pipeline_pb2.TrainAndEvalPipelineConfig proto.
-    checkpoint_path: Path to the checkpoint file to freeze.
-    inference_graph_path: Path to write inference graph to.
-    export_as_saved_model: If the model should be exported as a SavedModel. If
-                           false, it is saved as an inference graph.
+    trained_checkpoint_prefix: Path to the trained checkpoint file.
+    output_directory: Path to write outputs.
+    optimize_graph: Whether to optimize graph using Grappler.
+    output_collection_name: Name of collection to add output tensors to.
+      If None, does not add output tensors to a collection.
   """
   detection_model = model_builder.build(pipeline_config.model,
                                         is_training=False)
   _export_inference_graph(input_type, detection_model,
                           pipeline_config.eval_config.use_moving_averages,
-                          checkpoint_path, inference_graph_path,
-                          export_as_saved_model)
+                          trained_checkpoint_prefix, output_directory,
+                          optimize_graph, output_collection_name)

object_detection/exporter_test.py

@@ -43,18 +43,22 @@ class FakeModel(model.DetectionModel):
   def postprocess(self, prediction_dict):
     with tf.control_dependencies(prediction_dict.values()):
       postprocessed_tensors = {
-          '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)
+          '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),
+          'detection_classes': tf.constant([[0, 1],
+                                            [1, 0]], tf.float32),
+          'num_detections': tf.constant([2, 1], tf.float32)
       }
       if self._add_detection_masks:
         postprocessed_tensors['detection_masks'] = tf.constant(
-            np.arange(32).reshape([2, 4, 4]), tf.float32)
+            np.arange(64).reshape([2, 2, 4, 4]), tf.float32)
     return postprocessed_tensors
 
-  def restore_fn(self, checkpoint_path, from_detection_checkpoint):
+  def restore_map(self, checkpoint_path, from_detection_checkpoint):
     pass
 
   def loss(self, prediction_dict):
@@ -69,7 +73,7 @@ class ExportInferenceGraphTest(tf.test.TestCase):
     with g.as_default():
       mock_model = FakeModel()
       preprocessed_inputs = mock_model.preprocess(
-          tf.ones([1, 3, 4, 3], tf.float32))
+          tf.placeholder(tf.float32, shape=[None, None, None, 3]))
       predictions = mock_model.predict(preprocessed_inputs)
       mock_model.postprocess(predictions)
       if use_moving_averages:
@@ -103,71 +107,62 @@ class ExportInferenceGraphTest(tf.test.TestCase):
     return example
 
   def test_export_graph_with_image_tensor_input(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()
-      inference_graph_path = os.path.join(self.get_temp_dir(),
-                                          'exported_graph.pbtxt')
-
+      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,
-          checkpoint_path=None,
-          inference_graph_path=inference_graph_path)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
   def test_export_graph_with_tf_example_input(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()
-      inference_graph_path = os.path.join(self.get_temp_dir(),
-                                          'exported_graph.pbtxt')
+      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='tf_example',
           pipeline_config=pipeline_config,
-          checkpoint_path=None,
-          inference_graph_path=inference_graph_path)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
   def test_export_graph_with_encoded_image_string_input(self):
-    with mock.patch.object(
-        model_builder, 'build', autospec=True) as mock_builder:
-      mock_builder.return_value = FakeModel()
-      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='encoded_image_string_tensor',
-          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,
+    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)
-    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()
+      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',
+          input_type='encoded_image_string_tensor',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
-  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,
+  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,
                                           use_moving_averages=True)
-    inference_graph_path = os.path.join(self.get_temp_dir(),
-                                        'exported_graph.pb')
+    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()
@@ -176,15 +171,17 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='image_tensor',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
   def test_export_model_with_all_output_nodes(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')
+    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_masks=True)
@@ -192,8 +189,8 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='image_tensor',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          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')
@@ -204,11 +201,13 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       inference_graph.get_tensor_by_name('num_detections:0')
 
   def test_export_model_with_detection_only_nodes(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')
+    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_masks=False)
@@ -216,8 +215,8 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='image_tensor',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          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')
@@ -229,11 +228,13 @@ class ExportInferenceGraphTest(tf.test.TestCase):
         inference_graph.get_tensor_by_name('detection_masks:0')
 
   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')
+    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_masks=True)
@@ -242,8 +243,8 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='image_tensor',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          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) as sess:
@@ -253,15 +254,19 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       classes = inference_graph.get_tensor_by_name('detection_classes:0')
       masks = inference_graph.get_tensor_by_name('detection_masks:0')
       num_detections = inference_graph.get_tensor_by_name('num_detections:0')
-      (boxes, scores, classes, masks, num_detections) = sess.run(
+      (boxes_np, scores_np, classes_np, masks_np, num_detections_np) = sess.run(
           [boxes, scores, classes, masks, 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(masks, np.arange(32).reshape([2, 4, 4]))
-      self.assertAllClose(num_detections, [2])
+          feed_dict={image_tensor: np.ones((2, 4, 4, 3)).astype(np.uint8)})
+      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]])
+      self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
+      self.assertAllClose(num_detections_np, [2, 1])
 
   def _create_encoded_image_string(self, image_array_np, encoding_format):
     od_graph = tf.Graph()
@@ -276,11 +281,13 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       return encoded_string.eval()
 
   def test_export_and_run_inference_with_encoded_image_string_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')
+    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_masks=True)
@@ -289,8 +296,8 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='encoded_image_string_tensor',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
     inference_graph = self._load_inference_graph(inference_graph_path)
     jpg_image_str = self._create_encoded_image_string(
@@ -306,23 +313,69 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       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]:
+        image_str_batch_np = np.hstack([image_str]* 2)
         (boxes_np, scores_np, classes_np, masks_np,
          num_detections_np) = sess.run(
              [boxes, scores, classes, masks, num_detections],
-             feed_dict={image_str_tensor: image_str})
-        self.assertAllClose(boxes_np, [[0.0, 0.0, 0.5, 0.5],
-                                       [0.5, 0.5, 0.8, 0.8]])
-        self.assertAllClose(scores_np, [[0.7, 0.6]])
-        self.assertAllClose(classes_np, [[1, 2]])
-        self.assertAllClose(masks_np, np.arange(32).reshape([2, 4, 4]))
-        self.assertAllClose(num_detections_np, [2])
+             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]])
+        self.assertAllClose(classes_np, [[1, 2],
+                                         [2, 1]])
+        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:
+      mock_builder.return_value = FakeModel(add_detection_masks=True)
+      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')
+      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,
+                                   '^TensorArray has inconsistent shapes.'):
+        sess.run([boxes, scores, classes, masks, num_detections],
+                 feed_dict={image_str_tensor: image_str_batch_np})
 
   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')
+    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_masks=True)
@@ -331,10 +384,12 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='tf_example',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
     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')
@@ -342,23 +397,27 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       classes = inference_graph.get_tensor_by_name('detection_classes:0')
       masks = inference_graph.get_tensor_by_name('detection_masks:0')
       num_detections = inference_graph.get_tensor_by_name('num_detections:0')
-      (boxes, scores, classes, masks, num_detections) = sess.run(
+      (boxes_np, scores_np, classes_np, masks_np, num_detections_np) = sess.run(
           [boxes, scores, classes, masks, 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(masks, np.arange(32).reshape([2, 4, 4]))
-      self.assertAllClose(num_detections, [2])
+          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]])
+      self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
+      self.assertAllClose(num_detections_np, [2, 1])
 
   def test_export_saved_model_and_run_inference(self):
-    checkpoint_path = os.path.join(self.get_temp_dir(), 'model-ckpt')
-    self._save_checkpoint_from_mock_model(checkpoint_path,
+    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)
-    inference_graph_path = os.path.join(self.get_temp_dir(),
-                                        'saved_model')
+    output_directory = os.path.join(tmp_dir, 'output')
+    saved_model_path = os.path.join(output_directory, 'saved_model')
 
     with mock.patch.object(
         model_builder, 'build', autospec=True) as mock_builder:
@@ -368,30 +427,84 @@ class ExportInferenceGraphTest(tf.test.TestCase):
       exporter.export_inference_graph(
           input_type='tf_example',
           pipeline_config=pipeline_config,
-          checkpoint_path=checkpoint_path,
-          inference_graph_path=inference_graph_path,
-          export_as_saved_model=True)
+          trained_checkpoint_prefix=trained_checkpoint_prefix,
+          output_directory=output_directory)
 
+    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:
         tf.saved_model.loader.load(
-            sess, [tf.saved_model.tag_constants.SERVING], inference_graph_path)
-        tf_example = od_graph.get_tensor_by_name('import/tf_example:0')
-        boxes = od_graph.get_tensor_by_name('import/detection_boxes:0')
-        scores = od_graph.get_tensor_by_name('import/detection_scores:0')
-        classes = od_graph.get_tensor_by_name('import/detection_classes:0')
-        masks = od_graph.get_tensor_by_name('import/detection_masks:0')
-        num_detections = od_graph.get_tensor_by_name('import/num_detections:0')
-        (boxes, scores, classes, masks, num_detections) = sess.run(
-            [boxes, scores, classes, masks, 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(masks, np.arange(32).reshape([2, 4, 4]))
-        self.assertAllClose(num_detections, [2])
+            sess, [tf.saved_model.tag_constants.SERVING], saved_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')
+        masks = od_graph.get_tensor_by_name('detection_masks:0')
+        num_detections = od_graph.get_tensor_by_name('num_detections:0')
+        (boxes_np, scores_np, classes_np, masks_np,
+         num_detections_np) = sess.run(
+             [boxes, scores, classes, masks, num_detections],
+             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]])
+        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
+        self.assertAllClose(num_detections_np, [2, 1])
+
+  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:
+      mock_builder.return_value = FakeModel(add_detection_masks=True)
+      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)
+
+    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')
+        classes = od_graph.get_tensor_by_name('detection_classes:0')
+        masks = od_graph.get_tensor_by_name('detection_masks:0')
+        num_detections = od_graph.get_tensor_by_name('num_detections:0')
+        (boxes_np, scores_np, classes_np, masks_np,
+         num_detections_np) = sess.run(
+             [boxes, scores, classes, masks, num_detections],
+             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]])
+        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
+        self.assertAllClose(num_detections_np, [2, 1])
+
 
 if __name__ == '__main__':
   tf.test.main()

object_detection/g3doc/configuring_jobs.md

@@ -157,6 +157,6 @@ number of workers, gpu type).
 
 ## Configuring the Evaluator
 
-Currently evaluation is fixed to generating metrics as defined by the PASCAL
-VOC challenge. The parameters for `eval_config` are set to reasonable defaults
-and typically do not need to be configured.
+Currently evaluation is fixed to generating metrics as defined by the PASCAL VOC
+challenge. The parameters for `eval_config` are set to reasonable defaults and
+typically do not need to be configured.

object_detection/g3doc/installation.md

@@ -74,6 +74,6 @@ to avoid running this manually, you can add it as a new line to the end of your
 You can test that you have correctly installed the Tensorflow Object Detection\
 API by running the following command:
 
-``` bash
+```bash
 python object_detection/builders/model_builder_test.py
 ```

object_detection/g3doc/preparing_inputs.md

@@ -7,39 +7,51 @@ TFRecords.
 
 ## Generating the PASCAL VOC TFRecord files.
 
-The raw 2012 PASCAL VOC data set can be downloaded
+The raw 2012 PASCAL VOC data set is located
 [here](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar).
-Extract the tar file and run the `create_pascal_tf_record` script:
+To download, extract and convert it to TFRecords, run the following commands
+below:
 
 ```bash
-# From tensorflow/models/object_detection
+# From tensorflow/models
+wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
 tar -xvf VOCtrainval_11-May-2012.tar
-python create_pascal_tf_record.py --data_dir=VOCdevkit \
-    --year=VOC2012 --set=train --output_path=pascal_train.record
-python create_pascal_tf_record.py --data_dir=VOCdevkit \
-    --year=VOC2012 --set=val --output_path=pascal_val.record
+python object_detection/create_pascal_tf_record.py \
+    --label_map_path=object_detection/data/pascal_label_map.pbtxt \
+    --data_dir=VOCdevkit --year=VOC2012 --set=train \
+    --output_path=pascal_train.record
+python object_detection/create_pascal_tf_record.py \
+    --label_map_path=object_detection/data/pascal_label_map.pbtxt \
+    --data_dir=VOCdevkit --year=VOC2012 --set=val \
+    --output_path=pascal_val.record
 ```
 
 You should end up with two TFRecord files named `pascal_train.record` and
-`pascal_val.record` in the `tensorflow/models/object_detection` directory.
+`pascal_val.record` in the `tensorflow/models` directory.
 
 The label map for the PASCAL VOC data set can be found at
-`data/pascal_label_map.pbtxt`.
+`object_detection/data/pascal_label_map.pbtxt`.
 
 ## Generating the Oxford-IIIT Pet TFRecord files.
 
-The Oxford-IIIT Pet data set can be downloaded from
-[their website](http://www.robots.ox.ac.uk/~vgg/data/pets/). Extract the tar
-file and run the `create_pet_tf_record` script to generate TFRecords.
+The Oxford-IIIT Pet data set is located
+[here](http://www.robots.ox.ac.uk/~vgg/data/pets/). To download, extract and
+convert it to TFRecrods, run the following commands below:
 
 ```bash
-# From tensorflow/models/object_detection
+# From tensorflow/models
+wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz
+wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz
 tar -xvf annotations.tar.gz
 tar -xvf images.tar.gz
-python create_pet_tf_record.py --data_dir=`pwd` --output_dir=`pwd`
+python object_detection/create_pet_tf_record.py \
+    --label_map_path=object_detection/data/pet_label_map.pbtxt \
+    --data_dir=`pwd` \
+    --output_dir=`pwd`
 ```
 
 You should end up with two TFRecord files named `pet_train.record` and
-`pet_val.record` in the `tensorflow/models/object_detection` directory.
+`pet_val.record` in the `tensorflow/models` directory.
 
-The label map for the Pet dataset can be found at `data/pet_label_map.pbtxt`.
+The label map for the Pet dataset can be found at
+`object_detection/data/pet_label_map.pbtxt`.

object_detection/g3doc/running_locally.md

@@ -77,5 +77,5 @@ tensorboard --logdir=${PATH_TO_MODEL_DIRECTORY}
 ```
 
 where `${PATH_TO_MODEL_DIRECTORY}` points to the directory that contains the
-train and eval directories. Please note it make take Tensorboard a couple
-minutes to populate with data.
+train and eval directories. Please note it may take Tensorboard a couple minutes
+to populate with data.