Single file to process images to Annotation, Images, and ImageSets

process_data.py

@@ -1,11 +1,17 @@
+#!/usr/bin/python
+
 import json
 import numpy as np
 import cv2
+import os
+
 '''
 TODO: make buttons for saving stuff
 '''
 def main():
-    # Train, val, test set split (3:1:1)
+    view_bounding_box = False
+
+    # Train, val, test set index split (3:1:1)
     num_images  =  4388  # number of images with buildings in them
     total_range = range(num_images)
     train_range = total_range[:num_images // 5 * 3]
@@ -17,27 +23,23 @@ def main():
     print("num_sum: {}".format(len(train_range) + len(val_range) + len(test_range)))
     print("num_images: {}".format(num_images))
 
-    # test_value = 5555
-    # if test_value in train_range:
-    #     print('.')
-    # if test_value in val_range:
-    #     print('..')
-    # if test_value in test_range:
-    #     print('...')
-    # return
-
-    view_bounding_box = False
+    # Image chips source
+    image_clip_path = '/home/ubuntu/dataset/processedData/3band/'
 
+    # Bounding box annotaiton source
     json_path = '/home/ubuntu/dataset/processedData/vectorData/summaryData/AOI_1_Rio_polygons_solution_3Band.geojson'
-    #json_path = '/home/ubuntu/dataset/processedData/geoJson/013022223130_Public_img4_Geo.geojson'
-
-    image_clip_path = '/home/ubuntu/dataset/processedData/3band/'
 
+    # Output destinations
     dataset_output_path = '/home/ubuntu/fast-rcnn/spacenet/data/'
     annotation_path = dataset_output_path + 'Annotations/'
-    image_set_path = dataset_output_path + 'ImageSets/'
+    image_set_path  = dataset_output_path + 'ImageSets/'
+    image_path      = dataset_output_path + 'Images/'
+
+    # Outputs for bounding box overlay original image chip
+    chip_bb_path = 'chips/'
 
-    chip_path = 'chips/'
+    # Band
+    band_prefix = '3band_'
 
     with open(json_path, 'rb') as f:
         geo_json = json.load(f)
@@ -48,7 +50,7 @@ def main():
     image_counter = -1
     image_name_dict = {}
 
-    # Each iteration contains 1 bounding box
+    # Each iteration and entry contains 1 bounding box
     for feature_dict in feature_collection:
         assert(feature_dict["type"] == "Feature")
 
@@ -57,23 +59,25 @@ def main():
         if building_id == -1:
             continue
 
-        # Some features for some reason has no coordinate
+        # Some features for some reason has no coordinate, skip them
         if not feature_dict["geometry"]["coordinates"]:
             continue
 
+        # This shouldn't happen in our dataset.
         if len(feature_dict["geometry"]["coordinates"]) != 1:
             print("\tThis entry has odd polygon:")
+            continue
 
-        # Grab the image clip
+        # Grab the associated image clip
         image_id = feature_dict["properties"]["ImageId"]
-        print(image_id + "     building id: " + str(building_id) + "     image count:" + str(image_counter))
+        print(image_id + "     building id: " + str(building_id) +
+              "     image count:" + str(image_counter))
         img = cv2.imread(image_clip_path + "3band_" + image_id + ".tif")
 
-
         coordinate_list = []
         coordinate_list_np = []
         if feature_dict["geometry"]["type"] == "Polygon":
-            # We are making the assumption here that the bounding box is a simply polygon
+            # We are making the assumption here that the bounding box is a simple polygon
             for coordinate in feature_dict["geometry"]["coordinates"][0]:
                 x = int(coordinate[0])
                 y = int(coordinate[1])
@@ -89,6 +93,7 @@ def main():
             thickness = 2
             # cv2.line(img, coordinate_list[coordinate_index], coordinate_list[coordinate_index + 1], color, thickness)
 
+        # Find the rectangular bounding box for the target
         # Need to change into an odd format that boundingRect takes
         x, y, w, h = cv2.boundingRect(np.expand_dims(np.array(coordinate_list), axis=1))
         # cv2.rectangle(img, (x,y), (x+w, y+h), (0,255,0), 1)
@@ -101,26 +106,30 @@ def main():
             image_counter += 1
             if image_counter in train_range:
                 with open(image_set_path + "train.txt", 'ab') as f:
-                    f.write(image_id + "\n")
+                    f.write(band_prefix + image_id + "\n")
             if image_counter in val_range:
                 with open(image_set_path + "val.txt", 'ab') as f:
-                    f.write(image_id + "\n")
+                    f.write(band_prefix + image_id + "\n")
             if image_counter in test_range:
                 with open(image_set_path + "test.txt", 'ab') as f:
-                    f.write(image_id + "\n")
-
+                    f.write(band_prefix + image_id + "\n")
 
-        # Write x_min, y_min, x_max, y_max as annotation
+        # Write to disk
         if image_counter in train_range:
-            with open(annotation_path + str(image_counter) + ".txt", "ab") as f:
+            # Write x_min, y_min, x_max, y_max as annotation required for fast-rcnn
+            with open(annotation_path + band_prefix + image_id + ".txt", "ab") as f:
                 f.write("{} {} {} {}\n".format(x, y, x+w, y+h))
+            # Record converted image type file if it doesn't exist already
+            current_image_path = image_path + band_prefix + image_id + ".png"
+            if not os.path.isfile(current_image_path):
+                cv2.imwrite(current_image_path, img)
 
         if view_bounding_box:
             cv2.imshow('image',img)
             cv2.waitKey(0)
             cv2.destroyAllWindows()
+            # cv2.imwrite(chip_bb_path + image_id + "_" + str(building_id) + ".jpg", img[y:y+h, x:x+w])
 
-        # cv2.imwrite(chip_path + image_id + "_" + str(building_id) + ".jpg", img[y:y+h, x:x+w])