import argparse
import numpy as np
import os
import sys
import warnings
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=FutureWarning)
import tensorflow as tf
import keras
from keras.models import Model
from keras.layers import Input,merge, concatenate, Conv2D, MaxPooling2D, Activation, UpSampling2D,Dropout,Conv2DTranspose,add,multiply
from keras.layers.normalization import BatchNormalization as bn
from keras.optimizers import RMSprop, Adam
from keras import regularizers, losses, backend as K
from keras.callbacks import EarlyStopping, ModelCheckpoint, CSVLogger, ModelCheckpoint, TensorBoard
os.environ['CUDA_VISIBLE_DEVICES']="0"
smooth = 1.
def dice_coef(y_true, y_pred):
y_true_f = K.flatten(y_true)
y_pred_f = K.flatten(y_pred)
intersection = K.sum(y_true_f * y_pred_f)
return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
def dice_coef_test(y_true, y_pred):
y_true_f = np.array(y_true).flatten()
y_pred_f =np.array(y_pred).flatten()
intersection = np.sum(y_true_f * y_pred_f)
return (2. * intersection + smooth) / (np.sum(y_true_f) + np.sum(y_pred_f) + smooth)
def dice_coef_loss(y_true, y_pred):
return -dice_coef(y_true, y_pred)
def neg_dice_coef_loss(y_true, y_pred):
return dice_coef(y_true, y_pred)
#define the model
def Comp_U_Net(input_shape,learn_rate=1e-3):
l2_lambda = 0.0002
DropP = 0.3
kernel_size=3
inputs = Input(input_shape,name='ip0')
conv0a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(inputs)
conv0a = bn()(conv0a)
conv0b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv0a)
conv0b = bn()(conv0b)
pool0 = MaxPooling2D(pool_size=(2, 2))(conv0b)
pool0 = Dropout(DropP)(pool0)
conv1a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(pool0)
conv1a = bn()(conv1a)
conv1b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv1a)
conv1b = bn()(conv1b)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1b)
pool1 = Dropout(DropP)(pool1)
conv2a = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(pool1)
conv2a = bn()(conv2a)
conv2b = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv2a)
conv2b = bn()(conv2b)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2b)
pool2 = Dropout(DropP)(pool2)
conv3a = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(pool2)
conv3a = bn()(conv3a)
conv3b = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv3a)
conv3b = bn()(conv3b)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3b)
pool3 = Dropout(DropP)(pool3)
conv4a = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(pool3)
conv4a = bn()(conv4a)
conv4b = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv4a)
conv4b = bn()(conv4b)
pool4 = MaxPooling2D(pool_size=(2, 2))(conv4b)
pool4 = Dropout(DropP)(pool4)
conv5a = Conv2D(512, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(pool4)
conv5a = bn()(conv5a)
conv5b = Conv2D(512, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv5a)
conv5b = bn()(conv5b)
up6 = concatenate([Conv2DTranspose(256,(2, 2), strides=(2, 2), padding='same')(conv5b), (conv4b)],name='up6', axis=3)
up6 = Dropout(DropP)(up6)
conv6a = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(up6)
conv6a = bn()(conv6a)
conv6b = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv6a)
conv6b = bn()(conv6b)
up7 = concatenate([Conv2DTranspose(128,(2, 2), strides=(2, 2), padding='same')(conv6b),(conv3b)],name='up7', axis=3)
up7 = Dropout(DropP)(up7)
#add second output here
conv7a = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(up7)
conv7a = bn()(conv7a)
conv7b = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv7a)
conv7b = bn()(conv7b)
up8 = concatenate([Conv2DTranspose(64,(2, 2), strides=(2, 2), padding='same')(conv7b), (conv2b)],name='up8', axis=3)
up8 = Dropout(DropP)(up8)
conv8a = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(up8)
conv8a = bn()(conv8a)
conv8b = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv8a)
conv8b = bn()(conv8b)
up9 = concatenate([Conv2DTranspose(32,(2, 2), strides=(2, 2), padding='same')(conv8b),(conv1b)],name='up9',axis=3)
conv9a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(up9)
conv9a = bn()(conv9a)
conv9b = Conv2D(12, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv9a)
conv9b = bn()(conv9b)
up10 = concatenate([Conv2DTranspose(32,(2, 2), strides=(2, 2), padding='same')(conv9b),(conv0b)],name='up10',axis=3)
conv10a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(up10)
conv10a = bn()(conv10a)
conv10b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(conv10a)
conv10b = bn()(conv10b)
final_op=Conv2D(1, (1, 1), activation='sigmoid',name='final_op')(conv10b)
#----------------------------------------------------------------------------------------------------------------------------------
#second branch - brain
xup6 = concatenate([Conv2DTranspose(256,(2, 2), strides=(2, 2), padding='same')(conv5b), (conv4b)],name='xup6', axis=3)
xup6 = Dropout(DropP)(xup6)
xconv6a = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xup6)
xconv6a = bn()(xconv6a)
xconv6b = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xconv6a)
xconv6b = bn()(xconv6b)
xup7 = concatenate([Conv2DTranspose(128,(2, 2), strides=(2, 2), padding='same')(xconv6b),(conv3b)],name='xup7', axis=3)
xup7 = Dropout(DropP)(xup7)
xconv7a = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xup7)
xconv7a = bn()(xconv7a)
xconv7b = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xconv7a)
xconv7b = bn()(xconv7b)
xup8 = concatenate([Conv2DTranspose(64,(2, 2), strides=(2, 2), padding='same')(xconv7b),(conv2b)],name='xup8', axis=3)
xup8 = Dropout(DropP)(xup8)
#add third xoutxout here
xconv8a = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xup8)
xconv8a = bn()(xconv8a)
xconv8b = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xconv8a)
xconv8b = bn()(xconv8b)
xup9 = concatenate([Conv2DTranspose(32,(2, 2), strides=(2, 2), padding='same')(xconv8b), (conv1b)],name='xup9',axis=3)
xup9 = Dropout(DropP)(xup9)
xconv9a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xup9)
xconv9a = bn()(xconv9a)
xconv9b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xconv9a)
xconv9b = bn()(xconv9b)
xup10 = concatenate([Conv2DTranspose(32,(2, 2), strides=(2, 2), padding='same')(xconv9b), (conv0b)],name='xup10',axis=3)
xup10 = Dropout(DropP)(xup10)
xconv10a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xup10)
xconv10a = bn()(xconv10a)
xconv10b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(xconv10a)
xconv10b = bn()(xconv10b)
xfinal_op=Conv2D(1, (1, 1), activation='sigmoid',name='xfinal_op')(xconv10b)
#-----------------------------third branch
#Concatenation fed to the reconstruction layer of all 3
x_u_net_op0=keras.layers.concatenate([final_op,xfinal_op,keras.layers.add([final_op,xfinal_op])],name='res_a')
res_1_conv0a = Conv2D( 32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(x_u_net_op0)
res_1_conv0a = bn()(res_1_conv0a)
res_1_conv0b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv0a)
res_1_conv0b = bn()(res_1_conv0b)
res_1_pool0 = MaxPooling2D(pool_size=(2, 2))(res_1_conv0b)
res_1_pool0 = Dropout(DropP)(res_1_pool0)
res_1_conv1a = Conv2D( 32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_pool0)
res_1_conv1a = bn()(res_1_conv1a)
res_1_conv1b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv1a)
res_1_conv1b = bn()(res_1_conv1b)
res_1_pool1 = MaxPooling2D(pool_size=(2, 2))(res_1_conv1b)
res_1_pool1 = Dropout(DropP)(res_1_pool1)
res_1_conv2a = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_pool1)
res_1_conv2a = bn()(res_1_conv2a)
res_1_conv2b = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv2a)
res_1_conv2b = bn()(res_1_conv2b)
res_1_pool2 = MaxPooling2D(pool_size=(2, 2))(res_1_conv2b)
res_1_pool2 = Dropout(DropP)(res_1_pool2)
res_1_conv3a = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_pool2)
res_1_conv3a = bn()(res_1_conv3a)
res_1_conv3b = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv3a)
res_1_conv3b = bn()(res_1_conv3b)
res_1_pool3 = MaxPooling2D(pool_size=(2, 2))(res_1_conv3b)
res_1_pool3 = Dropout(DropP)(res_1_pool3)
res_1_conv4a = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_pool3)
res_1_conv4a = bn()(res_1_conv4a)
res_1_conv4b = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv4a)
res_1_conv4b = bn()(res_1_conv4b)
res_1_pool4 = MaxPooling2D(pool_size=(2, 2))(res_1_conv4b)
res_1_pool4 = Dropout(DropP)(res_1_pool4)
res_1_conv5a = Conv2D(512, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_pool4)
res_1_conv5a = bn()(res_1_conv5a)
res_1_conv5b = Conv2D(512, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv5a)
res_1_conv5b = bn()(res_1_conv5b)
res_1_up6 = concatenate([Conv2DTranspose(256,(2, 2), strides=(2, 2), padding='same')(res_1_conv5b), (res_1_conv4b)],name='res_1_up6', axis=3)
res_1_up6 = Dropout(DropP)(res_1_up6)
res_1_conv6a = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_up6)
res_1_conv6a = bn()(res_1_conv6a)
res_1_conv6b = Conv2D(256, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv6a)
res_1_conv6b = bn()(res_1_conv6b)
res_1_up7 = concatenate([Conv2DTranspose(128,(2, 2), strides=(2, 2), padding='same')(res_1_conv6b),(res_1_conv3b)],name='res_1_up7', axis=3)
res_1_up7 = Dropout(DropP)(res_1_up7)
#add second res_1_output here
res_1_conv7a = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_up7)
res_1_conv7a = bn()(res_1_conv7a)
res_1_conv7b = Conv2D(128, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv7a)
res_1_conv7b = bn()(res_1_conv7b)
res_1_up8 = concatenate([Conv2DTranspose(64,(2, 2), strides=(2, 2), padding='same')(res_1_conv7b),(res_1_conv2b)],name='res_1_up8', axis=3)
res_1_up8 = Dropout(DropP)(res_1_up8)
#add third outout here
res_1_conv8a = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_up8)
res_1_conv8a = bn()(res_1_conv8a)
res_1_conv8b = Conv2D(64, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv8a)
res_1_conv8b = bn()(res_1_conv8b)
res_1_up9 = concatenate([Conv2DTranspose(32,(2, 2), strides=(2, 2), padding='same')(res_1_conv8b), (res_1_conv1b)],name='res_1_up9',axis=3)
res_1_up9 = Dropout(DropP)(res_1_up9)
res_1_conv9a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_up9)
res_1_conv9a = bn()(res_1_conv9a)
res_1_conv9b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv9a)
res_1_conv9b = bn()(res_1_conv9b)
res_1_up10 = concatenate([Conv2DTranspose(32,(2, 2), strides=(2, 2), padding='same')(res_1_conv9b),(res_1_conv0b)],name='res_1_up10',axis=3)
res_1_up10 = Dropout(DropP)(res_1_up10)
res_1_conv10a = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_up10)
res_1_conv10a = bn()(res_1_conv10a)
res_1_conv10b = Conv2D(32, (kernel_size, kernel_size), activation='relu', padding='same',
kernel_regularizer=regularizers.l2(l2_lambda) )(res_1_conv10a)
res_1_conv10b = bn()(res_1_conv10b)
res_1_final_op=Conv2D(1, (1, 1), activation='sigmoid',name='res_1_final_op')(res_1_conv10b)
model=Model(inputs=[inputs],outputs=[final_op,
xfinal_op,
res_1_final_op,
])
model.compile(optimizer=keras.optimizers.Adam(lr=1e-5),loss={'final_op':dice_coef_loss,
'xfinal_op':neg_dice_coef_loss,
'res_1_final_op':'mse'})
return model
#----------------------------------------------------Main--------------------------------------------------#
def train_model(data_params, train_params, common_params):
training_data_folder = data_params['data_dir'].rstrip('/')
train_x = training_data_folder + '/' + data_params['train_data_file']
train_y = training_data_folder + '/' + data_params['train_label_file']
model = Comp_U_Net(input_shape=(256,256,1), learn_rate=train_params['learning_rate'])
# print(model.summary())
x_train = np.load(train_x)
y_train = np.load(train_y)
x_train=x_train.reshape(x_train.shape+(1,))
y_train=y_train.reshape(y_train.shape+(1,))
# Log output
print ("Training dwi volume shape: ", x_train.shape)
print ("Training dwi mask volume shape: ", y_train.shape)
view = train_params['principal_axis']
os.makedirs(common_params['log_dir'], exist_ok= True)
csv_logger = CSVLogger(common_params['log_dir'] + '/' + view + '.csv', append=True, separator=';')
# checkpoint
os.makedirs(common_params['save_model_dir'], exist_ok= True)
filepath = common_params['save_model_dir'] + "/weights-" + view + "-improvement-{epoch:02d}.h5"
checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=False, save_weights_only=True)
# Trains the model for a given number of epochs (iterations on a dataset).
history_callback = model.fit([x_train],
[y_train,y_train,y_train],
validation_split=train_params['validation_split'],
batch_size=train_params['train_batch_size'],
epochs=train_params['num_epochs'],
shuffle=train_params['shuffle_data'],
verbose=1,
callbacks=[csv_logger, checkpoint])
import h5py
# serialize model to JSON
model_json = model.to_json()
with open(common_params['save_model_dir'] + "/CompNetBasicModel.json", "w") as json_file:
json_file.write(model_json)
# serialize weights to HDF5
model.save_weights(common_params['save_model_dir'] + "/" + view + "-compnet_final_weight.h5")
print("Saved model to disk location: ", common_params['save_model_dir'])
