"""
This code is to build and train 2D U-Net
"""
import numpy as np
import sys
import subprocess
import argparse
import os
from keras.models import Model
from keras.layers import Input, Activation, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose, ZeroPadding2D, add
from keras.optimizers import Adam, SGD
from keras.callbacks import ModelCheckpoint, CSVLogger
from keras import backend as K
from keras import losses
import tensorflow as tf
import matplotlib.pyplot as plt
import pandas as pd
import csv
from utils import *
from data import load_train_data
K.set_image_data_format('channels_last') # Tensorflow dimension ordering
# ----- paths setting -----
data_path = sys.argv[1] + "/"
model_path = data_path + "models/"
log_path = data_path + "logs/"
# ----- params for training and testing -----
batch_size = 1
cur_fold = sys.argv[2]
plane = sys.argv[3]
epoch = int(sys.argv[4])
init_lr = float(sys.argv[5])
# ----- Dice Coefficient and cost function for training -----
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.0 * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
def dice_coef_loss(y_true, y_pred):
return -dice_coef(y_true, y_pred)
def get_unet((img_rows, img_cols), flt=64, pool_size=(2, 2, 2), init_lr=1.0e-5):
"""build and compile Neural Network"""
print "start building NN"
inputs = Input((img_rows, img_cols, 1))
conv1 = Conv2D(flt, (3, 3), activation='relu', padding='same')(inputs)
conv1 = Conv2D(flt, (3, 3), activation='relu', padding='same')(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
conv2 = Conv2D(flt*2, (3, 3), activation='relu', padding='same')(pool1)
conv2 = Conv2D(flt*2, (3, 3), activation='relu', padding='same')(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
conv3 = Conv2D(flt*4, (3, 3), activation='relu', padding='same')(pool2)
conv3 = Conv2D(flt*4, (3, 3), activation='relu', padding='same')(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
conv4 = Conv2D(flt*8, (3, 3), activation='relu', padding='same')(pool3)
conv4 = Conv2D(flt*8, (3, 3), activation='relu', padding='same')(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)
conv5 = Conv2D(flt*16, (3, 3), activation='relu', padding='same')(pool4)
conv5 = Conv2D(flt*8, (3, 3), activation='relu', padding='same')(conv5)
up6 = concatenate([Conv2DTranspose(flt*8, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=3)
conv6 = Conv2D(flt*8, (3, 3), activation='relu', padding='same')(up6)
conv6 = Conv2D(flt*4, (3, 3), activation='relu', padding='same')(conv6)
up7 = concatenate([Conv2DTranspose(flt*4, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=3)
conv7 = Conv2D(flt*4, (3, 3), activation='relu', padding='same')(up7)
conv7 = Conv2D(flt*2, (3, 3), activation='relu', padding='same')(conv7)
up8 = concatenate([Conv2DTranspose(flt*2, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=3)
conv8 = Conv2D(flt*2, (3, 3), activation='relu', padding='same')(up8)
conv8 = Conv2D(flt, (3, 3), activation='relu', padding='same')(conv8)
up9 = concatenate([Conv2DTranspose(flt, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=3)
conv9 = Conv2D(flt, (3, 3), activation='relu', padding='same')(up9)
conv9 = Conv2D(flt, (3, 3), activation='relu', padding='same')(conv9)
conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)
model = Model(inputs=[inputs], outputs=[conv10])
model.compile(optimizer=Adam(lr=init_lr), loss=dice_coef_loss, metrics=[dice_coef])
return model
def train(fold, plane, batch_size, nb_epoch,init_lr):
"""
train an Unet model with data from load_train_data()
Parameters
----------
fold : string
which fold is experimenting in 4-fold. It should be one of 0/1/2/3
plane : char
which plane is experimenting. It is from 'X'/'Y'/'Z'
batch_size : int
size of mini-batch
nb_epoch : int
number of epochs to train NN
init_lr : float
initial learning rate
"""
print "number of epoch: ", nb_epoch
print "learning rate: ", init_lr
# --------------------- load and preprocess training data -----------------
print '-'*80
print ' Loading and preprocessing train data...'
print '-'*80
imgs_train, imgs_mask_train = load_train_data(fold, plane)
imgs_row = imgs_train.shape[1]
imgs_col = imgs_train.shape[2]
imgs_train = preprocess(imgs_train)
imgs_mask_train = preprocess(imgs_mask_train)
imgs_train = imgs_train.astype('float32')
imgs_mask_train = imgs_mask_train.astype('float32')
# ---------------------- Create, compile, and train model ------------------------
print '-'*80
print ' Creating and compiling model...'
print '-'*80
model = get_unet((imgs_row, imgs_col), pool_size=(2, 2, 2), init_lr=init_lr)
print model.summary()
print '-'*80
print ' Fitting model...'
print '-'*80
ver = 'unet_fd%s_%s_ep%s_lr%s.csv'%(cur_fold, plane, epoch, init_lr)
csv_logger = CSVLogger(log_path + ver)
model_checkpoint = ModelCheckpoint(model_path + ver + ".h5",
monitor='loss',
save_best_only=False,
period=10)
history = model.fit(imgs_train, imgs_mask_train,
batch_size= batch_size, epochs= nb_epoch, verbose=1, shuffle=True,
callbacks=[model_checkpoint, csv_logger])
if __name__ == "__main__":
train(cur_fold, plane, batch_size, epoch, init_lr)
print "training done"
