-
Notifications
You must be signed in to change notification settings - Fork 34
/
Copy pathimage_loader.py
95 lines (77 loc) · 5.28 KB
/
image_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
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
def get_stereo_image_generators(train_folder, test_folder, img_rows=128, img_cols=512, batch_size=16, shuffle=True):
train_image_gen = ImageDataGenerator(rescale=1.0 / 255.0,
# TODO: when network is complete, add image transformations to improve training
# rotation_range=5,
# shear_range=0.01,
# zoom_range=0.01,
# height_shift_range=0.01,
# width_shift_range=0.01
)
test_image_gen = ImageDataGenerator(rescale=1.0 / 255.0)
train_generator_left = train_image_gen.flow_from_directory(train_folder,
target_size=(img_rows, img_cols),
batch_size=batch_size,
seed=10,
shuffle=shuffle,
classes=['left'],
class_mode=None,
follow_links=True)
train_generator_right = train_image_gen.flow_from_directory(train_folder,
target_size=(img_rows, img_cols),
batch_size=batch_size,
seed=10,
shuffle=shuffle,
classes=['right'],
class_mode=None,
follow_links=True)
test_generator_left = test_image_gen.flow_from_directory(test_folder,
target_size=(img_rows, img_cols),
batch_size=batch_size,
seed=10,
shuffle=shuffle,
classes=['left'],
class_mode=None,
follow_links=True)
test_generator_right = test_image_gen.flow_from_directory(test_folder,
target_size=(img_rows, img_cols),
batch_size=batch_size,
seed=10,
shuffle=shuffle,
classes=['right'],
class_mode=None,
follow_links=True)
def train_generator_func():
while True:
left_image = train_generator_left.next()
right_image = train_generator_right.next()
yield [left_image, right_image], [left_image,
right_image,
np.zeros((left_image.shape[0],
left_image.shape[1],
left_image.shape[2],
1)),
np.zeros((right_image.shape[0],
right_image.shape[1],
right_image.shape[2],
1))]
def test_generator_func():
while True:
left_image = test_generator_left.next()
right_image = test_generator_right.next()
yield [left_image, right_image], [left_image,
right_image,
np.zeros((left_image.shape[0],
left_image.shape[1],
left_image.shape[2],
1)),
np.zeros((right_image.shape[0],
right_image.shape[1],
right_image.shape[2],
1))]
train_generator = train_generator_func()
test_generator = test_generator_func()
train_length = len(train_generator_left.filenames)
test_length = len(test_generator_left.filenames)
return train_generator, test_generator, train_length, test_length