# Keras implementation of the paper:
# 3D MRI Brain Tumor Segmentation Using Autoencoder Regularization
# by Myronenko A. (https://arxiv.org/pdf/1810.11654.pdf)
# Author of this code: Suyog Jadhav (https://github.com/IAmSUyogJadhav)
import keras.backend as K
from keras.losses import mse
from keras.layers import Conv3D, Activation, Add, UpSampling3D, Lambda, Dense
from keras.layers import Input, Reshape, Flatten, Dropout, SpatialDropout3D
from keras.optimizers import adam
from keras.models import Model
try:
from group_norm import GroupNormalization
except ImportError:
import urllib.request
print('Downloading group_norm.py in the current directory...')
url = 'https://raw.githubusercontent.com/titu1994/Keras-Group-Normalization/master/group_norm.py'
urllib.request.urlretrieve(url, "group_norm.py")
from group_norm import GroupNormalization
def green_block(inp, filters, data_format='channels_first', name=None):
"""
green_block(inp, filters, name=None)
------------------------------------
Implementation of the special residual block used in the paper. The block
consists of two (GroupNorm --> ReLu --> 3x3x3 non-strided Convolution)
units, with a residual connection from the input `inp` to the output. Used
internally in the model. Can be used independently as well.
Parameters
----------
`inp`: An keras.layers.layer instance, required
The keras layer just preceding the green block.
`filters`: integer, required
No. of filters to use in the 3D convolutional block. The output
layer of this green block will have this many no. of channels.
`data_format`: string, optional
The format of the input data. Must be either 'chanels_first' or
'channels_last'. Defaults to `channels_first`, as used in the paper.
`name`: string, optional
The name to be given to this green block. Defaults to None, in which
case, keras uses generated names for the involved layers. If a string
is provided, the names of individual layers are generated by attaching
a relevant prefix from [GroupNorm_, Res_, Conv3D_, Relu_, ], followed
by _1 or _2.
Returns
-------
`out`: A keras.layers.Layer instance
The output of the green block. Has no. of channels equal to `filters`.
The size of the rest of the dimensions remains same as in `inp`.
"""
inp_res = Conv3D(
filters=filters,
kernel_size=(1, 1, 1),
strides=1,
data_format=data_format,
name=f'Res_{name}' if name else None)(inp)
# axis=1 for channels_first data format
# No. of groups = 8, as given in the paper
x = GroupNormalization(
groups=8,
axis=1 if data_format == 'channels_first' else 0,
name=f'GroupNorm_1_{name}' if name else None)(inp)
x = Activation('relu', name=f'Relu_1_{name}' if name else None)(x)
x = Conv3D(
filters=filters,
kernel_size=(3, 3, 3),
strides=1,
padding='same',
data_format=data_format,
name=f'Conv3D_1_{name}' if name else None)(x)
x = GroupNormalization(
groups=8,
axis=1 if data_format == 'channels_first' else 0,
name=f'GroupNorm_2_{name}' if name else None)(x)
x = Activation('relu', name=f'Relu_2_{name}' if name else None)(x)
x = Conv3D(
filters=filters,
kernel_size=(3, 3, 3),
strides=1,
padding='same',
data_format=data_format,
name=f'Conv3D_2_{name}' if name else None)(x)
out = Add(name=f'Out_{name}' if name else None)([x, inp_res])
return out
# From keras-team/keras/blob/master/examples/variational_autoencoder.py
def sampling(args):
"""Reparameterization trick by sampling from an isotropic unit Gaussian.
# Arguments
args (tensor): mean and log of variance of Q(z|X)
# Returns
z (tensor): sampled latent vector
"""
z_mean, z_var = args
batch = K.shape(z_mean)[0]
dim = K.int_shape(z_mean)[1]
# by default, random_normal has mean = 0 and std = 1.0
epsilon = K.random_normal(shape=(batch, dim))
return z_mean + K.exp(0.5 * z_var) * epsilon
def dice_coefficient(y_true, y_pred):
intersection = K.sum(K.abs(y_true * y_pred), axis=[-3,-2,-1])
dn = K.sum(K.square(y_true) + K.square(y_pred), axis=[-3,-2,-1]) + 1e-8
return K.mean(2 * intersection / dn, axis=[0,1])
def loss_gt(e=1e-8):
"""
loss_gt(e=1e-8)
------------------------------------------------------
Since keras does not allow custom loss functions to have arguments
other than the true and predicted labels, this function acts as a wrapper
that allows us to implement the custom loss used in the paper. This function
only calculates - L<dice> term of the following equation. (i.e. GT Decoder part loss)
L = - L<dice> + weight_L2 ∗ L<L2> + weight_KL ∗ L<KL>
Parameters
----------
`e`: Float, optional
A small epsilon term to add in the denominator to avoid dividing by
zero and possible gradient explosion.
Returns
-------
loss_gt_(y_true, y_pred): A custom keras loss function
This function takes as input the predicted and ground labels, uses them
to calculate the dice loss.
"""
def loss_gt_(y_true, y_pred):
intersection = K.sum(K.abs(y_true * y_pred), axis=[-3,-2,-1])
dn = K.sum(K.square(y_true) + K.square(y_pred), axis=[-3,-2,-1]) + e
return - K.mean(2 * intersection / dn, axis=[0,1])
return loss_gt_
def loss_VAE(input_shape, z_mean, z_var, weight_L2=0.1, weight_KL=0.1):
"""
loss_VAE(input_shape, z_mean, z_var, weight_L2=0.1, weight_KL=0.1)
------------------------------------------------------
Since keras does not allow custom loss functions to have arguments
other than the true and predicted labels, this function acts as a wrapper
that allows us to implement the custom loss used in the paper. This function
calculates the following equation, except for -L<dice> term. (i.e. VAE decoder part loss)
L = - L<dice> + weight_L2 ∗ L<L2> + weight_KL ∗ L<KL>
Parameters
----------
`input_shape`: A 4-tuple, required
The shape of an image as the tuple (c, H, W, D), where c is
the no. of channels; H, W and D is the height, width and depth of the
input image, respectively.
`z_mean`: An keras.layers.Layer instance, required
The vector representing values of mean for the learned distribution
in the VAE part. Used internally.
`z_var`: An keras.layers.Layer instance, required
The vector representing values of variance for the learned distribution
in the VAE part. Used internally.
`weight_L2`: A real number, optional
The weight to be given to the L2 loss term in the loss function. Adjust to get best
results for your task. Defaults to 0.1.
`weight_KL`: A real number, optional
The weight to be given to the KL loss term in the loss function. Adjust to get best
results for your task. Defaults to 0.1.
Returns
-------
loss_VAE_(y_true, y_pred): A custom keras loss function
This function takes as input the predicted and ground labels, uses them
to calculate the L2 and KL loss.
"""
def loss_VAE_(y_true, y_pred):
c, H, W, D = input_shape
n = c * H * W * D
loss_L2 = K.mean(K.square(y_true - y_pred), axis=(1, 2, 3, 4)) # original axis value is (1,2,3,4).
loss_KL = (1 / n) * K.sum(
K.exp(z_var) + K.square(z_mean) - 1. - z_var,
axis=-1
)
return weight_L2 * loss_L2 + weight_KL * loss_KL
return loss_VAE_
def build_model(input_shape=(4, 160, 192, 128), output_channels=3, weight_L2=0.1, weight_KL=0.1, dice_e=1e-8):
"""
build_model(input_shape=(4, 160, 192, 128), output_channels=3, weight_L2=0.1, weight_KL=0.1)
-------------------------------------------
Creates the model used in the BRATS2018 winning solution
by Myronenko A. (https://arxiv.org/pdf/1810.11654.pdf)
Parameters
----------
`input_shape`: A 4-tuple, optional.
Shape of the input image. Must be a 4D image of shape (c, H, W, D),
where, each of H, W and D are divisible by 2^4, and c is divisible by 4.
Defaults to the crop size used in the paper, i.e., (4, 160, 192, 128).
`output_channels`: An integer, optional.
The no. of channels in the output. Defaults to 3 (BraTS 2018 format).
`weight_L2`: A real number, optional
The weight to be given to the L2 loss term in the loss function. Adjust to get best
results for your task. Defaults to 0.1.
`weight_KL`: A real number, optional
The weight to be given to the KL loss term in the loss function. Adjust to get best
results for your task. Defaults to 0.1.
`dice_e`: Float, optional
A small epsilon term to add in the denominator of dice loss to avoid dividing by
zero and possible gradient explosion. This argument will be passed to loss_gt function.
Returns
-------
`model`: A keras.models.Model instance
The created model.
"""
c, H, W, D = input_shape
assert len(input_shape) == 4, "Input shape must be a 4-tuple"
assert (c % 4) == 0, "The no. of channels must be divisible by 4"
assert (H % 16) == 0 and (W % 16) == 0 and (D % 16) == 0, \
"All the input dimensions must be divisible by 16"
# -------------------------------------------------------------------------
# Encoder
# -------------------------------------------------------------------------
## Input Layer
inp = Input(input_shape)
## The Initial Block
x = Conv3D(
filters=32,
kernel_size=(3, 3, 3),
strides=1,
padding='same',
data_format='channels_first',
name='Input_x1')(inp)
## Dropout (0.2)
x = SpatialDropout3D(0.2, data_format='channels_first')(x)
## Green Block x1 (output filters = 32)
x1 = green_block(x, 32, name='x1')
x = Conv3D(
filters=32,
kernel_size=(3, 3, 3),
strides=2,
padding='same',
data_format='channels_first',
name='Enc_DownSample_32')(x1)
## Green Block x2 (output filters = 64)
x = green_block(x, 64, name='Enc_64_1')
x2 = green_block(x, 64, name='x2')
x = Conv3D(
filters=64,
kernel_size=(3, 3, 3),
strides=2,
padding='same',
data_format='channels_first',
name='Enc_DownSample_64')(x2)
## Green Blocks x2 (output filters = 128)
x = green_block(x, 128, name='Enc_128_1')
x3 = green_block(x, 128, name='x3')
x = Conv3D(
filters=128,
kernel_size=(3, 3, 3),
strides=2,
padding='same',
data_format='channels_first',
name='Enc_DownSample_128')(x3)
## Green Blocks x4 (output filters = 256)
x = green_block(x, 256, name='Enc_256_1')
x = green_block(x, 256, name='Enc_256_2')
x = green_block(x, 256, name='Enc_256_3')
x4 = green_block(x, 256, name='x4')
# -------------------------------------------------------------------------
# Decoder
# -------------------------------------------------------------------------
## GT (Groud Truth) Part
# -------------------------------------------------------------------------
### Green Block x1 (output filters=128)
x = Conv3D(
filters=128,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_GT_ReduceDepth_128')(x4)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_GT_UpSample_128')(x)
x = Add(name='Input_Dec_GT_128')([x, x3])
x = green_block(x, 128, name='Dec_GT_128')
### Green Block x1 (output filters=64)
x = Conv3D(
filters=64,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_GT_ReduceDepth_64')(x)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_GT_UpSample_64')(x)
x = Add(name='Input_Dec_GT_64')([x, x2])
x = green_block(x, 64, name='Dec_GT_64')
### Green Block x1 (output filters=32)
x = Conv3D(
filters=32,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_GT_ReduceDepth_32')(x)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_GT_UpSample_32')(x)
x = Add(name='Input_Dec_GT_32')([x, x1])
x = green_block(x, 32, name='Dec_GT_32')
### Blue Block x1 (output filters=32)
x = Conv3D(
filters=32,
kernel_size=(3, 3, 3),
strides=1,
padding='same',
data_format='channels_first',
name='Input_Dec_GT_Output')(x)
### Output Block
out_GT = Conv3D(
filters=output_channels, # No. of tumor classes is 3
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
activation='sigmoid',
name='Dec_GT_Output')(x)
## VAE (Variational Auto Encoder) Part
# -------------------------------------------------------------------------
### VD Block (Reducing dimensionality of the data)
x = GroupNormalization(groups=8, axis=1, name='Dec_VAE_VD_GN')(x4)
x = Activation('relu', name='Dec_VAE_VD_relu')(x)
x = Conv3D(
filters=16,
kernel_size=(3, 3, 3),
strides=2,
padding='same',
data_format='channels_first',
name='Dec_VAE_VD_Conv3D')(x)
# Not mentioned in the paper, but the author used a Flattening layer here.
x = Flatten(name='Dec_VAE_VD_Flatten')(x)
x = Dense(256, name='Dec_VAE_VD_Dense')(x)
### VDraw Block (Sampling)
z_mean = Dense(128, name='Dec_VAE_VDraw_Mean')(x)
z_var = Dense(128, name='Dec_VAE_VDraw_Var')(x)
x = Lambda(sampling, name='Dec_VAE_VDraw_Sampling')([z_mean, z_var])
### VU Block (Upsizing back to a depth of 256)
x = Dense((c//4) * (H//16) * (W//16) * (D//16))(x)
x = Activation('relu')(x)
x = Reshape(((c//4), (H//16), (W//16), (D//16)))(x)
x = Conv3D(
filters=256,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_VAE_ReduceDepth_256')(x)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_VAE_UpSample_256')(x)
### Green Block x1 (output filters=128)
x = Conv3D(
filters=128,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_VAE_ReduceDepth_128')(x)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_VAE_UpSample_128')(x)
x = green_block(x, 128, name='Dec_VAE_128')
### Green Block x1 (output filters=64)
x = Conv3D(
filters=64,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_VAE_ReduceDepth_64')(x)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_VAE_UpSample_64')(x)
x = green_block(x, 64, name='Dec_VAE_64')
### Green Block x1 (output filters=32)
x = Conv3D(
filters=32,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_VAE_ReduceDepth_32')(x)
x = UpSampling3D(
size=2,
data_format='channels_first',
name='Dec_VAE_UpSample_32')(x)
x = green_block(x, 32, name='Dec_VAE_32')
### Blue Block x1 (output filters=32)
x = Conv3D(
filters=32,
kernel_size=(3, 3, 3),
strides=1,
padding='same',
data_format='channels_first',
name='Input_Dec_VAE_Output')(x)
### Output Block
out_VAE = Conv3D(
filters=4,
kernel_size=(1, 1, 1),
strides=1,
data_format='channels_first',
name='Dec_VAE_Output')(x)
# Build and Compile the model
out = out_GT
model = Model(inp, outputs=[out, out_VAE]) # Create the model
model.compile(
adam(lr=1e-4),
[loss_gt(dice_e), loss_VAE(input_shape, z_mean, z_var, weight_L2=weight_L2, weight_KL=weight_KL)],
metrics=[dice_coefficient]
)
return model
