forked from cedricoeldorf/ConditionalStyleGAN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
create_pickle.py
68 lines (47 loc) · 1.91 KB
/
create_pickle.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
import argparse
import glob
import pickle
import numpy as np
from PIL import Image
from os import listdir, makedirs
from os.path import join
def pickle_data(file, data_dict):
with open(file, 'wb') as fo:
pickle.dump(data_dict, fo)
def unpickle(file):
with open(file, 'rb') as fo:
_dict = pickle.load(fo, encoding='bytes')
return _dict
def main(dataset_path, path, size):
makedirs(path, exist_ok=True)
out_path = join(path, 'mypickle.pickle')
classes_names = sorted(listdir(dataset_path))
mypickle = {"Filenames": [], "Labels": []}
for index, classname in enumerate(classes_names):
for image in glob.glob(join(dataset_path, classname, '*.*')):
pil_image = Image.open(image)
pil_image = pil_image.resize((size, size))
pil_image.save(image)
np_pil_image = np.asarray(Image.open(image))
shape = np_pil_image.shape
resolution_log2 = int(np.log2(shape[0]))
if shape[2] not in [1, 3]:
continue
if shape[0] != 2**resolution_log2:
res = 2**resolution_log2
pil_image = pil_image.resize((res, res))
pil_image.save(image)
elif shape[0] != shape[1]:
res = min(shape[0], shape[1])
pil_image = pil_image.resize((res, res))
pil_image.save(image)
mypickle["Filenames"].append(image)
mypickle["Labels"].append(index)
pickle_data(out_path, mypickle)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--dataset_path", "-d", help="set dataset folder path")
parser.add_argument("--size", "-s", default=512, help="set dataset folder path")
parser.add_argument("--output_path", "-o", default='../data/', help="set output path")
args = parser.parse_args()
main(args.dataset_path, args.output_path, int(args.size))