import numpy as np
from keras.models import Model,load_model
from keras.layers.advanced_activations import PReLU
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.layers import Dropout,GaussianNoise, Input,Activation
from keras.layers.normalization import BatchNormalization
from keras.layers import Conv2DTranspose,UpSampling2D,concatenate,add
from keras.optimizers import SGD
import keras.backend as K
from losses import *
K.set_image_data_format("channels_last")
#u-net model
class Unet_model(object):
def __init__(self,img_shape,load_model_weights=None):
self.img_shape=img_shape
self.load_model_weights=load_model_weights
self.model =self.compile_unet()
def compile_unet(self):
"""
compile the U-net model
"""
i = Input(shape=self.img_shape)
#add gaussian noise to the first layer to combat overfitting
i_=GaussianNoise(0.01)(i)
i_ = Conv2D(64, 2, padding='same',data_format = 'channels_last')(i_)
out=self.unet(inputs=i_)
model = Model(input=i, output=out)
sgd = SGD(lr=0.08, momentum=0.9, decay=5e-6, nesterov=False)
model.compile(loss=gen_dice_loss, optimizer=sgd, metrics=[dice_whole_metric,dice_core_metric,dice_en_metric])
#load weights if set for prediction
if self.load_model_weights is not None:
model.load_weights(self.load_model_weights)
return model
def unet(self,inputs, nb_classes=4, start_ch=64, depth=3, inc_rate=2. ,activation='relu', dropout=0.0, batchnorm=True, upconv=True,format_='channels_last'):
"""
the actual u-net architecture
"""
o = self.level_block(inputs,start_ch, depth, inc_rate,activation, dropout, batchnorm, upconv,format_)
o = BatchNormalization()(o)
#o = Activation('relu')(o)
o=PReLU(shared_axes=[1, 2])(o)
o = Conv2D(nb_classes, 1, padding='same',data_format = format_)(o)
o = Activation('softmax')(o)
return o
def level_block(self,m, dim, depth, inc, acti, do, bn, up,format_="channels_last"):
if depth > 0:
n = self.res_block_enc(m,0.0,dim,acti, bn,format_)
#using strided 2D conv for donwsampling
m = Conv2D(int(inc*dim), 2,strides=2, padding='same',data_format = format_)(n)
m = self.level_block(m,int(inc*dim), depth-1, inc, acti, do, bn, up )
if up:
m = UpSampling2D(size=(2, 2),data_format = format_)(m)
m = Conv2D(dim, 2, padding='same',data_format = format_)(m)
else:
m = Conv2DTranspose(dim, 3, strides=2,padding='same',data_format = format_)(m)
n=concatenate([n,m])
#the decoding path
m = self.res_block_dec(n, 0.0,dim, acti, bn, format_)
else:
m = self.res_block_enc(m, 0.0,dim, acti, bn, format_)
return m
def res_block_enc(self,m, drpout,dim,acti, bn,format_="channels_last"):
"""
the encoding unit which a residual block
"""
n = BatchNormalization()(m) if bn else n
#n= Activation(acti)(n)
n=PReLU(shared_axes=[1, 2])(n)
n = Conv2D(dim, 3, padding='same',data_format = format_)(n)
n = BatchNormalization()(n) if bn else n
#n= Activation(acti)(n)
n=PReLU(shared_axes=[1, 2])(n)
n = Conv2D(dim, 3, padding='same',data_format =format_ )(n)
n=add([m,n])
return n
def res_block_dec(self,m, drpout,dim,acti, bn,format_="channels_last"):
"""
the decoding unit which a residual block
"""
n = BatchNormalization()(m) if bn else n
#n= Activation(acti)(n)
n=PReLU(shared_axes=[1, 2])(n)
n = Conv2D(dim, 3, padding='same',data_format = format_)(n)
n = BatchNormalization()(n) if bn else n
#n= Activation(acti)(n)
n=PReLU(shared_axes=[1, 2])(n)
n = Conv2D(dim, 3, padding='same',data_format =format_ )(n)
Save = Conv2D(dim, 1, padding='same',data_format = format_,use_bias=False)(m)
n=add([Save,n])
return n
