Create images/ and labels/ subdirectories for converted KITTI formatted data, and also suffle images before splitting train/val sets

Author: jkjung-avt

create_kitti_tf_record.py

@@ -17,12 +17,16 @@
 # ========================================================================
 
 
-r"""Convert KITTI detection dataset to TFRecord for object_detection.
+"""Convert KITTI detection dataset to TFRecord for object detection.
 
 Example usage:
+
     python3 create_kitti_tf_record.py \
             --data_dir=egohands_kitti_formatted \
-            --output_path=egohands
+            --output_path=data/egohands \
+            --classes_to_use=hand \
+            --label_map_path=data/egohands_label_map.pbtxt \
+            --validation_set_size=500
 """
 
 
@@ -34,6 +38,7 @@
 import hashlib
 import io
 import os
+import random
 
 import numpy as np
 import PIL.Image as pil
@@ -72,7 +77,9 @@ def convert_kitti_to_tfrecords(data_dir, output_path, classes_to_use,
   """Convert the KITTI detection dataset to TFRecords.
 
   Args:
-    data_dir: The full path containing all jpg and txt files.
+    data_dir: The full path containing KITTI formatted data, with all
+      jpg files located in <data_dir>/images/ and all txt files in
+      <data_dir>/labels/.
     output_path: The path to which TFRecord files will be written.
       The TFRecord with the training set will be located at:
       <output_path>_train.tfrecord
@@ -92,8 +99,8 @@ def convert_kitti_to_tfrecords(data_dir, output_path, classes_to_use,
   train_count = 0
   val_count = 0
 
-  annotation_dir = data_dir
-  image_dir = data_dir
+  annotation_dir = os.path.join(data_dir, 'labels')
+  image_dir = os.path.join(data_dir, 'images')
 
   train_writer = tf.python_io.TFRecordWriter('%s_train.tfrecord'%
                                              output_path)
@@ -102,6 +109,8 @@ def convert_kitti_to_tfrecords(data_dir, output_path, classes_to_use,
 
   images = sorted(tf.gfile.ListDirectory(image_dir))
   images = [f for f in images if f.endswith('jpg')]  # only keep jpg files
+  assert len(images) > 0
+  random.shuffle(images)
   for idx, img_name in enumerate(images):
     img_num = img_name.split('.')[0]
     is_validation_img = idx < validation_set_size

bootstrap.py prepare_egohands.py

@@ -1,16 +1,23 @@
-"""bootstrap.py
+"""prepare_egohands.py
 
-This script downloads the 'egohands' dataset and convert the annotations
+This script downloads the 'egohands' dataset and convert its annotations
 into bounding boxes in KITTI format.
 
-Output of this script is data in KITTI format:
+Output of this script:
 
+  ./egohands_data.zip
+  ./egohands
+    ├── (egohands dataset unzipped)
+    └── ......
   ./egohands_kitti_formatted
-    ├── CARDS_COURTYARD_B_T_frame_0011.jpg
-    ├── CARDS_COURTYARD_B_T_frame_0011.txt
-    ├── ......
-    ├── PUZZLE_OFFICE_T_S_frame_2697.jpg
-    └── PUZZLE_OFFICE_T_S_frame_2697.txt
+    ├── images
+    │   ├── CARDS_COURTYARD_B_T_frame_0011.jpg
+    │   ├── ......
+    │   └── PUZZLE_OFFICE_T_S_frame_2697.jpg
+    └── labels
+        ├── CARDS_COURTYARD_B_T_frame_0011.txt
+        ├── ......
+        └── PUZZLE_OFFICE_T_S_frame_2697.txt
 """
 
 
@@ -22,15 +29,18 @@
 from zipfile import ZipFile
 from shutil import rmtree, copyfile
 
-import cv2
+import numpy as np
 from scipy.io import loadmat
+import cv2
 
 
 EGOHANDS_DATASET_URL = \
     'http://vision.soic.indiana.edu/egohands_files/egohands_data.zip'
 EGOHANDS_DIR = './egohands'
 EGOHANDS_DATA_DIR = './egohands/_LABELLED_SAMPLES'
 CONVERTED_DIR = './egohands_kitti_formatted'
+CONVERTED_IMG_DIR = './egohands_kitti_formatted/images'
+CONVERTED_LBL_DIR = './egohands_kitti_formatted/labels'
 
 VISUALIZE = False  # visualize each image (for debugging)
 
@@ -83,10 +93,10 @@ def polygon_to_box(polygon):
     if len(polygon) < 3:  # a polygon has at least 3 vertices
         return None
 
-    x_min = min(polygon[:, 0])
-    y_min = min(polygon[:, 1])
-    x_max = max(polygon[:, 0])
-    y_max = max(polygon[:, 1])
+    x_min = np.min(polygon[:, 0])
+    y_min = np.min(polygon[:, 1])
+    x_max = np.max(polygon[:, 0])
+    y_max = np.max(polygon[:, 1])
 
     x_min = int(math.floor(x_min))
     y_min = int(math.floor(y_min))
@@ -159,19 +169,19 @@ def convert_one_folder(folder):
         src_jpg = frame + '.jpg'
         dst_jpg = folder + '_' + src_jpg
         copyfile(os.path.join(folder_path, src_jpg),
-                 os.path.join(CONVERTED_DIR, dst_jpg))
+                 os.path.join(CONVERTED_IMG_DIR, dst_jpg))
         # generate txt (the KITTI annotation corresponding to the jpg)
         dst_txt = folder + '_' + frame + '.txt'
         boxes = []
-        with open(os.path.join(CONVERTED_DIR, dst_txt), 'w') as f:
+        with open(os.path.join(CONVERTED_LBL_DIR, dst_txt), 'w') as f:
             for polygon in polygons[i]:
                 box = polygon_to_box(polygon)
                 if box:
                     boxes.append(box)
                     f.write(box_to_line(box) + '\n')
 
         if VISUALIZE:
-            img = cv2.imread(os.path.join(CONVERTED_DIR, dst_jpg))
+            img = cv2.imread(os.path.join(CONVERTED_IMG_DIR, dst_jpg))
             for box in boxes:
                 cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]),
                               (0, 224, 0), 2)
@@ -190,6 +200,9 @@ def egohands_to_kitti():
       3. convert the original annotations ('polygon.mat') into
          bounding boxes and write a KITTI txt file for each image.
     """
+    rmtree(CONVERTED_DIR, ignore_errors=True)
+    os.makedirs(CONVERTED_IMG_DIR)
+    os.makedirs(CONVERTED_LBL_DIR)
     for folder in os.listdir(EGOHANDS_DATA_DIR):
         convert_one_folder(folder)
 
@@ -209,8 +222,6 @@ def main():
             zf.extractall(EGOHANDS_DIR)
 
     logging.info('Copying jpg files and converting annotations...')
-    rmtree(CONVERTED_DIR, ignore_errors=True)
-    os.makedirs(CONVERTED_DIR)
     egohands_to_kitti()
 
     logging.info('All done.')