-
Notifications
You must be signed in to change notification settings - Fork 75
/
build_dataset.py
114 lines (94 loc) · 4.48 KB
/
build_dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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
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")
parser.add_argument('--up_scale', default='4', help="Where to write the new data")
def crop_and_save(filename, output_dir, out_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, out_size, up_scale):
"""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 // up_scale, out_size // up_scale)) # upscaling factor 4
blur = resize(image_resized, (in_size, in_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)
UP_SCALE = int(args.up_scale)
# 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, OUTPUT_SIZE)
#blur image
for filename in tqdm(filenames[split]):
blur_and_save(filename, output_dir_split_blur, INPUT_SIZE, OUTPUT_SIZE, UP_SCALE)
print("Done building dataset")