/
augmentations.py
208 lines (167 loc) · 8.63 KB
/
augmentations.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
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
201
202
203
204
205
206
207
208
import math
import tensorflow as tf
from data_utils import exterior_exclusion
def random_rotation(image, max_degrees, bbox=None, prob=0.5):
"""Applies random rotation to image and bbox"""
def _rotation(image, bbox):
# Get random angle
degrees = tf.random.uniform([], minval=-max_degrees, maxval=max_degrees, dtype=tf.float32)
radians = degrees * math.pi / 180.
if bbox is not None:
# Get offset from image center
image_shape = tf.cast(tf.shape(image), tf.float32)
image_height, image_width = image_shape[0], image_shape[1]
bbox = tf.cast(bbox, tf.float32)
center_x = image_width / 2.
center_y = image_height / 2.
bbox_center_x = (bbox[0] + bbox[2]) / 2.
bbox_center_y = (bbox[1] + bbox[3]) / 2.
trans_x = center_x - bbox_center_x
trans_y = center_y - bbox_center_y
# Apply rotation
image = _translate_image(image, trans_x, trans_y)
bbox = _translate_bbox(bbox, image_height, image_width, trans_x, trans_y)
image = tf.contrib.image.rotate(image, radians, interpolation='BILINEAR')
bbox = _rotate_bbox(bbox, image_height, image_width, radians)
image = _translate_image(image, -trans_x, -trans_y)
bbox = _translate_bbox(bbox, image_height, image_width, -trans_x, -trans_y)
bbox = tf.cast(bbox, tf.int32)
return image, bbox
return tf.contrib.image.rotate(image, radians, interpolation='BILINEAR')
retval = image if bbox is None else (image, bbox)
return tf.cond(_should_apply(prob), lambda: _rotation(image, bbox), lambda: retval)
def random_bbox_jitter(bbox, image_height, image_width, max_fraction, prob=0.5):
"""Randomly jitters bbox coordinates by +/- jitter_fraction of the width/height"""
def _bbox_jitter(bbox):
bbox = tf.cast(bbox, tf.float32)
width_jitter = max_fraction*(bbox[2] - bbox[0])
height_jitter = max_fraction*(bbox[3] - bbox[1])
xmin = bbox[0] + tf.random.uniform([], minval=-width_jitter, maxval=width_jitter, dtype=tf.float32)
ymin = bbox[1] + tf.random.uniform([], minval=-height_jitter, maxval=height_jitter, dtype=tf.float32)
xmax = bbox[2] + tf.random.uniform([], minval=-width_jitter, maxval=width_jitter, dtype=tf.float32)
ymax = bbox[3] + tf.random.uniform([], minval=-height_jitter, maxval=height_jitter, dtype=tf.float32)
xmin, ymin, xmax, ymax = _clip_bbox(xmin, ymin, xmax, ymax, image_height, image_width)
bbox = tf.cast(tf.stack([xmin, ymin, xmax, ymax]), tf.int32)
return bbox
return tf.cond(_should_apply(prob), lambda: _bbox_jitter(bbox), lambda: bbox)
def random_shift_and_scale(image, max_shift, max_scale_change, prob=0.5):
"""Applies random shift and scale to pixel values"""
def _shift_and_scale(image):
shift = tf.cast(tf.random.uniform([], minval=-max_shift, maxval=max_shift, dtype=tf.int32), tf.float32)
scale = tf.random.uniform([], minval=(1. - max_scale_change),
maxval=(1. + max_scale_change), dtype=tf.float32)
image = scale*(tf.cast(image, tf.float32) + shift)
image = tf.cast(tf.clip_by_value(image, 0., 255.), tf.uint8)
return image
return tf.cond(_should_apply(prob), lambda: _shift_and_scale(image), lambda: image)
def random_shear(image, max_lambda, bbox=None, prob=0.5):
"""Applies random shear in either the x or y direction"""
shear_lambda = tf.random.uniform([], minval=-max_lambda, maxval=max_lambda, dtype=tf.float32)
image_shape = tf.cast(tf.shape(image), tf.float32)
image_height, image_width = image_shape[0], image_shape[1]
def _shear_x(image, bbox):
image = _shear_x_image(image, shear_lambda)
if bbox is not None:
bbox = _shear_bbox(bbox, image_height, image_width, shear_lambda, horizontal=True)
bbox = tf.cast(bbox, tf.int32)
return image, bbox
return image
def _shear_y(image, bbox):
image = _shear_y_image(image, shear_lambda)
if bbox is not None:
bbox = _shear_bbox(bbox, image_height, image_width, shear_lambda, horizontal=False)
bbox = tf.cast(bbox, tf.int32)
return image, bbox
return image
def _shear(image, bbox):
return tf.cond(_should_apply(0.5), lambda: _shear_x(image, bbox), lambda: _shear_y(image, bbox))
retval = image if bbox is None else (image, bbox)
return tf.cond(_should_apply(prob), lambda: _shear(image, bbox), lambda: retval)
def random_exterior_exclusion(image, prob=0.5):
"""Randomly removes visual features exterior to the patient's body"""
def _exterior_exclusion(image):
shape = image.get_shape()
image = tf.py_func(exterior_exclusion, [image], tf.uint8)
image.set_shape(shape)
return image
return tf.cond(_should_apply(prob), lambda: _exterior_exclusion(image), lambda: image)
def _translate_image(image, delta_x, delta_y):
"""Translate an image"""
return tf.contrib.image.translate(image, [delta_x, delta_y], interpolation='BILINEAR')
def _translate_bbox(bbox, image_height, image_width, delta_x, delta_y):
"""Translate an bbox, ensuring coordinates lie in the image"""
bbox = bbox + tf.stack([delta_x, delta_y, delta_x, delta_y])
xmin, ymin, xmax, ymax = _clip_bbox(bbox[0], bbox[1], bbox[2], bbox[3], image_height, image_width)
bbox = tf.stack([xmin, ymin, xmax, ymax])
return bbox
def _rotate_bbox(bbox, image_height, image_width, radians):
"""Rotates the bbox by the given angle"""
# Shift bbox to origin
xmin, ymin, xmax, ymax = bbox[0], bbox[1], bbox[2], bbox[3]
center_x = (xmin + xmax) / 2.
center_y = (ymin + ymax) / 2.
xmin = xmin - center_x
xmax = xmax - center_x
ymin = ymin - center_y
ymax = ymax - center_y
# Rotate bbox coordinates
radians = -radians # negate direction since y-axis is flipped
coords = tf.stack([[xmin, ymin], [xmax, ymin], [xmin, ymax], [xmax, ymax]])
coords = tf.transpose(tf.cast(coords, tf.float32))
rotation_matrix = tf.stack(
[[tf.cos(radians), -tf.sin(radians)],
[tf.sin(radians), tf.cos(radians)]])
new_coords = tf.matmul(rotation_matrix, coords)
# Find new bbox coordinates and clip to image size
xmin = tf.reduce_min(new_coords[0, :]) + center_x
ymin = tf.reduce_min(new_coords[1, :]) + center_y
xmax = tf.reduce_max(new_coords[0, :]) + center_x
ymax = tf.reduce_max(new_coords[1, :]) + center_y
xmin, ymin, xmax, ymax = _clip_bbox(xmin, ymin, xmax, ymax, image_height, image_width)
bbox = tf.stack([xmin, ymin, xmax, ymax])
return bbox
def _shear_x_image(image, shear_lambda):
"""Shear image in x-direction"""
tform = tf.stack([1., shear_lambda, 0., 0., 1., 0., 0., 0.])
image = tf.contrib.image.transform(
image, tform, interpolation='BILINEAR')
return image
def _shear_y_image(image, shear_lambda):
"""Shear image in y-direction"""
tform = tf.stack([1., 0., 0., shear_lambda, 1., 0., 0., 0.])
image = tf.contrib.image.transform(
image, tform, interpolation='BILINEAR')
return image
def _shear_bbox(bbox, image_height, image_width, shear_lambda, horizontal=True):
"""Shear bbox in x- or y-direction"""
# Shear bbox coordinates
xmin, ymin, xmax, ymax = bbox[0], bbox[1], bbox[2], bbox[3]
coords = tf.stack([[xmin, ymin], [xmax, ymin], [xmin, ymax], [xmax, ymax]])
coords = tf.transpose(tf.cast(coords, tf.float32))
if horizontal:
shear_matrix = tf.stack(
[[1., -shear_lambda],
[0., 1.]])
else:
shear_matrix = tf.stack(
[[1., 0.],
[-shear_lambda, 1.]])
new_coords = tf.matmul(shear_matrix, coords)
# Find new bbox coordinates and clip to image size
xmin = tf.reduce_min(new_coords[0, :])
ymin = tf.reduce_min(new_coords[1, :])
xmax = tf.reduce_max(new_coords[0, :])
ymax = tf.reduce_max(new_coords[1, :])
xmin, ymin, xmax, ymax = _clip_bbox(xmin, ymin, xmax, ymax, image_height, image_width)
bbox = tf.stack([xmin, ymin, xmax, ymax])
return bbox
def _clip_bbox(xmin, ymin, xmax, ymax, image_height, image_width):
"""Clip bbox to valid image coordinates"""
xmin = tf.clip_by_value(xmin, 0, image_width)
ymin = tf.clip_by_value(ymin, 0, image_height)
xmax = tf.clip_by_value(xmax, 0, image_width)
ymax = tf.clip_by_value(ymax, 0, image_height)
return xmin, ymin, xmax, ymax
def _should_apply(prob):
"""Helper function to create bool tensor with probability"""
return tf.cast(tf.floor(tf.random_uniform([], dtype=tf.float32) + prob), tf.bool)