-
Notifications
You must be signed in to change notification settings - Fork 1
/
data_loader.py
116 lines (85 loc) · 3.97 KB
/
data_loader.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
# This code is heavily inspired by the works of
# https://github.com/eriklindernoren/Keras-GAN/blob/master/cyclegan
# under the MIT License Copyright (c) 2017 Erik Linder-Norén
import scipy
import cv2
from matplotlib import pyplot as plt
from glob import glob
import numpy as np
class DataLoader():
def __init__(self, img_res=(128, 128)):
self.img_res = img_res
def load_data(self, nb_samples=1):
handbag_path = glob('dataset/edges2handbags_sample/*.jpg')
shoe_path = glob("dataset/shoes-athletic/*.jpg")
handbags = np.zeros((nb_samples, self.img_res[0], self.img_res[1], 3), dtype=np.float32)
shoes = np.zeros((nb_samples, self.img_res[0], self.img_res[1], 3), dtype=np.float32)
for i in range(nb_samples):
handbag = handbag_path[i + 100] # 100 shoe/handbag because 1st are same color
shoe = shoe_path[i + 200]
# Importing handbags
try:
handbag = self.imread(handbag)
h, w, _ = handbag.shape
# Separate edge and image
half_w = int(w/2)
# handbag_edge = handbag[:, :half_w, :]
handbag_img = handbag[:, half_w:, :]
# handbag_edge = cv2.resize(handbag_edge, self.img_res)
handbags[i,:,:,:] = cv2.resize(handbag_img, self.img_res)
except Exception as e:
print("Couldn't load handbag")
print(e)
# Importing shoes
try:
shoe_img = self.imread(shoe)
shoes[i,:,:,:] = cv2.resize(shoe_img, self.img_res)
except Exception as e:
print("Couldn't load shoe")
print(e)
# Normalizing arrays to [-1, 1]
shoes_norm = (shoes - 255/2)/(255/2)
handbags_norm = (handbags - 255/2)/(255/2)
return shoes_norm, handbags_norm
def load_batch(self, batch_size=1):
handbag_path = glob('./dataset/edges2handbags_sample/*.jpg')
shoe_path = glob("./dataset/shoes-athletic/*.jpg")
handbag_n_batches = int(len(handbag_path) / batch_size)
shoe_n_batches = int(len(shoe_path) / batch_size)
self.n_batches = handbag_n_batches if handbag_n_batches < shoe_n_batches else shoe_n_batches
for i in range(self.n_batches-1):
handbags = np.zeros((batch_size, self.img_res[0], self.img_res[1], 3), dtype=np.float32)
shoes = np.zeros((batch_size, self.img_res[0], self.img_res[1], 3), dtype=np.float32)
handbag_batch = handbag_path[i*batch_size:(i+1)*batch_size]
shoe_batch = shoe_path[i*batch_size:(i+1)*batch_size]
# Importing handbags
for k,handbag in enumerate(handbag_batch):
try:
handbag = self.imread(handbag)
h, w, _ = handbag.shape
# Separate edge and image
half_w = int(w/2)
# handbag_edge = handbag[:, :half_w, :]
handbag_img = handbag[:, half_w:, :]
# handbag_edge = cv2.resize(handbag_edge, self.img_res)
handbags[k,:,:,:] = cv2.resize(handbag_img, self.img_res)
except Exception as e:
print("Couldn't load handbag")
print(e)
# Importing shoes
for k, shoe in enumerate(shoe_batch):
try:
shoe_img = self.imread(shoe)
shoes[k,:,:,:] = cv2.resize(shoe_img, self.img_res)
except Exception as e:
print("Couldn't load shoe")
print(e)
# Normalizing arrays to [-1, 1]
shoes_norm = (shoes- 255/2)/(255/2)
handbags_norm = (handbags - 255/2)/(255/2)
yield shoes_norm, handbags_norm
def imread(self, path):
return cv2.imread(path)
if __name__=='__main__':
dl = DataLoader()
dl.load_data()