-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathutils.py
129 lines (95 loc) · 4 KB
/
utils.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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct 12 14:08:56 2018
@author: andrea
"""
import os
import glob
import cv2
import numpy as np
import pickle
import os.path
def load_dataset(image_size, create_new, flip_img, roteate_img):
c1 = os.path.isfile("saved_dataset/X_train.pickle")
c2 = os.path.isfile("saved_dataset/y_train.pickle")
if c1 and c2 and not create_new:
with open('saved_dataset/X_train.pickle', 'rb') as data:
X = pickle.load(data)
with open('saved_dataset/y_train.pickle', 'rb') as data:
y = pickle.load(data)
number_of_classes = max(y)[0] +1
print("Dataset loaded successfully")
else:
X, y, number_of_classes = create_dataset(image_size, flip_img, roteate_img)
return X, y, number_of_classes
def create_dataset(image_size, flip_img, roteate_img):
"""
Create the dataset in the required format.
It assumes that the images are devided per class in subfolders inside the
folder called data
"""
print("Creating the dataset (might take a little)")
label_ref = "Mapping from folder class to numberic label used: \n\n"
X = []
y = []
folder_list = glob.glob("data/*")
number_of_classes = len([folder for folder in folder_list if os.path.isdir(folder)])
# Parse through each folder:
folder_counter = -1
for folder_name in folder_list:
if os.path.isdir(folder_name):
folder_counter += 1
label_ref += folder_name+" : "+str(folder_counter )+"\n"
image_list = glob.glob(folder_name+"/*")
# Parse through each image in the current folder:
for image_name in image_list:
X, y = add_image(X, y, image_name, image_size,
folder_counter, roteate_img, flip_img)
#Normalize and format the data:
X = np.array(X)
X = (X-127.5)/127.5
y = np.array(y)
y = y.reshape(-1,1)
#Shuffle:
seed = 547
np.random.seed(seed)
np.random.shuffle(X)
np.random.seed(seed)
np.random.shuffle(y)
with open('saved_dataset/X_train.pickle', 'wb') as output:
pickle.dump(X, output)
with open('saved_dataset/y_train.pickle', 'wb') as output:
pickle.dump(y, output)
#txt file showing the mapping from folder class to numberic label used
with open("labels.txt", "w") as text_file:
print(label_ref , file=text_file)
print("""Dataset correctly created and saved.
Total number of samples: """+str(X.shape[0]) )
return X, y, number_of_classes
def add_image(X, y, image_name, image_size, folder_counter,
roteate = False, mirror = False):
img = cv2.imread(image_name)
if img is not None:
img = cv2.resize(img, (image_size, image_size), interpolation = cv2.INTER_AREA)
X.append(img.astype(np.float32))
y.append(np.uint8(folder_counter))
# Roteate 90-180-270 degrees:
if roteate:
for i in range(3):
img = np.rot90(img)
X.append(img.astype(np.float32))
y.append(np.uint8(folder_counter))
#restore image to original orientation
img = np.rot90(img)
# Mirror horizontal and vertical
if mirror:
img_hor = cv2.flip( img, 0 )
X.append(img_hor.astype(np.float32))
y.append(np.uint8(folder_counter))
img_ver = cv2.flip( img, 1 )
X.append(img_ver.astype(np.float32))
y.append(np.uint8(folder_counter))
else:
print("Could not load ",image_name,"Is it an image?")
return X, y