#
# Script which predicts segmented image of kidney tumor of CT image.
# To successfully run the script, you need to have installed following packages:
                                                                                # tensorflow 2.0
                                                                                # Keras
                                                                                # numpy
                                                                                # opencv-python
                                                                                # matplotlib
# This script can be run in the command window: python model_prediction.py
# Author: Hedvika Urbanova



from tensorflow.keras.models import *
from tensorflow.keras import backend as K
import numpy as np
import cv2
import matplotlib.pyplot as plt


def weighted_bce(y_true, y_pred):
    weights = (y_true * 59.) + 1.
    bce = K.binary_crossentropy(y_true, y_pred)
    weighted_bce = K.mean(bce * weights)
    return weighted_bce


complete_path = 'model.h5'
model = load_model(complete_path, custom_objects={'weighted_bce': weighted_bce})
model.summary()

Xtest = cv2.imread('test.png', 0) / 255
Xtest_predict = np.expand_dims(Xtest, axis=0)
Xtest_predict = np.expand_dims(Xtest_predict, axis=3)

Ytest = cv2.imread('expectedResult.png', 0) / 255

predikce = model.predict(Xtest_predict)
predikce[predikce < 0.5] = 0
predikce[predikce >= 0.5] = 1
predikce = np.squeeze(predikce, axis=0)
predikce = np.squeeze(predikce, axis=2)

combined = np.stack((predikce, Ytest, predikce), axis=2)
pom = np.copy(Ytest)

Ytest = np.append(np.atleast_3d(Ytest), np.atleast_3d(pom), axis=2)
Ytest = np.append(np.atleast_3d(Ytest), np.atleast_3d(pom), axis=2)

pom = np.copy(predikce)
predikce = np.append(np.atleast_3d(predikce), np.atleast_3d(pom), axis=2)
predikce = np.append(np.atleast_3d(predikce), np.atleast_3d(pom), axis=2)

pom = np.copy(Xtest)
Xtest = np.append(np.atleast_3d(Xtest), np.atleast_3d(pom), axis=2)
Xtest = np.append(np.atleast_3d(Xtest), np.atleast_3d(pom), axis=2)

combined = Xtest + combined
plt.figure(1)
plt.subplot(2, 2, 1)
plt.imshow(Xtest)
plt.title('Original image')
plt.axis('off')
plt.subplot(2, 2, 2)
plt.imshow(predikce)
plt.title('Predicted image')
plt.axis('off')
plt.subplot(2, 2, 3)
plt.imshow(Ytest)
plt.title('Expected result')
plt.axis('off')
plt.subplot(2, 2, 4)
plt.imshow(combined)
plt.title('Combined images')
plt.axis('off')
plt.show()
