forked from ykamikawa/tf-keras-SegNet
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathgenerator.py
More file actions
200 lines (173 loc) · 7.25 KB
/
Copy pathgenerator.py
File metadata and controls
200 lines (173 loc) · 7.25 KB
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
# -*- coding: utf-8 -*-
"""Generator module for Segnet"""
from random import randint, uniform
from keras.utils import to_categorical
import numpy as np
import cv2
def domain_generator(img_dir,
domain_dir,
img_list,
domain_list,
batch_size,
dims,
crop,
flip,
motion_blur,
sp_noise):
"""Continous generator"""
while True:
imgs = []
labels = []
# Choose random sequences of training images and domain images
ix = np.random.choice(np.arange(len(img_list)), batch_size)
iy = np.random.choice(np.arange(len(domain_list)), batch_size)
for i in range(0, batch_size):
img_path = ''
if randint(0, 1):
# training image
img_path = img_dir + img_list.iloc[ix[i], 0]
labels.append(0)
else:
# domain adaptation image
img_path = domain_dir + domain_list.iloc[iy[i], 0]
labels.append(1)
original_img = cv2.imread(img_path)
# switch colors to RGB
original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)
# Geometric transformations to fit the network
transformed_img, transformed_mask = transform_data(original_img, original_img, dims, crop, flip)
# Data augmentations
if uniform(0, 1) < motion_blur:
transformed_img = motion_blur_image(transformed_img)
if uniform(0, 1) < sp_noise:
transformed_img = sp_noise_image(transformed_img)
# Append image to main list
imgs.append(transformed_img)
yield np.array(imgs), to_categorical(labels, 2)
def segnet_generator(img_dir,
mask_dir,
lists,
batch_size,
dims,
n_labels,
crop,
flip,
motion_blur,
sp_noise):
"""Continous generator"""
while True:
yield single_batch_generator(img_dir,
mask_dir,
lists,
batch_size,
dims,
n_labels,
crop,
flip,
motion_blur,
sp_noise)
def single_batch_generator(img_dir,
mask_dir,
lists,
batch_size,
dims,
n_labels,
crop,
flip,
motion_blur=0,
sp_noise=0,
empty_mask=False):
"""Generate one batch of data"""
ix = np.random.choice(np.arange(len(lists)), batch_size)
imgs = []
labels = []
for i in ix:
# images
original_img = cv2.imread(img_dir + lists.iloc[i, 0])
# switch colors to RGB
original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)
# masks
original_mask = cv2.imread(mask_dir + lists.iloc[i, 0][:-4] + ".png")
# Geometric transformations to fit the network
transformed_img, transformed_mask = transform_data(original_img, original_mask, dims, crop, flip)
# Data augmentations
if uniform(0, 1) < motion_blur:
transformed_img = motion_blur_image(transformed_img)
if uniform(0, 1) < sp_noise:
transformed_img = sp_noise_image(transformed_img)
# Convert mask to labels
if empty_mask:
array_mask = np.zeros([dims[0], dims[1], n_labels])
else:
array_mask = to_categorical(transformed_mask[:, :, 0], n_labels)
array_mask = array_mask.reshape(dims[0] * dims[1], n_labels)
# Append image and mask to main lists
imgs.append(transformed_img)
labels.append(array_mask)
imgs = np.array(imgs)
labels = np.array(labels)
return imgs, labels
def transform_data(original_img, original_mask, dims, crop, flip):
"""Geometric transformations of images and mask"""
# Random crop or resize, openCV starts with height, not width
if crop:
random_x = randint(0, original_img.shape[0] - dims[1])
random_y = randint(0, original_img.shape[1] - dims[0])
transformed_img = original_img[random_y:random_y + dims[0], random_x:random_x + dims[1]]
transformed_mask = original_mask[random_y:random_y + dims[0], random_x:random_x + dims[1]]
else:
transformed_img = cv2.resize(original_img, (dims[1], dims[0]))
transformed_mask = cv2.resize(original_mask, (dims[1], dims[0]))
#Flip randomly images and masks
if flip:
orientation = randint(0, 3)
if orientation == 0: #horizontal
transformed_img = cv2.flip(transformed_img, 0)
transformed_mask = cv2.flip(transformed_mask, 0)
elif orientation == 1: #vertical
transformed_img = cv2.flip(transformed_img, 1)
transformed_mask = cv2.flip(transformed_mask, 1)
elif orientation == 2: #horizontal and vertical
transformed_img = cv2.flip(transformed_img, -1)
transformed_mask = cv2.flip(transformed_mask, -1)
elif orientation == 3: #none
pass
return transformed_img, transformed_mask
def motion_blur_image(image):
"""Degrade image quality - motion blur"""
# Motion blur
# generating the kernel
orientation = randint(0, 4)
kernel_motion_blur = np.zeros((3, 3))
if orientation == 0: #horizontal
kernel_motion_blur = np.array([[0, 0, 0],
[1, 1, 1],
[0, 0, 0]])
elif orientation == 1: #vertical
kernel_motion_blur = np.array([[0, 1, 0],
[0, 1, 0],
[0, 1, 0]])
elif orientation == 2: #diagonal 1
kernel_motion_blur = np.array([[1, 0, 0],
[0, 1, 0],
[0, 0, 1]])
elif orientation == 3: #diagonal 2
kernel_motion_blur = np.array([[0, 0, 1],
[0, 1, 0],
[1, 0, 0]])
kernel_motion_blur = kernel_motion_blur / 3
# applying the kernel to the input image
return cv2.filter2D(image, -1, kernel_motion_blur)
def sp_noise_image(image, salt_value=40):
noise = np.random.randint(salt_value + 1, size = (image.shape[0], image.shape[1]))
#---------- Pepper ----------#
index = np.where(noise == 0)
A = index[0]
B = index[1]
image[A,B,:] = 0
#---------- Salt ----------#
index = np.where(noise == salt_value)
A = index[0]
B = index[1]
image[A,B,:] = 255
return image