-
Notifications
You must be signed in to change notification settings - Fork 4
/
data_generator.py
27 lines (22 loc) · 977 Bytes
/
data_generator.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
from keras.preprocessing.image import ImageDataGenerator
class FlippedImageDataGenerator(ImageDataGenerator):
flip_indices = [
(0, 2), (1, 3),
(4, 8), (5, 9), (6, 10), (7, 11),
(12, 16), (13, 17), (14, 18), (15, 19),
(22, 24), (23, 25),
]
def next(self):
X_batch, y_batch = super(FlippedImageDataGenerator, self).next()
batch_size = X_batch.shape[0]
indices = np.random.choice(batch_size, batch_size/2, replace=False)
X_batch[indices] = X_batch[indices, :, :, ::-1]
if y_batch is not None:
# x座標をフリップ
y_batch[indices, ::2] = y_batch[indices, ::2] * -1
# left_eye_center_x -> right_eye_center_x のようにフリップ
for a, b in self.flip_indices:
y_batch[indices, a], y_batch[indices, b] = (
y_batch[indices, b], y_batch[indices, a]
)
return X_batch, y_batch