build datasets

build_dataset.py

@@ -0,0 +1,117 @@
+import argparse
+import random
+import os
+
+from PIL import Image
+from tqdm import tqdm
+
+from skimage import io
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import numpy as np
+import matplotlib.pyplot as plt
+from skimage import data, color
+from skimage.transform import rescale, resize, downscale_local_mean
+from skimage import util
+
+INPUT_SIZE = 144
+OUTPUT_SIZE = 144
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--data_dir', default='../img_align_celeba_test', help="Directory with the SIGNS dataset")
+parser.add_argument('--output_dir', default='../data/cnn_faces', help="Where to write the new data")
+parser.add_argument('--input_size', default='144', help="Where to write the new data")
+parser.add_argument('--output_size', default='144', help="Where to write the new data")
+
+
+def crop_and_save(filename, output_dir, out_size=OUTPUT_SIZE):
+    """crop the image contained in `filename` and save it to the `output_dir`"""
+    image = io.imread(filename)
+    vert_start = (218 - out_size) // 2
+    vert_end = vert_start + out_size
+    horiz_start = (178 - out_size) // 2
+    horiz_end = horiz_start + out_size
+    cropped = image[vert_start:vert_end, horiz_start:horiz_end] # 218*178 -> 144*144
+    io.imsave(os.path.join(output_dir, filename.split('/')[-1]), cropped)
+
+def blur_and_save(filename, output_dir, in_size=INPUT_SIZE, out_size=OUTPUT_SIZE):
+    """Blur the image contained in `filename` and save it to the `output_dir`"""
+    image = io.imread(filename)
+    vert_start = (218 - out_size) // 2
+    vert_end = vert_start + out_size
+    horiz_start = (178 - out_size) // 2
+    horiz_end = horiz_start + out_size
+
+    cropped = image[vert_start:vert_end, horiz_start:horiz_end] # 218*178 -> 144*144    
+    image_resized = resize(cropped, (out_size // 4, out_size // 4)) # upscaling factor 4
+    blur = resize(image_resized, (out_size, out_size)) # rescale back to 144 * 144
+    io.imsave(os.path.join(output_dir, filename.split('/')[-1]), blur)
+
+if __name__ == '__main__':
+    args = parser.parse_args()
+
+    assert os.path.isdir(args.data_dir), "Couldn't find the dataset at {}".format(args.data_dir)
+
+    # get args
+    data_dir = args.data_dir
+    INPUT_SIZE = int(args.input_size)
+    OUTPUT_SIZE = int(args.output_size)
+
+    # Get the filenames in data directory
+    filenames = os.listdir(data_dir)
+    filenames = [os.path.join(data_dir, f) for f in filenames if f.endswith('.jpg')]
+
+    # Split the images into 98% train, 1% val, and 1% test
+    # Make sure to always shuffle with a fixed seed so that the split is reproducible
+    random.seed(230)
+    filenames.sort()
+    random.shuffle(filenames)
+
+    split1 = int(0.98 * len(filenames))
+    split2 = (len(filenames) - split1) // 2 + split1
+    train_filenames = filenames[:split1]
+    val_filenames = filenames[split1:split2]
+    test_filenames = filenames[split2:]
+
+    print("train", len(train_filenames))
+    print("val", len(val_filenames))
+    print("test", len(test_filenames))
+
+    filenames = {'train': train_filenames,
+                 'val': val_filenames,
+                 'test': test_filenames}
+
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+    else:
+        print("Warning: output dir {} already exists".format(args.output_dir))
+
+
+    # Preprocess train, val and test
+    for split in ['train', 'val', 'test']:
+        # clear image
+        output_dir_split_clear = os.path.join(args.output_dir, '{}_clear'.format(split))
+        # blur image
+        output_dir_split_blur = os.path.join(args.output_dir, '{}_blur'.format(split))
+
+        # clear image
+        if not os.path.exists(output_dir_split_clear):
+            os.mkdir(output_dir_split_clear)
+        else:
+            print("Warning: dir {} already exists".format(output_dir_split_clear))
+        # blur image
+        if not os.path.exists(output_dir_split_blur):
+            os.mkdir(output_dir_split_blur)
+        else:
+            print("Warning: dir {} already exists".format(output_dir_split_blur))
+
+        print("Processing {} data, saving to {} and {}".format(split, output_dir_split_clear, output_dir_split_blur))
+        # clear image
+        for filename in tqdm(filenames[split]):
+            crop_and_save(filename, output_dir_split_clear, out_size=OUTPUT_SIZE)
+        #blur image
+        for filename in tqdm(filenames[split]):
+            blur_and_save(filename, output_dir_split_blur, in_size=INPUT_SIZE, out_size=OUTPUT_SIZE)
+
+    print("Done building dataset")
+