Updated python examples

ryouchinsa · Aug 26, 2018 · 5b04909 · 5b04909
1 parent 34f6564
commit 5b04909
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 12 deletions.
diff --git a/rectlabel_create_coco_tf_record.py b/rectlabel_create_coco_tf_record.py
@@ -1,11 +1,11 @@
 r"""
 python object_detection/dataset_tools/rectlabel_create_coco_tf_record.py \
-  --train_image_dir="${TRAIN_IMAGE_DIR}" \
-  --val_image_dir="${VAL_IMAGE_DIR}" \
-  --train_annotations_file="${TRAIN_ANNOTATIONS_FILE}" \
-  --val_annotations_file="${VAL_ANNOTATIONS_FILE}" \
-  --output_dir="${OUTPUT_DIR}" \
-  --include_masks
+    --train_image_dir="${TRAIN_IMAGE_DIR}" \
+    --val_image_dir="${VAL_IMAGE_DIR}" \
+    --train_annotations_file="${TRAIN_ANNOTATIONS_FILE}" \
+    --val_annotations_file="${VAL_ANNOTATIONS_FILE}" \
+    --output_dir="${OUTPUT_DIR}" \
+    --include_masks
 """
 import hashlib
 import io

diff --git a/rectlabel_create_pascal_tf_record.py b/rectlabel_create_pascal_tf_record.py
@@ -21,11 +21,11 @@
 
 
 flags = tf.app.flags
+flags.DEFINE_boolean('include_masks', False, 'Add image/object/mask to TFRecord using png images in annotations folder')
 flags.DEFINE_string('images_dir', '', 'Full path to the images directory.')
 flags.DEFINE_string('annotations_dir', 'annotations', '(Relative) path to annotations directory.')
 flags.DEFINE_string('label_map_path', 'data/pascal_label_map.pbtxt', 'Path to label map proto')
 flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
-flags.DEFINE_boolean('include_masks', False, 'Add image/object/mask to TFRecord using png images in annotations folder')
 flags.DEFINE_boolean('ignore_difficult_instances', False, 'Whether to ignore difficult instances')
 FLAGS = flags.FLAGS
 
@@ -48,7 +48,6 @@ def dict_to_tf_example(data, annotations_dir, images_dir, label_map_dict, includ
     image = PIL.Image.open(encoded_jpg_io)
     if image.format != 'JPEG':
         raise ValueError('Image format not JPEG')
-
     key = hashlib.sha256(encoded_jpg).hexdigest()
     width = int(data['size']['width'])
     height = int(data['size']['height'])
@@ -90,8 +89,6 @@ def dict_to_tf_example(data, annotations_dir, images_dir, label_map_dict, includ
                 mask_np = np.asarray(mask)
                 mask_remapped = (mask_np == 255).astype(np.uint8)
                 masks.append(mask_remapped)
-                # print(mask_remapped)
-
     feature_dict = {
         'image/height': dataset_util.int64_feature(height),
         'image/width': dataset_util.int64_feature(width),
@@ -120,9 +117,7 @@ def dict_to_tf_example(data, annotations_dir, images_dir, label_map_dict, includ
             output = io.BytesIO()
             img.save(output, format='PNG')
             encoded_mask_png_list.append(output.getvalue())
-            # print(output.getvalue())
         feature_dict['image/object/mask'] = (dataset_util.bytes_list_feature(encoded_mask_png_list))
-
     example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
     return example