1

Custom_Mask_RCNN/create_pet_tf_record.py

@@ -46,7 +46,7 @@
 flags.DEFINE_string('output_dir', '', 'Path to directory to output TFRecords.')
 flags.DEFINE_string('label_map_path', 'data/pet_label_map.pbtxt',
                     'Path to label map proto')
-flags.DEFINE_boolean('faces_only', True, 'If True, generates bounding boxes '
+flags.DEFINE_boolean('faces_only', False, 'If True, generates bounding boxes '
                      'for pet faces.  Otherwise generates bounding boxes (as '
                      'well as segmentations for full pet bodies).  Note that '
                      'in the latter case, the resulting files are much larger.')
@@ -65,7 +65,7 @@ def get_class_name_from_filename(file_name):
   Returns:
     A string of the class name.
   """
-  match = re.match(r'([A-Za-z_]+)(_[0-9]+\.jpg)', file_name, re.I)
+  match = re.match(r'([A-Za-z_]+)([0-9]+\.jpg)', file_name, re.I)
   return match.groups()[0]
 
 
@@ -118,10 +118,7 @@ def dict_to_tf_example(data,
     raise ValueError('Mask format not PNG')
 
   mask_np = np.asarray(mask)
-  nonbackground_indices_x = np.any(mask_np != 2, axis=0)
-  nonbackground_indices_y = np.any(mask_np != 2, axis=1)
-  nonzero_x_indices = np.where(nonbackground_indices_x)
-  nonzero_y_indices = np.where(nonbackground_indices_y)
+  nonzero = np.where(mask_np != 0)
 
   width = int(data['size']['width'])
   height = int(data['size']['height'])
@@ -148,10 +145,10 @@ def dict_to_tf_example(data,
       ymin = float(obj['bndbox']['ymin'])
       ymax = float(obj['bndbox']['ymax'])
     else:
-      xmin = float(np.min(nonzero_x_indices))
-      xmax = float(np.max(nonzero_x_indices))
-      ymin = float(np.min(nonzero_y_indices))
-      ymax = float(np.max(nonzero_y_indices))
+      xmin = float(np.min(nonzero[1]))
+      xmax = float(np.max(nonzero[1]))
+      ymin = float(np.min(nonzero[0]))
+      ymax = float(np.max(nonzero[0]))
 
     xmins.append(xmin / width)
     ymins.append(ymin / height)
@@ -163,7 +160,7 @@ def dict_to_tf_example(data,
     truncated.append(int(obj['truncated']))
     poses.append(obj['pose'].encode('utf8'))
     if not faces_only:
-      mask_remapped = (mask_np != 2).astype(np.uint8)
+      mask_remapped = (mask_np != 0).astype(np.uint8)
       masks.append(mask_remapped)
 
   feature_dict = {
@@ -231,7 +228,7 @@ def create_tf_record(output_filename,
     if idx % 100 == 0:
       logging.info('On image %d of %d', idx, len(examples))
     xml_path = os.path.join(annotations_dir, 'xmls', example + '.xml')
-    mask_path = os.path.join(annotations_dir, 'trimaps', example + '.png')
+    mask_path = os.path.join(annotations_dir, 'masks', example + '.png')
 
     if not os.path.exists(xml_path):
       logging.warning('Could not find %s, ignoring example.', xml_path)
@@ -259,7 +256,8 @@ def create_tf_record(output_filename,
 # TODO: Add test for pet/PASCAL main files.
 def main(_):
   data_dir = FLAGS.data_dir
-  label_map_dict = label_map_util.get_label_map_dict(FLAGS.label_map_path)
+  #label_map_dict = label_map_util.get_label_map_dict(FLAGS.label_map_path)
+  label_map_dict = label_map_util.get_label_map_dict('./label_map.pbtxt')
 
   logging.info('Reading from Pet dataset.')
   image_dir = os.path.join(data_dir, 'images')

Custom_Mask_RCNN/label_map.pbtxt

@@ -0,0 +1,4 @@
+item {
+  id: 1
+  name: 'toy'
+}

Custom_Mask_RCNN/test.py

@@ -0,0 +1,156 @@
+import numpy as np
+import os
+import six.moves.urllib as urllib
+import sys
+import tarfile
+import tensorflow as tf
+import zipfile
+import datetime
+#import cv2
+
+from collections import defaultdict
+from io import StringIO
+
+import matplotlib
+matplotlib.use('Agg')
+
+from matplotlib import pyplot as plt
+from PIL import Image
+
+
+# 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
+
+#if tf.__version__ < '1.4.0':
+#      raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!')
+
+# This is needed to display the images.
+#%matplotlib inline
+
+from utils import label_map_util
+
+from utils import visualization_utils as vis_util
+
+# What model to download.
+MODEL_NAME = 'out4'
+#MODEL_NAME = 'mask_rcnn_inception_v2_coco_2018_01_28'
+#MODEL_FILE = MODEL_NAME + '.tar.gz'
+
+# Path to frozen detection graph. This is the actual model that is used for the object detection.
+PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
+
+# List of the strings that is used to add correct label for each box.
+#PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')
+PATH_TO_LABELS = 'label_map.pbtxt'
+NUM_CLASSES = 1
+
+detection_graph = tf.Graph()
+with detection_graph.as_default():
+  od_graph_def = tf.GraphDef()
+  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
+    serialized_graph = fid.read()
+    od_graph_def.ParseFromString(serialized_graph)
+    tf.import_graph_def(od_graph_def, name='')
+
+label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
+categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
+category_index = label_map_util.create_category_index(categories)
+
+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:
+# image1.jpg
+# image2.jpg
+# If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS.
+#PATH_TO_TEST_IMAGES_DIR = 'test_images'
+#TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 7) ]
+img_path='images'
+TEST_IMAGE_PATHS = os.listdir(img_path)
+os.chdir(img_path)
+
+# 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[0], image.shape[1])
+        detection_masks_reframed = tf.cast(
+            tf.greater(detection_masks_reframed, 0.1), 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
+      print("start")
+      b = datetime.datetime.now()
+      output_dict = sess.run(tensor_dict,
+                             feed_dict={image_tensor: np.expand_dims(image, 0)})
+
+      e = datetime.datetime.now()
+      k = e - b
+      print("%f ms" % (k.total_seconds()*1000))
+      # 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.uint8)
+      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)
+image_path='/home/xiaoqife/models/models-master/research/object_detection/images/toy6.jpg'
+image = Image.open(image_path)
+
+print(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, detection_graph)
+# Visualization of the results of a detection.
+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'),
+  use_normalized_coordinates=True,
+  line_thickness=8)
+plt.figure(figsize=IMAGE_SIZE)
+filename = os.path.split(image_path)[1]
+plt.imshow(image_np)
+plt.show()
+plt.savefig(filename + '.jpg')