Adding source code for detection

Author: Tutj

corgi_detection.py

@@ -0,0 +1,115 @@
+from object_detection.utils import label_map_util
+from object_detection.utils import visualization_utils as vis_util
+from matplotlib import pyplot as plt
+from PIL import Image
+
+import numpy as np
+import os
+import sys
+import tensorflow as tf
+import matplotlib
+matplotlib.use('MacOSX') # Change to relevant OS, check matplotlib backend
+
+# This is needed since the notebook is stored in the object_detection folder.
+sys.path.append("..")
+from object_detection.utils import ops as utils_ops
+
+
+# Path to labels
+MODEL_NAME = '/Users/hluong/PycharmProjects/CorgiRecognition/models/research/object_detection/inference_graph_corgi'
+PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'
+PATH_TO_LABELS = '/Users/hluong/PycharmProjects/CorgiRecognition/models/research/object_detection/corgi_training/corgi.pbtxt'
+
+# Initialize a tensorflow graph
+detection_graph = tf.Graph()
+with detection_graph.as_default():
+    od_graph_def = tf.GraphDef()
+    with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
+        serialized_graph = fid.read()
+        od_graph_def.ParseFromString(serialized_graph)
+        tf.import_graph_def(od_graph_def, name='')
+category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
+
+
+# Load images to a numpy array
+def load_image_into_numpy_array(image):
+    (im_width, im_height) = image.size
+    return np.array(image.getdata()).reshape(
+        (im_height, im_width, 3)).astype(np.uint8)
+
+
+# For the sake of simplicity we will use only 2 images:
+PATH_TO_TEST_IMAGES_DIR = '/Users/hluong/PycharmProjects/CorgiRecognition/test_corgi_images'
+TEST_IMAGE_PATHS = [os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 3)]
+
+# Size, in inches, of the output images.
+IMAGE_SIZE = (12, 8)
+
+
+def run_inference_for_single_image(image, graph):
+    with graph.as_default():
+        with tf.Session() as sess:
+            # Get handles to input and output tensors
+            ops = tf.get_default_graph().get_operations()
+            all_tensor_names = {output.name for op in ops for output in op.outputs}
+            tensor_dict = {}
+            for key in [
+              'num_detections', 'detection_boxes', 'detection_scores',
+              'detection_classes', 'detection_masks'
+            ]:
+                tensor_name = key + ':0'
+                if tensor_name in all_tensor_names:
+                  tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
+                      tensor_name)
+            if 'detection_masks' in tensor_dict:
+                # The following processing is only for single image
+                detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
+                detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
+                # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
+                real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
+                detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
+                detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
+                detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
+                    detection_masks, detection_boxes, image.shape[1], image.shape[2])
+                detection_masks_reframed = tf.cast(
+                    tf.greater(detection_masks_reframed, 0.5), tf.uint8)
+                # Follow the convention by adding back the batch dimension
+                tensor_dict['detection_masks'] = tf.expand_dims(
+                    detection_masks_reframed, 0)
+            image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
+            # Run inference
+            output_dict = sess.run(tensor_dict, feed_dict={image_tensor: image})
+
+            # all outputs are float32 numpy arrays, so convert types as appropriate
+            output_dict['num_detections'] = int(output_dict['num_detections'][0])
+            output_dict['detection_classes'] = output_dict[
+              'detection_classes'][0].astype(np.int64)
+            output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
+            output_dict['detection_scores'] = output_dict['detection_scores'][0]
+            if 'detection_masks' in output_dict:
+                output_dict['detection_masks'] = output_dict['detection_masks'][0]
+    return output_dict
+
+
+for image_path in TEST_IMAGE_PATHS:
+    image = Image.open(image_path)
+    # the array based representation of the image will be used later in order to prepare the
+    # result image with boxes and labels on it.
+    image_np = load_image_into_numpy_array(image)
+    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
+    image_np_expanded = np.expand_dims(image_np, axis=0)
+    # Actual detection.
+    output_dict = run_inference_for_single_image(image_np_expanded, detection_graph)
+    # Visualization of the results of a detection.
+    image_np = vis_util.visualize_boxes_and_labels_on_image_array(
+          image_np,
+          output_dict['detection_boxes'],
+          output_dict['detection_classes'],
+          output_dict['detection_scores'],
+          category_index,
+          instance_masks=output_dict.get('detection_masks_reframed', None),
+          use_normalized_coordinates=True,
+          line_thickness=8)
+    plt.figure(figsize=IMAGE_SIZE)
+    plt.imshow(image_np)
+    plt.show()