"""
Utilities for real-time multi-thread data generator
"""
import scipy
import numpy as np
from tensorflow.keras.utils import Sequence, to_categorical
class CustomDataGenerator(Sequence):
def __init__(self, hdf5_file, brain_idx, batch_size=16, view="axial", mode='train', horizontal_flip=False,
vertical_flip=False, rotation_range=0, zoom_range=0., shuffle=True):
"""
Custom data generator based on Keras Sequance class.
This implementation enables multiprocessing and on-the-fly data augmentation
which will speed up training, especially in the task of brain tumor segmentation
that suffers from time-consuming data processing.
Parameters
----------
hdf5_file : file.File
An opend hdf5 file that contains all data.
brain_idx : array
The brain indexes corresponing to a specific fold. All of these
brain indexes will be use for training and the ones which are
not in 'brain_idx' will be used for validation
batch_size : int
The number of input/output arrays that will be generated each
time. The default is 16.
view : str
'axial', 'sagittal' or 'coronal'. The generator will extract
2D slices and perform normalization with respect to the chosen view.
The defualt is axial.
mode : str
Prepare the DataGenerator for 'train' or 'validation' phase.
The default is 'train'.
horizontal_flip : bool
Whether to use horizontal flip for data augmentation. The default is False.
vertical_flip : bool
Whether to use vertical flip for data augmentation. The default is False.
rotation_range : float
Random rotation for data augmentation. The default is 0.
zoom_range : float
Random zoom for data augmentation. The default is 0.
shuffle : bool
Whether to shuffle data. The default is True. Note that if mode='validation'
it will not shufflw tha data.
"""
self.data_storage = hdf5_file.root.data
self.truth_storage = hdf5_file.root.truth
total_brains = self.data_storage.shape[0]
self.brain_idx = self.get_brain_idx(brain_idx, mode, total_brains)
self.batch_size = batch_size
if view == 'axial':
self.view_axes = (0, 1, 2, 3)
elif view == 'sagittal':
self.view_axes = (2, 1, 0, 3)
elif view == 'coronal':
self.view_axes = (1, 2, 0, 3)
else:
ValueError('unknown input view => {}'.format(view))
self.mode = mode
self.horizontal_flip = horizontal_flip
self.vertical_flip = vertical_flip
self.rotation_range = rotation_range
self.zoom_range = [1 - zoom_range, 1 + zoom_range]
self.shuffle = shuffle
self.data_shape = tuple(np.array(self.data_storage.shape[1:])[np.array(self.view_axes)])
print('Using {} out of {} brains'.format(len(self.brain_idx), total_brains), end=' ')
print('({} out of {} 2D slices)'.format(len(self.brain_idx) * self.data_shape[0], total_brains * self.data_shape[0]))
print('the generated data shape in "{}" view: {}'.format(view, str(self.data_shape[1:])))
print('-----'*10)
self.on_epoch_end()
@staticmethod
def get_brain_idx(brain_idx, mode, total_brains):
"""
Getting the brain indexes that will be used by the generator.
if mode=='train' => the original indexes will be used (because we built these
npy files based on training indexes in 'prepare_data.py' for k-fold, remember? :)
if mode=='validation' => the indexes which are not in the brain_idx will
be used.
"""
if mode=='validation':
brain_idx = np.array([i for i in np.arange(total_brains) if i not in brain_idx])
print('DataGenerator is preparing for validation mode ...')
elif mode=='train':
brain_idx = brain_idx
print('DataGenerator is preparing for training mode ...')
else:
raise ValueError('unknown "{}" mode'.format(mode))
return brain_idx
def __len__(self):
return int(np.floor( len(self.indexes) / self.batch_size))
def __getitem__(self, index):
# Generate indexes of the batch
idx = self.indexes[index*self.batch_size:(index+1)*self.batch_size]
# Generate data
X_batch, Y_batch = self.data_load_and_preprocess(idx)
return X_batch, Y_batch
def on_epoch_end(self):
"""
Updates indexes after each epoch
"""
tmp=[]
for i in self.brain_idx:
for j in range(self.data_shape[0]):
tmp.append((i,j))
self.indexes = tmp
if self.mode=='train' and self.shuffle:
np.random.shuffle(self.indexes)
def data_load_and_preprocess(self, idx):
"""
Generates data containing batch_size samples
"""
slice_batch = []
label_batch = []
# Generate data
for i in idx:
brain_number = i[0]
slice_number = i[1]
slice_, label_ = self.read_data(brain_number, slice_number)
slice_ = self.normalize_modalities(slice_)
slice_and_label = np.concatenate((slice_, label_) , axis=-1)
params = self.get_random_transform()
slice_and_label = self.apply_transform(slice_and_label, params)
slice_ = slice_and_label[...,:4]
label_ = slice_and_label[..., 4]
label_ = to_categorical(label_, 4)
slice_batch.append(slice_)
label_batch.append(label_)
return np.array(slice_batch), np.array(label_batch)
def read_data(self, brain_number, slice_number):
"""
Reads data from table with respect to the 'view'
"""
slice_ = self.data_storage[brain_number].transpose(self.view_axes)[slice_number]
label_ = self.truth_storage[brain_number].transpose(self.view_axes[:3])[slice_number]
label_ = np.expand_dims(label_, axis=-1)
return slice_, label_
def normalize_slice(self, slice):
"""
Removes 1% of the top and bottom intensities and perform
normalization on the input 2D slice.
"""
b = np.percentile(slice, 99)
t = np.percentile(slice, 1)
slice = np.clip(slice, t, b)
if np.std(slice)==0:
return slice
else:
slice = (slice - np.mean(slice)) / np.std(slice)
return slice
def normalize_modalities(self, Slice):
"""
Performs normalization on each modalities of input
"""
normalized_slices = np.zeros_like(Slice).astype(np.float32)
for slice_ix in range(4):
normalized_slices[..., slice_ix] = self.normalize_slice(Slice[..., slice_ix])
return normalized_slices
def flip_axis(self, x, axis):
x = np.asarray(x).swapaxes(axis, 0)
x = x[::-1, ...]
x = x.swapaxes(0, axis)
return x
def apply_transform(self, x, transform_parameters):
x = apply_affine_transform(x, transform_parameters.get('theta', 0),
transform_parameters.get('tx', 0),
transform_parameters.get('ty', 0),
transform_parameters.get('shear', 0),
transform_parameters.get('zx', 1),
transform_parameters.get('zy', 1),
row_axis=0,
col_axis=1,
channel_axis=2)
if transform_parameters.get('flip_horizontal', False):
x = self.flip_axis(x, 1)
if transform_parameters.get('flip_vertical', False):
x = self.flip_axis(x, 0)
return x
def get_random_transform(self):
if self.rotation_range:
theta = np.random.uniform(-self.rotation_range,self.rotation_range)
else:
theta = 0
if self.zoom_range[0] == 1 and self.zoom_range[1] == 1:
zx, zy = 1, 1
else:
zx, zy = np.random.uniform(self.zoom_range[0],self.zoom_range[1], 2)
flip_horizontal = (np.random.random() < 0.5) * self.horizontal_flip
flip_vertical = (np.random.random() < 0.5) * self.vertical_flip
transform_parameters = {'flip_horizontal': flip_horizontal,
'flip_vertical':flip_vertical,
'theta': theta,
'zx': zx,
'zy': zy}
return transform_parameters
"""
The two following functions are from ImageDataGenerator class of keras.
https://github.com/keras-team/keras/blob/master/keras/preprocessing/image.py
"""
def apply_affine_transform(x, theta=0, tx=0, ty=0, shear=0, zx=1, zy=1,
row_axis=0, col_axis=1, channel_axis=2,
fill_mode='nearest', cval=0.):
"""Applies an affine transformation specified by the parameters given.
# Arguments
x: 2D numpy array, single image.
theta: Rotation angle in degrees.
tx: Width shift.
ty: Heigh shift.
shear: Shear angle in degrees.
zx: Zoom in x direction.
zy: Zoom in y direction
row_axis: Index of axis for rows in the input image.
col_axis: Index of axis for columns in the input image.
channel_axis: Index of axis for channels in the input image.
fill_mode: Points outside the boundaries of the input
are filled according to the given mode
(one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
cval: Value used for points outside the boundaries
of the input if `mode='constant'`.
# Returns
The transformed version of the input.
"""
transform_matrix = None
if theta != 0:
theta = np.deg2rad(theta)
rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
[np.sin(theta), np.cos(theta), 0],
[0, 0, 1]])
transform_matrix = rotation_matrix
if tx != 0 or ty != 0:
shift_matrix = np.array([[1, 0, tx],
[0, 1, ty],
[0, 0, 1]])
if transform_matrix is None:
transform_matrix = shift_matrix
else:
transform_matrix = np.dot(transform_matrix, shift_matrix)
if shear != 0:
shear = np.deg2rad(shear)
shear_matrix = np.array([[1, -np.sin(shear), 0],
[0, np.cos(shear), 0],
[0, 0, 1]])
if transform_matrix is None:
transform_matrix = shear_matrix
else:
transform_matrix = np.dot(transform_matrix, shear_matrix)
if zx != 1 or zy != 1:
zoom_matrix = np.array([[zx, 0, 0],
[0, zy, 0],
[0, 0, 1]])
if transform_matrix is None:
transform_matrix = zoom_matrix
else:
transform_matrix = np.dot(transform_matrix, zoom_matrix)
if transform_matrix is not None:
h, w = x.shape[row_axis], x.shape[col_axis]
transform_matrix = transform_matrix_offset_center(
transform_matrix, h, w)
x = np.rollaxis(x, channel_axis, 0)
final_affine_matrix = transform_matrix[:2, :2]
final_offset = transform_matrix[:2, 2]
channel_images = [scipy.ndimage.interpolation.affine_transform(
x_channel,
final_affine_matrix,
final_offset,
order=1,
mode=fill_mode,
cval=cval) for x_channel in x]
x = np.stack(channel_images, axis=0)
x = np.rollaxis(x, 0, channel_axis + 1)
return x
def transform_matrix_offset_center(matrix, x, y):
o_x = float(x) / 2 + 0.5
o_y = float(y) / 2 + 0.5
offset_matrix = np.array([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])
reset_matrix = np.array([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])
transform_matrix = np.dot(np.dot(offset_matrix, matrix), reset_matrix)
return transform_matrix
