<reponame>kft-upc/image_DL<filename>unet/train.py
from __future__ import print_function

import os
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from keras.models import Model
from keras.layers import Input, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint
from keras import backend as K
from keras.utils import multi_gpu_model,plot_model
import warnings
from matplotlib import pyplot as plt

from data import load_train_data, load_test_data

K.set_image_data_format('channels_last')  # TF dimension ordering in this code

img_cols = 256
img_rows = 256


smooth = 0

save_checkpoint = os.path.join(os.getcwd(),'weights','weights_far_1.h5')
weights_path = os.path.join(os.getcwd(),'weights','weights_nerve_1.h5')

model_save = os.path.join(os.getcwd(),'models','model_far_1.png')

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_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)


def get_unet():
    inputs = Input((img_rows, img_cols, 1))
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(inputs)
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool1)
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool2)
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool3)
    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool4)
    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv5)

    up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=3)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up6)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv6)

    up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=3)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up7)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv7)

    up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=3)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up8)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv8)

    up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=3)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up9)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv9)

    conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)  

    model = Model(inputs=[inputs], outputs=[conv10])
    
    model.load_weights(weights_path,by_name=True)
    
    for i,layer in enumerate(model.layers):
        print(i,layer.name)
        
    plot_model(model,to_file=model_save,show_shapes = True)
    
    parallele_model = multi_gpu_model(model,gpus=2)    

    parallele_model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[dice_coef])

    return parallele_model

def preprocess(imgs):  
        imgs_p = np.ndarray((imgs.shape[0],img_rows,img_cols), dtype = np.uint8)
        for i in range(imgs.shape[0]):
            imgs_p[i] = resize(imgs[i], (img_cols, img_rows), preserve_range = True)
            
        imgs_p = imgs_p[..., np.newaxis]
        return imgs_p

def train():
    
    print('-'*30)
    print('loading and preprocessing train data...')
    print('-'*30)
    imgs_train,imgs_mask_train = load_train_data()
    imgs_train = preprocess(imgs_train)
    imgs_mask_train = preprocess(imgs_mask_train)
    
    imgs_train = imgs_train.astype('float32')
    #mean = np.mean(imgs_train)
    #std = np.std(imgs_train)
    
    #imgs_train -= mean
    imgs_train /= 255
    
    imgs_mask_train = imgs_mask_train.astype('float32')
    imgs_mask_train /=255
    
    print('-'*30)
    print('Creating and compiling model...')
    print('-'*30)
    model = get_unet()
    model_checkpoint = ModelCheckpoint(save_checkpoint,monitor='val_loss', save_best_only=True)

    print('-'*30)
    print('Fitting model...')
    print('-'*30)
      
    
    hist = model.fit(imgs_train, imgs_mask_train, batch_size=32, nb_epoch=20, verbose=1, shuffle=True,
                    validation_split=0.2,
                    callbacks=[model_checkpoint])
    
    return hist
    
    
def predict():
    print('-'*30)
    print('Loading and preprocessing test data...')
    print('-'*30)
    imgs_test ,imgs_id_test = load_test_data()
    imgs_test = preprocess(imgs_test)
    
    imgs_test = imgs_test.astype('float32')
    mean = np.mean(imgs_test)
    std = np.std(imgs_test)
    
    imgs_test -= mean
    imgs_test /= std
    
    pred_dir = os.path.join(os.getcwd(),'preds','ultra')
    if not os.path.exists(pred_dir):
        os.mkdir(pred_dir)
        
    model = get_unet()

    print('-'*30)
    print('Loading saved weights...')
    print('-'*30)
    model.load_weights(save_checkpoint)

    print('-'*30)
    print('Predicting masks on test data...')
    print('-'*30)
    imgs_mask_test = model.predict(imgs_test,verbose=1)    
    np.save('imgs_mask_test.npy',imgs_mask_test)

    print('-' * 30)
    print('Saving predicted masks to files...')
    print('-' * 30)       
    for image, image_id in zip(imgs_mask_test, imgs_id_test):
        image = (image[:,:,0] *255.).astype('np.uint8')
        imsave(os.path.join(pred_dir,str(image_id) + '_pred.png'),image)

def resvisual(history):
    
    plt.figure(figsize=(6,12))
    plt.subplot(211)
    plt.plot(history.history['val_dice_coef'])
    plt.title('model dice_coef')
    plt.ylabel('dice_coef')
    plt.xlabel('epoch')
    
    plt.subplot(212)
    plt.plot(history.history['val_loss'])
    plt.title('model loss')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    
    plt.show()
        


if __name__ == '__main__':
    warnings.filterwarnings('ignore')
    hist = train()
    resvisual(hist)
    #predict()

from __future__ import print_function

