add inference tools for Open Image dataset.

research/object_detection/inference/BUILD

@@ -0,0 +1,40 @@
+# Tensorflow Object Detection API: main runnables.
+
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])
+
+# Apache 2.0
+
+py_library(
+    name = "detection_inference",
+    srcs = ["detection_inference.py"],
+    deps = [
+        "//tensorflow",
+        "//tensorflow_models/object_detection/core:standard_fields",
+    ],
+)
+
+py_test(
+    name = "detection_inference_test",
+    srcs = ["detection_inference_test.py"],
+    deps = [
+        ":detection_inference",
+        "//third_party/py/PIL:pil",
+        "//third_party/py/numpy",
+        "//tensorflow",
+        "//tensorflow_models/object_detection/core:standard_fields",
+        "//tensorflow_models/object_detection/utils:dataset_util",
+    ],
+)
+
+py_binary(
+    name = "infer_detections",
+    srcs = ["infer_detections.py"],
+    deps = [
+        ":detection_inference",
+        "//tensorflow",
+    ],
+)

research/object_detection/inference/detection_inference.py

@@ -0,0 +1,141 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions for detection inference."""
+from __future__ import division
+
+import tensorflow as tf
+
+from object_detection.core import standard_fields
+
+
+def build_input(tfrecord_paths):
+  """Builds the graph's input.
+
+  Args:
+    tfrecord_paths: List of paths to the input TFRecords
+
+  Returns:
+    serialized_example_tensor: The next serialized example. String scalar Tensor
+    image_tensor: The decoded image of the example. Uint8 tensor,
+        shape=[1, None, None,3]
+  """
+  filename_queue = tf.train.string_input_producer(
+      tfrecord_paths, shuffle=False, num_epochs=1)
+
+  tf_record_reader = tf.TFRecordReader()
+  _, serialized_example_tensor = tf_record_reader.read(filename_queue)
+  features = tf.parse_single_example(
+      serialized_example_tensor,
+      features={
+          standard_fields.TfExampleFields.image_encoded:
+              tf.FixedLenFeature([], tf.string),
+      })
+  encoded_image = features[standard_fields.TfExampleFields.image_encoded]
+  image_tensor = tf.image.decode_image(encoded_image, channels=3)
+  image_tensor.set_shape([None, None, 3])
+  image_tensor = tf.expand_dims(image_tensor, 0)
+
+  return serialized_example_tensor, image_tensor
+
+
+def build_inference_graph(image_tensor, inference_graph_path):
+  """Loads the inference graph and connects it to the input image.
+
+  Args:
+    image_tensor: The input image. uint8 tensor, shape=[1, None, None, 3]
+    inference_graph_path: Path to the inference graph with embedded weights
+
+  Returns:
+    detected_boxes_tensor: Detected boxes. Float tensor,
+        shape=[num_detections, 4]
+    detected_scores_tensor: Detected scores. Float tensor,
+        shape=[num_detections]
+    detected_labels_tensor: Detected labels. Int64 tensor,
+        shape=[num_detections]
+  """
+  with tf.gfile.Open(inference_graph_path, 'r') as graph_def_file:
+    graph_content = graph_def_file.read()
+  graph_def = tf.GraphDef()
+  graph_def.MergeFromString(graph_content)
+
+  tf.import_graph_def(
+      graph_def, name='', input_map={'image_tensor': image_tensor})
+
+  g = tf.get_default_graph()
+
+  num_detections_tensor = tf.squeeze(
+      g.get_tensor_by_name('num_detections:0'), 0)
+  num_detections_tensor = tf.cast(num_detections_tensor, tf.int32)
+
+  detected_boxes_tensor = tf.squeeze(
+      g.get_tensor_by_name('detection_boxes:0'), 0)
+  detected_boxes_tensor = detected_boxes_tensor[:num_detections_tensor]
+
+  detected_scores_tensor = tf.squeeze(
+      g.get_tensor_by_name('detection_scores:0'), 0)
+  detected_scores_tensor = detected_scores_tensor[:num_detections_tensor]
+
+  detected_labels_tensor = tf.squeeze(
+      g.get_tensor_by_name('detection_classes:0'), 0)
+  detected_labels_tensor = tf.cast(detected_labels_tensor, tf.int64)
+  detected_labels_tensor = detected_labels_tensor[:num_detections_tensor]
+
+  return detected_boxes_tensor, detected_scores_tensor, detected_labels_tensor
+
+
+def infer_detections_and_add_to_example(
+    serialized_example_tensor, detected_boxes_tensor, detected_scores_tensor,
+    detected_labels_tensor, discard_image_pixels):
+  """Runs the supplied tensors and adds the inferred detections to the example.
+
+  Args:
+    serialized_example_tensor: Serialized TF example. Scalar string tensor
+    detected_boxes_tensor: Detected boxes. Float tensor,
+        shape=[num_detections, 4]
+    detected_scores_tensor: Detected scores. Float tensor,
+        shape=[num_detections]
+    detected_labels_tensor: Detected labels. Int64 tensor,
+        shape=[num_detections]
+    discard_image_pixels: If true, discards the image from the result
+  Returns:
+    The de-serialized TF example augmented with the inferred detections.
+  """
+  tf_example = tf.train.Example()
+  (serialized_example, detected_boxes, detected_scores,
+   detected_classes) = tf.get_default_session().run([
+       serialized_example_tensor, detected_boxes_tensor, detected_scores_tensor,
+       detected_labels_tensor
+   ])
+  detected_boxes = detected_boxes.T
+
+  tf_example.ParseFromString(serialized_example)
+  feature = tf_example.features.feature
+  feature[standard_fields.TfExampleFields.
+          detection_score].float_list.value[:] = detected_scores
+  feature[standard_fields.TfExampleFields.
+          detection_bbox_ymin].float_list.value[:] = detected_boxes[0]
+  feature[standard_fields.TfExampleFields.
+          detection_bbox_xmin].float_list.value[:] = detected_boxes[1]
+  feature[standard_fields.TfExampleFields.
+          detection_bbox_ymax].float_list.value[:] = detected_boxes[2]
+  feature[standard_fields.TfExampleFields.
+          detection_bbox_xmax].float_list.value[:] = detected_boxes[3]
+  feature[standard_fields.TfExampleFields.
+          detection_class_label].int64_list.value[:] = detected_classes
+
+  if discard_image_pixels:
+    del feature[standard_fields.TfExampleFields.image_encoded]
+
+  return tf_example

research/object_detection/inference/detection_inference_test.py

@@ -0,0 +1,176 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+r"""Tests for detection_inference.py."""
+
+import os
+import StringIO
+
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+
+from object_detection.core import standard_fields
+from object_detection.inference import detection_inference
+from object_detection.utils import dataset_util
+
+
+def get_mock_tfrecord_path():
+  return os.path.join(tf.test.get_temp_dir(), 'mock.tfrec')
+
+
+def create_mock_tfrecord():
+  pil_image = Image.fromarray(np.array([[[123, 0, 0]]], dtype=np.uint8), 'RGB')
+  image_output_stream = StringIO.StringIO()
+  pil_image.save(image_output_stream, format='png')
+  encoded_image = image_output_stream.getvalue()
+
+  feature_map = {
+      'test_field':
+          dataset_util.float_list_feature([1, 2, 3, 4]),
+      standard_fields.TfExampleFields.image_encoded:
+          dataset_util.bytes_feature(encoded_image),
+  }
+
+  tf_example = tf.train.Example(features=tf.train.Features(feature=feature_map))
+  with tf.python_io.TFRecordWriter(get_mock_tfrecord_path()) as writer:
+    writer.write(tf_example.SerializeToString())
+
+
+def get_mock_graph_path():
+  return os.path.join(tf.test.get_temp_dir(), 'mock_graph.pb')
+
+
+def create_mock_graph():
+  g = tf.Graph()
+  with g.as_default():
+    in_image_tensor = tf.placeholder(
+        tf.uint8, shape=[1, None, None, 3], name='image_tensor')
+    tf.constant([2.0], name='num_detections')
+    tf.constant(
+        [[[0, 0.8, 0.7, 1], [0.1, 0.2, 0.8, 0.9], [0.2, 0.3, 0.4, 0.5]]],
+        name='detection_boxes')
+    tf.constant([[0.1, 0.2, 0.3]], name='detection_scores')
+    tf.identity(
+        tf.constant([[1.0, 2.0, 3.0]]) *
+        tf.reduce_sum(tf.cast(in_image_tensor, dtype=tf.float32)),
+        name='detection_classes')
+    graph_def = g.as_graph_def()
+
+  with tf.gfile.Open(get_mock_graph_path(), 'w') as fl:
+    fl.write(graph_def.SerializeToString())
+
+
+class InferDetectionsTests(tf.test.TestCase):
+
+  def test_simple(self):
+    create_mock_graph()
+    create_mock_tfrecord()
+
+    serialized_example_tensor, image_tensor = detection_inference.build_input(
+        [get_mock_tfrecord_path()])
+    self.assertAllEqual(image_tensor.get_shape().as_list(), [1, None, None, 3])
+
+    (detected_boxes_tensor, detected_scores_tensor,
+     detected_labels_tensor) = detection_inference.build_inference_graph(
+         image_tensor, get_mock_graph_path())
+
+    with self.test_session(use_gpu=False) as sess:
+      sess.run(tf.global_variables_initializer())
+      sess.run(tf.local_variables_initializer())
+      tf.train.start_queue_runners()
+
+      tf_example = detection_inference.infer_detections_and_add_to_example(
+          serialized_example_tensor, detected_boxes_tensor,
+          detected_scores_tensor, detected_labels_tensor, False)
+
+    self.assertProtoEquals(r"""
+        features {
+          feature {
+            key: "image/detection/bbox/ymin"
+            value { float_list { value: [0.0, 0.1] } } }
+          feature {
+            key: "image/detection/bbox/xmin"
+            value { float_list { value: [0.8, 0.2] } } }
+          feature {
+            key: "image/detection/bbox/ymax"
+            value { float_list { value: [0.7, 0.8] } } }
+          feature {
+            key: "image/detection/bbox/xmax"
+            value { float_list { value: [1.0, 0.9] } } }
+          feature {
+            key: "image/detection/label"
+            value { int64_list { value: [123, 246] } } }
+          feature {
+            key: "image/detection/score"
+            value { float_list { value: [0.1, 0.2] } } }
+          feature {
+            key: "image/encoded"
+            value { bytes_list { value:
+              "\211PNG\r\n\032\n\000\000\000\rIHDR\000\000\000\001\000\000"
+              "\000\001\010\002\000\000\000\220wS\336\000\000\000\022IDATx"
+              "\234b\250f`\000\000\000\000\377\377\003\000\001u\000|gO\242"
+              "\213\000\000\000\000IEND\256B`\202" } } }
+          feature {
+            key: "test_field"
+            value { float_list { value: [1.0, 2.0, 3.0, 4.0] } } } }
+    """, tf_example)
+
+  def test_discard_image(self):
+    create_mock_graph()
+    create_mock_tfrecord()
+
+    serialized_example_tensor, image_tensor = detection_inference.build_input(
+        [get_mock_tfrecord_path()])
+    (detected_boxes_tensor, detected_scores_tensor,
+     detected_labels_tensor) = detection_inference.build_inference_graph(
+         image_tensor, get_mock_graph_path())
+
+    with self.test_session(use_gpu=False) as sess:
+      sess.run(tf.global_variables_initializer())
+      sess.run(tf.local_variables_initializer())
+      tf.train.start_queue_runners()
+
+      tf_example = detection_inference.infer_detections_and_add_to_example(
+          serialized_example_tensor, detected_boxes_tensor,
+          detected_scores_tensor, detected_labels_tensor, True)
+
+    self.assertProtoEquals(r"""
+        features {
+          feature {
+            key: "image/detection/bbox/ymin"
+            value { float_list { value: [0.0, 0.1] } } }
+          feature {
+            key: "image/detection/bbox/xmin"
+            value { float_list { value: [0.8, 0.2] } } }
+          feature {
+            key: "image/detection/bbox/ymax"
+            value { float_list { value: [0.7, 0.8] } } }
+          feature {
+            key: "image/detection/bbox/xmax"
+            value { float_list { value: [1.0, 0.9] } } }
+          feature {
+            key: "image/detection/label"
+            value { int64_list { value: [123, 246] } } }
+          feature {
+            key: "image/detection/score"
+            value { float_list { value: [0.1, 0.2] } } }
+          feature {
+            key: "test_field"
+            value { float_list { value: [1.0, 2.0, 3.0, 4.0] } } } }
+    """, tf_example)
+
+
+if __name__ == '__main__':
+  tf.test.main()