import keras
from keras.models import Model,Sequential
from keras.layers import Input,Dense, Dropout, Activation, Flatten, Reshape, Conv2D, MaxPooling2D, AveragePooling2D
from keras import regularizers
from keras.losses import mean_squared_error
from keras import losses
import matplotlib.patches as patches
import numpy as np
import dicom
import cv2
from utils import *
from train_cnn import run
def open_data_AE(X_fullsize, y_pred, contour_mask):
"""
Open dataset from the output of the CNN and
unroll it as 64*64 = vector of 4096 elements
:param X_fullsize: full size training set (256x256)
:param y_pred: CNN output
:return: input AE, output
"""
input_AE = []
contour_experts = []
for j in range(y_pred.shape[0]):
in_AE = cv2.resize(compute_roi_pred(X_fullsize, y_pred, contour_mask, j, roi_shape=32)[0],(64 , 64))
contour = cv2.resize(compute_roi_pred(X_fullsize, y_pred, contour_mask, j)[2], (64,64), interpolation = cv2.INTERSECT_NONE)
input_AE.append(in_AE)
contour_experts.append(contour)
return np.array(input_AE).reshape((-1, 64*64)), np.array(contour_experts).reshape((-1, 64*64))
def customized_loss(y_true, y_pred, alpha=0.0001, beta=3):
"""
Create a customized loss for the stacked AE.
Linear combination of MSE and KL divergence.
"""
#customize your own loss components
loss1 = losses.mean_absolute_error(y_true, y_pred)
loss2 = losses.kullback_leibler_divergence(y_true, y_pred)
#adjust the weight between loss components
return (alpha/2) * loss1 + beta * loss2
def model1(X_train, param_reg, get_history=False, verbose=0,loss = "customized_loss"):
"""
First part of the stacked AE.
Train the AE on the ROI input images.
:param X_train: ROI input image
:param get_history: boolean to return the loss history
:param loss: if "customized_loss" -> customized_loss
:return: encoded ROI image
"""
autoencoder_0 = Sequential()
encoder_0 = Dense(input_dim=4096, units=100, kernel_regularizer=regularizers.l2(param_reg))
decoder_0 = Dense(input_dim=100, units=4096, kernel_regularizer=regularizers.l2(param_reg))
autoencoder_0.add(encoder_0)
autoencoder_0.add(decoder_0)
if (loss == "customized_loss"):
loss = customized_loss
autoencoder_0.compile(loss = loss,optimizer='adam', metrics=['accuracy'])
h = autoencoder_0.fit(X_train, X_train, epochs=100, verbose=verbose)
temp_0 = Sequential()
temp_0.add(encoder_0)
temp_0.compile(loss=loss, optimizer='adam', metrics=['accuracy'])
encoded_X = temp_0.predict(X_train, verbose=0)
if get_history:
return h.history['loss'], encoded_X, encoder_0
else:
return encoded_X, encoder_0
def model2(encoder_0, encoded_X, X_train, param_reg, get_history=False, verbose=0,loss = "customized_loss"):
"""
Second part of the stacked AE.
:param X_train: encoder ROI image
:param get_history: boolean to return the loss history
:return: encoding layer
"""
autoencoder_1 = Sequential()
encoder_1 = Dense(input_dim=100, units=100, kernel_regularizer=regularizers.l2(param_reg))
decoder_1 = Dense(input_dim=100, units=100, kernel_regularizer=regularizers.l2(param_reg))
autoencoder_1.add(encoder_1)
autoencoder_1.add(decoder_1)
if (loss == "customized_loss"):
loss = customized_loss
autoencoder_1.compile(loss= loss, optimizer='adam', metrics=['accuracy'])
h = autoencoder_1.fit(encoded_X, encoded_X, epochs=100, verbose=verbose)
temp_0 = Sequential()
temp_0.add(encoder_0)
temp_0.compile(loss= loss, optimizer='adam', metrics=['accuracy'])
encoded_X = temp_0.predict(X_train, verbose=0)
if get_history:
return h.history['loss'], encoder_1
else:
return encoder_1
def model3(X_train, Y_train, encoder_0, encoder_1, init, param_reg,
get_history=False, verbose=0,loss = "MSE"):
"""
Last part of the stacked AE.
:param X_train: ROI input image
:param init: set the initial kernel weights (None for uniform)
:param get_history: boolean to return the loss history
:return: final model
"""
model = Sequential()
model.add(encoder_0)
model.add(encoder_1)
model.add(Dense(input_dim=100, units=4096, kernel_initializer=init, kernel_regularizer=regularizers.l2(param_reg)))
if (loss == "customized_loss"):
loss = customized_loss
model.compile(optimizer = 'adam', loss = loss, metrics=['accuracy'])
h = model.fit(X_train, Y_train, epochs=20, verbose=verbose)
if get_history:
return h.history['loss'], model
else:
return model
def run(X_fullsize, y_pred, contour_mask, X_to_pred=None, verbose=0, param_reg=3*0.001, init_3='zero',history=False ,loss1 = "customized_loss",loss2 = "customized_loss",loss3 = "MSE"):
"""
Full pipeline for CNN: load the dataset, train the model and predict ROIs
:param X_fullsize: full size training set (256x256)
:param y_pred: CNN output for training
:param X_to_pred: input for predictions after training (X_train if not specified)
:param verbose: int for verbose
:return: X_train, Y_train, contours_pred
"""
X_train, Y_train = open_data_AE(X_fullsize, y_pred, contour_mask)
encoded_X, encoder_0 = model1(X_train, param_reg=param_reg,loss = loss1)
encoder_1 = model2(encoder_0, encoded_X, X_train, param_reg=param_reg,loss = loss2)
h, model = model3(X_train, Y_train, encoder_0, encoder_1, param_reg=param_reg,
init=init_3, get_history=True, verbose=verbose,loss = loss3)
if not X_to_pred:
X_to_pred = X_train
contours = model.predict(X_to_pred)
binarys = np.array([cv2.threshold(contour, 0, 1, cv2.INTERSECT_NONE)[1].reshape((64,64)) for contour in contours])
if history:
return X_train, Y_train, binarys, model, h
else:
return X_train, Y_train, binarys, model
def inference(X_fullsize, y_pred, contour_mask, model):
X_test, Y_test = open_data_AE(X_fullsize, y_pred, contour_mask)
contours = model.predict(X_test)
binarys = np.array([cv2.threshold(contour, 0, 1, cv2.INTERSECT_NONE)[1].reshape((64,64)) for contour in contours])
return X_test, Y_test, binarys