import os
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from keras.models import Model
from keras.layers import Input, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint
from keras import backend as K
from keras.utils import multi_gpu_model,plot_model
import warnings
from matplotlib import pyplot as plt

from data import load_train_data, load_test_data

K.set_image_data_format('channels_last')  # TF dimension ordering in this code

img_cols = 256
img_rows = 256


smooth = 0

save_checkpoint = os.path.join(os.getcwd(),'weights','weights_far_1.h5')
weights_path = os.path.join(os.getcwd(),'weights','weights_nerve_1.h5')

model_save = os.path.join(os.getcwd(),'models','model_far_1.png')

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_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)


def get_unet():
    inputs = Input((img_rows, img_cols, 1))
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(inputs)
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool1)
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool2)
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool3)
    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(pool4)
    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv5)

    up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=3)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up6)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv6)

    up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=3)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up7)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv7)

    up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=3)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up8)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv8)

    up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=3)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(up9)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same',kernel_initializer='he_normal')(conv9)

    conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)  

    model = Model(inputs=[inputs], outputs=[conv10])
    
    model.load_weights(weights_path,by_name=True)
    
    for i,layer in enumerate(model.layers):
        print(i,layer.name)
        
    plot_model(model,to_file=model_save,show_shapes = True)
    
    parallele_model = multi_gpu_model(model,gpus=2)    

    parallele_model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[dice_coef])

    return parallele_model

def preprocess(imgs):  
        imgs_p = np.ndarray((imgs.shape[0],img_rows,img_cols), dtype = np.uint8)
        for i in range(imgs.shape[0]):
            imgs_p[i] = resize(imgs[i], (img_cols, img_rows), preserve_range = True)
            
        imgs_p = imgs_p[..., np.newaxis]
        return imgs_p

def train():
    
    print('-'*30)
    print('loading and preprocessing train data...')
    print('-'*30)
    imgs_train,imgs_mask_train = load_train_data()
    imgs_train = preprocess(imgs_train)
    imgs_mask_train = preprocess(imgs_mask_train)
    
    imgs_train = imgs_train.astype('float32')
    #mean = np.mean(imgs_train)
    #std = np.std(imgs_train)
    
    #imgs_train -= mean
    imgs_train /= 255
    
    imgs_mask_train = imgs_mask_train.astype('float32')
    imgs_mask_train /=255
    
    print('-'*30)
    print('Creating and compiling model...')
    print('-'*30)
    model = get_unet()
    model_checkpoint = ModelCheckpoint(save_checkpoint,monitor='val_loss', save_best_only=True)

    print('-'*30)
    print('Fitting model...')
    print('-'*30)
      
    
    hist = model.fit(imgs_train, imgs_mask_train, batch_size=32, nb_epoch=20, verbose=1, shuffle=True,
                    validation_split=0.2,
                    callbacks=[model_checkpoint])
    
    return hist
    
    
def predict():
    print('-'*30)
    print('Loading and preprocessing test data...')
    print('-'*30)
    imgs_test ,imgs_id_test = load_test_data()
    imgs_test = preprocess(imgs_test)
    
    imgs_test = imgs_test.astype('float32')
    mean = np.mean(imgs_test)
    std = np.std(imgs_test)
    
    imgs_test -= mean
    imgs_test /= std
    
    pred_dir = os.path.join(os.getcwd(),'preds','ultra')
    if not os.path.exists(pred_dir):
        os.mkdir(pred_dir)
        
    model = get_unet()

    print('-'*30)
    print('Loading saved weights...')
    print('-'*30)
    model.load_weights(save_checkpoint)

    print('-'*30)
    print('Predicting masks on test data...')
    print('-'*30)
    imgs_mask_test = model.predict(imgs_test,verbose=1)    
    np.save('imgs_mask_test.npy',imgs_mask_test)

    print('-' * 30)
    print('Saving predicted masks to files...')
    print('-' * 30)       
    for image, image_id in zip(imgs_mask_test, imgs_id_test):
        image = (image[:,:,0] *255.).astype('np.uint8')
        imsave(os.path.join(pred_dir,str(image_id) + '_pred.png'),image)

def resvisual(history):
    
    plt.figure(figsize=(6,12))
    plt.subplot(211)
    plt.plot(history.history['val_dice_coef'])
    plt.title('model dice_coef')
    plt.ylabel('dice_coef')
    plt.xlabel('epoch')
    
    plt.subplot(212)
    plt.plot(history.history['val_loss'])
    plt.title('model loss')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    
    plt.show()
        


if __name__ == '__main__':
    warnings.filterwarnings('ignore')
    hist = train()
    resvisual(hist)
    #predict()
# TF dimension ordering in this code #mean = np.mean(imgs_train) #std = np.std(imgs_train) #imgs_train -= mean #predict()