EvaluationWindow.py

import ntpath

import PySimpleGUI as sg

import json

from lovasz_losses_tf import *
import tensorflow
import nibabel as nib
import numpy as np
from glob import glob
import os
import cv2

sg.theme('Dark Purple 3')

layout = [[sg.Text('Select Inputs')],
          [sg.Text('Model.h5', justification='right', size=(15, 1)), sg.Input(key='Model'), sg.FileBrowse()],
          [sg.Text('Eval Image Folder', size=(15, 1), auto_size_text=False, justification='right'),
           sg.InputText(key='inputFolder', enable_events=True), sg.FolderBrowse(key='inputFolder2')],
          [sg.Text('MedML.json', justification='right', size=(15, 1)), sg.Input(key='MedMLjson'), sg.FileBrowse()],
          [sg.Submit(key='Submit'), sg.Cancel()]]

window = sg.Window('EvalMedML', layout, default_element_size=(40, 1))

while True:
    event, values = window.read()

    # print(event, values)
    if event is None:  # always,  always give a way out!
        break

    #extract filename from path
    def path_leaf(path):
        head, tail = ntpath.split(path)
        return tail or ntpath.basename(head)

    if event == 'Submit':
        # MedML.json
        f = open(path_leaf(values['MedMLjson']))
        MedML = json.load(f)

        def dice_loss(y_true, y_pred):
            numerator = 2 * tensorflow.reduce_sum(y_true * y_pred, axis=-1)
            denominator = tensorflow.reduce_sum(y_true + y_pred, axis=-1)
            return 1 - (numerator + 1) / (denominator + 1)


        def balanced_cross_entropy(beta):
            def convert_to_logits(y_pred):
                # see https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/python/keras/backend.py#L3525
                y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon())

                return tf.log(y_pred / (1 - y_pred))

            def loss(y_true, y_pred):
                y_pred = convert_to_logits(y_pred)
                pos_weight = beta / (1 - beta)
                loss = tf.nn.weighted_cross_entropy_with_logits(logits=y_pred, targets=y_true, pos_weight=pos_weight)

                # or reduce_sum and/or axis=-1
                return tf.reduce_mean(loss * (1 - beta))

            return loss


        def weighted_cross_entropy(beta):
            def convert_to_logits(y_pred):
                # see https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/python/keras/backend.py#L3525
                y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon())

                return tf.log(y_pred / (1 - y_pred))

            def loss(y_true, y_pred):
                y_pred = convert_to_logits(y_pred)
                loss = tf.nn.weighted_cross_entropy_with_logits(logits=y_pred, targets=y_true, pos_weight=beta)

                # or reduce_sum and/or axis=-1
                return tf.reduce_mean(loss)

            return loss


        def intersection_over_union(num_classes):
            return tf.keras.metrics.MeanIoU(num_classes=num_classes)


        def tversky_loss(beta):
            def loss(y_true, y_pred):
                numerator = tensorflow.reduce_sum(y_true * y_pred, axis=-1)
                denominator = y_true * y_pred + beta * (1 - y_true) * y_pred + (1 - beta) * y_true * (1 - y_pred)

                return 1 - (numerator + 1) / (tensorflow.reduce_sum(denominator, axis=-1) + 1)

            return loss


        def lovasz_softmax(y_true, y_pred):
            return lovasz_hinge(labels=y_true, logits=y_pred)


        def execute():
            # TODO read in your parameters from the JSON file
            X = MedML.get('X')
            Y = MedML.get('Y')
            sliceSamples = int(MedML.get('Slice samples'))
            # This will be the name of the data eval folder
            input_folder = path_leaf(values['inputFolder'])
            model_input_size = (int(X), int(Y))

            lossDict = MedML.get('Loss function')

            # check and call the selected loss and insert parameter
            if (lossDict == "dice_loss"):
                lossVal = dice_loss
            if (lossDict == "balanced_cross_entropy"):
                lossVal = balanced_cross_entropy
            if (lossDict == "weighted_cross_entropy"):
                lossVal = weighted_cross_entropy
            if (lossDict == "intersection_over_union"):
                lossVal = intersection_over_union(MedML.get('Number of segmentation classes'))
            if (lossDict == "tversky_loss"):
                lossVal = tversky_loss
            if (lossDict == "lovasz_softmax"):
                lossVal = lovasz_softmax

            print('load model')
            model = tensorflow.keras.models.load_model(path_leaf(values['Model']), custom_objects={lossDict: lossVal})
            # optionally consider loading best weights here as well

            print('searching for data')
            inputFiles = glob(os.path.join(input_folder, '*.nii.gz'), recursive=True) + glob(
                os.path.join(input_folder, '*.nii'), recursive=True)

            print('evaluating model')
            for f in inputFiles:
                print('processing {}'.format(f))

                # load data
                nii = nib.load(f)
                nii_data = nii.get_fdata()
                data_input_size = (nii_data.shape[0], nii_data.shape[1])
                out_data = np.zeros_like(nii_data)

                # loop through slices
                for i in range(0, nii_data.shape[2] - sliceSamples):

                    # determine sampling range
                    if sliceSamples == 1:
                        z = i
                        seg_z = z + sliceSamples // 2
                    else:
                        # warning, sliceSamples is assumed to be 1,3,5,7,9
                        z = range(i, i + sliceSamples)

                        # output segmentation will be the center slice
                        seg_z = z[0] + sliceSamples // 2

                    # resize input data to fit model
                    curr_input = cv2.resize(nii_data[:, :, z], dsize=model_input_size, interpolation=cv2.INTER_CUBIC)
                    if sliceSamples == 1:
                        curr_input = curr_input[np.newaxis, ..., np.newaxis]  # make input size= [1,X,Y,1]
                    else:
                        curr_input = curr_input[np.newaxis, ...]  # make input size= [1,X,Y,sliceSamples]

                    # predict segmentation
                    print("outside")
                    try:
                        print("no need")
                        curr_output = model.predict(curr_input)

                    except:
                        # rotate
                        print("rotate")
                        np.rot90(curr_input)
                        curr_output = model.predict(curr_input)

                    # convert from channel-encoding to integer encoding
                    curr_output_flat = np.zeros((curr_output.shape[1], curr_output.shape[2]))
                    for c in range(0, curr_output.shape[3]):
                        curr_output_flat += np.squeeze(curr_output[0, :, :, c])

                    # reshape segmentation to match input data size
                    curr_output_flat = cv2.resize(curr_output_flat, dsize=data_input_size,
                                                  interpolation=cv2.INTER_NEAREST)

                    # store in output array
                    out_data[:, :, seg_z] = curr_output_flat


                # save file
                # first create new Nifti object based on the input
                new_nii = nib.Nifti1Image(out_data, nii.affine, nii.header)

                # next determine the filename
                drive, filepath = os.path.splitdrive(f)
                path, filename = os.path.split(filepath)
                new_base = filename.replace('.nii', '_seg.nii')
                new_file = os.path.join(drive, path, new_base)

                # now write to disk
                nib.save(new_nii, new_file)

            print('done')


        if __name__ == "__main__":
            execute()