# -*- coding: utf-8 -*-
"""MRI-tumor-segmentation-Brats.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1BTXwPQ3kXM3cODnJeay1LuqCoiSiHZnH
"""
from tqdm import tqdm

# !pip install medpy

# Dataset BraTS 2019 training em https://drive.google.com/drive/folders/1fdHMkpGIuMNJTZioDGskxmueIkFF_Ssn?usp=sharing
# Para funcionar sem modificar o path abaixo, copie esta pasta no seu Google Drive em <raiz>/datasets/
# Ou modifique o path abaixo para o seu diretório, onde /content/drive/My Drive/ é a raiz do seu Google Drive

# from google.colab import drive
# drive.mount('/content/drive')

# Commented out IPython magic to ensure Python compatibility.
# %tensorflow_version 1.x
# import argparse
import tensorflow as tf

# from tensorflow.contrib.keras.python.keras.layers import *
from keras.layers import *
from tensorboardX import SummaryWriter


def Conv3DWithBN(
    x,
    filters,
    ksize,
    strides,
    name,
    padding="same",
    dilation_rate=1,
    center=True,
    scale=True,
    decay=0.99,
):
    x = Conv3D(
        filters=filters,
        kernel_size=ksize,
        strides=strides,
        padding=padding,
        dilation_rate=dilation_rate,
        use_bias=False,
        kernel_initializer="he_normal",
        name=name + "_conv",
    )(x)
    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn"
    )(x)
    x = Activation("relu", name=name + "_relu")(x)
    return x


def Conv2DWithBN(
    x,
    filters,
    ksize,
    strides,
    name,
    padding="same",
    dilation_rate=1,
    center=True,
    scale=True,
    decay=0.99,
):
    x = Conv2D(
        filters=filters,
        kernel_size=ksize,
        strides=strides,
        padding=padding,
        dilation_rate=dilation_rate,
        use_bias=False,
        kernel_initializer="he_normal",
        name=name + "_conv",
    )(x)
    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn"
    )(x)
    x = Activation("relu", name=name + "_relu")(x)
    return x


def Conv1DWithBN(
    x,
    filters,
    ksize,
    strides,
    name,
    padding="same",
    dilation_rate=1,
    center=True,
    scale=True,
    decay=0.99,
):
    x = Conv1D(
        filters=filters,
        kernel_size=ksize,
        strides=strides,
        padding=padding,
        dilation_rate=dilation_rate,
        use_bias=False,
        kernel_initializer="he_normal",
        name=name + "_conv",
    )(x)
    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn"
    )(x)
    x = Activation("relu", name=name + "_relu")(x)
    return x


def DenseWithBN(
    x, units, name, kernel_regularizer=None, center=True, scale=True, decay=0.99
):
    x = Dense(
        units=units,
        use_bias=False,
        kernel_regularizer=kernel_regularizer,
        name=name + "_weight",
    )(x)
    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bias"
    )(x)
    x = Activation("relu", name=name + "_relu")(x)
    return x


def ResNetUnit2D(
    x, filters, ksize, name, end=False, center=True, scale=True, decay=0.99
):
    identity = x

    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn_1"
    )(x)
    x = Activation("relu", name=name + "_relu_1")(x)
    x = Conv2D(
        filters,
        kernel_size=ksize,
        strides=1,
        padding="same",
        kernel_initializer="he_normal",
        name=name + "_conv_1",
    )(x)

    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn_2"
    )(x)
    x = Activation("relu", name=name + "_relu_2")(x)
    x = Conv2D(
        filters,
        kernel_size=ksize,
        strides=1,
        padding="same",
        kernel_initializer="he_normal",
        name=name + "_conv2",
    )(x)

    x = add([x, identity])
    if end:
        x = BatchNormalization(center=center, scale=scale, momentum=decay)(x)
        x = Activation("relu")(x)
    return x


def ResNetUnitIncreasingDims2D(
    x, filters, ksize, strides, name, begin=False, center=True, scale=True, decay=0.99
):
    identity = x

    if not begin:
        x = BatchNormalization(
            center=center, scale=scale, momentum=decay, name=name + "_bn_1"
        )(x)
        x = Activation("relu", name=name + "_relu_1")(x)
    x = Conv2D(
        filters,
        kernel_size=ksize,
        strides=strides[0],
        padding="same",
        kernel_initializer="he_normal",
        use_bias=False,
        name=name + "_conv_1",
    )(x)

    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn_2"
    )(x)
    x = Activation("relu", name=name + "_relu_2")(x)
    x = Conv2D(
        filters,
        kernel_size=ksize,
        strides=strides[1],
        padding="same",
        kernel_initializer="he_normal",
        use_bias=False,
        name=name + "_conv_2",
    )(x)

    identity = Conv2D(
        filters,
        kernel_size=1,
        strides=strides[0],
        padding="same",
        kernel_initializer="he_normal",
        name=name + "_conv_identity",
    )(identity)
    x = add([x, identity])
    return x


def ResNetUnit1D(
    x, filters, ksize, name, end=False, center=True, scale=True, decay=0.99
):
    identity = x

    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn_1"
    )(x)
    x = Activation("relu", name=name + "_relu_1")(x)
    x = Conv1D(
        filters,
        kernel_size=ksize,
        strides=1,
        padding="same",
        kernel_initializer="he_normal",
        use_bias=False,
        name=name + "_conv_1",
    )(x)

    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn_2"
    )(x)
    x = Activation("relu", name=name + "_relu_2")(x)
    x = Conv1D(
        filters,
        kernel_size=ksize,
        strides=1,
        padding="same",
        kernel_initializer="he_normal",
        use_bias=False,
        name=name + "_conv_2",
    )(x)

    x = add([x, identity])
    if end:
        x = BatchNormalization(
            center=center, scale=scale, momentum=decay, name=name + "_bn_3"
        )(x)
        x = Activation("relu", name=name + "_relu_3")(x)
    return x


def ResNetUnitIncreasingDims1D(
    x, filters, ksize, strides, name, begin=False, center=True, scale=True, decay=0.99
):
    """
    ResNet unit without BottleNeck. 2 layers
    :param x:
    :param filters:
    :param ksize:
    :param strides: list with 2 elements, stride for each layer
    :param begin:
    :return:
    """

    identity = x

    if not begin:
        x = BatchNormalization(
            center=center, scale=scale, momentum=decay, name=name + "_bn_1"
        )(x)
        x = Activation("relu", name=name + "_relu_1")(x)
    x = Conv1D(
        filters,
        kernel_size=ksize,
        strides=strides[0],
        padding="same",
        kernel_initializer="he_normal",
        use_bias=False,
        name=name + "_conv_1",
    )(x)

    x = BatchNormalization(
        center=center, scale=scale, momentum=decay, name=name + "_bn_2"
    )(x)
    x = Activation("relu", name=name + "_relu_2")(x)
    x = Conv1D(
        filters,
        kernel_size=ksize,
        strides=strides[1],
        padding="same",
        kernel_initializer="he_normal",
        use_bias=False,
        name=name + "_conv_2",
    )(x)

    identity = Conv1D(
        filters,
        kernel_size=1,
        strides=strides[0],
        padding="same",
        kernel_initializer="he_normal",
        name=name + "_conv_identity",
    )(identity)
    x = add([x, identity])
    return x


def ContextualAvgPooling(x, ksizes, strides):
    """
    Concatenate input with pooling result
    :param x:
    :param ksizes:
    :param strides:
    :return:
    """
    out = None
    for ks in ksizes:
        x_pooled = AvgPool1D(pool_size=ks, strides=strides, padding="same")
        if out is None:
            out = x_pooled
        else:
            out = concatenate([out, x_pooled], axis=-1)
    x = concatenate([x, out], axis=-1)
    return x


def ContextualAtrousConv1D(x, filters, ksize, strides, dilation_rates, name):
    """
    Retrieve contextual information with Atrous Convolution
    :param x:
    :param filters:
    :param ksize:
    :param strides:
    :param dilation_rates:
    :return:
    """
    concat = x
    for dr in dilation_rates:
        x_atrous = Conv1DWithBN(
            x,
            filters=filters,
            ksize=ksize,
            strides=strides,
            dilation_rate=dr,
            name=name + "a_conv_s" + str(dr),
        )
        concat = concatenate([concat, x_atrous], axis=-1)
    concat = Conv1DWithBN(
        concat,
        filters=256,
        ksize=ksize,
        strides=strides,
        name=name + "atrou_post_conv1",
    )
    return concat


def ContextualAtrousConv3D(x, filters, ksize, strides, dilation_rates, name):
    """
    Retrieve contextual information with Atrous Convolution
    :param x:
    :param filters:
    :param ksize:
    :param strides:
    :param dilation_rates:
    :return:
    """
    concat = None
    for dr in dilation_rates:
        x_atrous = Conv3DWithBN(
            x,
            filters=filters,
            ksize=ksize,
            strides=strides,
            dilation_rate=dr,
            name=name + "a_conv_s" + str(dr),
        )
        if concat is None:
            concat = x_atrous
        else:
            concat = concatenate([concat, x_atrous], axis=-1)
    concat = Conv3DWithBN(
        concat,
        filters=filters,
        ksize=ksize,
        strides=strides,
        name=name + "atrou_post_conv1",
    )
    return concat


def SharedAtrousConv1D(x, SharedConvs, PostConv):
    concat = x
    for SharedConv in SharedConvs:
        x_atrous = SharedConv(x)
        x_atrous = BatchNormalization()(x_atrous)
        x_atrous = Activation("relu")(x_atrous)
        concat = concatenate([concat, x_atrous], axis=-1)
    concat = PostConv(concat)
    concat = BatchNormalization()(concat)
    concat = Activation("relu")(concat)
    return concat


def densenet_block3d(x, k, rep):
    dense_input = x
    for i in range(rep):
        x_dense = Conv3D(
            filters=k, kernel_size=3, strides=1, padding="same", activation="relu"
        )(dense_input)
        dense_input = concatenate([dense_input, x_dense])
    return dense_input


def DenseNetTransit(x, rate=1, name=None):
    if rate != 1:
        out_features = x.get_shape().as_list()[-1] * rate
        x = BatchNormalization(center=True, scale=True, name=name + "_bn")(x)
        x = Activation("relu", name=name + "_relu")(x)
        x = Conv3D(
            filters=out_features,
            kernel_size=1,
            strides=1,
            padding="same",
            kernel_initializer="he_normal",
            use_bias=False,
            name=name + "_conv",
        )(x)
    x = AveragePooling3D(pool_size=2, strides=2, padding="same")(x)
    return x


def DenseNetUnit3D(x, growth_rate, ksize, rep, bn_decay=0.99, name=None):
    for i in range(rep):
        concat = x
        x = BatchNormalization(
            center=True, scale=True, momentum=bn_decay, name=name + "_bn_rep_" + str(i)
        )(x)
        x = Activation("relu")(x)
        x = Conv3D(
            filters=growth_rate,
            kernel_size=ksize,
            padding="same",
            kernel_initializer="glorot_normal",
            use_bias=False,
            name=name + "_conv_rep_" + str(i),
        )(x)
        x = concatenate([concat, x])
    return x



from tf_layers import (
    BilinearUpsampling1D,
    BilinearUpsampling3D,

)


from tf_models import (
    BraTS2ScaleDenseNetConcat,
    BraTS2ScaleDenseNetConcat_large,
    BraTS2ScaleDenseNet,
    BraTS3ScaleDenseNet,
    BraTS1ScaleDenseNet,
    PlainCounterpart,
)





from sklearn.model_selection import train_test_split


from pathlib import Path
from typing import List


def read_file(file_path: Path) -> str:
    with open(file_path) as text_file:
        return text_file.read()


def read_multiline(file_path: Path) -> List[str]:
    """Read, split by LF, and pop final empty line if present."""
    text = read_file(file_path)
    lines = text.split("\n")
    if lines[-1] == "":
        lines.pop()
    return lines


files = read_multiline("train.txt")
val_files = read_multiline("val.txt")

files = [file_name[:-1] for file_name in files]

val_files = [file_name[:-1] for file_name in val_files]

# print(files[2])
# print(len(files))

# create training and testing vars
files, val_files = train_test_split(files, test_size=0.8)
print("%d training   samples" % len(files))
print("%d validation samples" % len(val_files))

import glob

import numpy as np
from keras.backend import learning_phase
from keras.layers import concatenate, Conv3D

# import tf_models
from sklearn.preprocessing import scale
import tensorflow.compat.v1 as tf

tf.disable_v2_behavior()
# from tensorflow.contrib.keras.python.keras.backend import learning_phase
# from tensorflow.contrib.keras.python.keras.layers import concatenate, Conv3D
from nibabel import load as load_nii
import os
import argparse
import keras


def parse_inputs():

    parser = argparse.ArgumentParser(description="train the model")
    parser.add_argument(
        "-r",
        "--root-path",
        dest="root_path",
        default="./MICCAI_BraTS_2019_Data_Training/HGG/"
    )
    parser.add_argument(
        "-sp", "--save-path", dest="save_path", default="dense24_correction"
    )
    parser.add_argument(
        "-lp", "--load-path", dest="load_path", default="dense24_correction"
    )
    parser.add_argument("-ow", "--offset-width", dest="offset_w", type=int, default=12)
    parser.add_argument("-oh", "--offset-height", dest="offset_h", type=int, default=12)
    parser.add_argument(
        "-oc", "--offset-channel", dest="offset_c", nargs="+", type=int, default=12
    )
    parser.add_argument("-ws", "--width-size", dest="wsize", type=int, default=38)
    parser.add_argument("-hs", "--height-size", dest="hsize", type=int, default=38)
    parser.add_argument("-cs", "--channel-size", dest="csize", type=int, default=38)
    parser.add_argument("-ps", "--pred-size", dest="psize", type=int, default=12)
    parser.add_argument("-bs", "--batch-size", dest="batch_size", type=int, default=4)
    parser.add_argument("-e", "--num-epochs", dest="num_epochs", type=int, default=5)
    parser.add_argument(
        "-c", "--continue-training", dest="continue_training", type=bool, default=False
    )
    parser.add_argument(
        "-mn", "--model_name", dest="model_name", type=str, default="dense24"
    )  # dense48 no_dense
    parser.add_argument(
        "-nc", "--n4correction", dest="correction", type=bool, default=False
    )
    parser.add_argument("-gpu", "--gpu_id", dest="gpu_id", type=str, default="0")
    return vars(parser.parse_args())


options = parse_inputs()

os.environ["CUDA_VISIBLE_DEVICES"] = options['gpu_id']


def acc_tf(y_pred, y_true):
    correct_prediction = tf.equal(
        tf.cast(tf.argmax(y_pred, -1), tf.int32),
        tf.cast(tf.argmax(y_true, -1), tf.int32),
    )
    return 100 * tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


def get_patches_3d(data, labels, centers, hsize, wsize, csize, psize, preprocess=True):
    """
    :param data: 4D nparray (h, w, c, ?)
    :param centers:
    :param hsize:
    :param wsize:
    :param csize:
    :return:
    """
    patches_x, patches_y = [], []
    offset_p = (hsize - psize) / 2
    for i in range(len(centers[0])):
        h, w, c = centers[0, i], centers[1, i], centers[2, i]
        h_beg = int(min(max(0, h - hsize / 2), 240 - hsize))
        w_beg = int(min(max(0, w - wsize / 2), 240 - wsize))
        c_beg = int(min(max(0, c - csize / 2), 155 - csize))
        ph_beg = int(h_beg + offset_p)
        pw_beg = int(w_beg + offset_p)
        pc_beg = int(c_beg + offset_p)
        vox = data[
            h_beg : h_beg + hsize, w_beg : w_beg + wsize, c_beg : c_beg + csize, :
        ]
        vox_labels = labels[
            ph_beg : ph_beg + psize, pw_beg : pw_beg + psize, pc_beg : pc_beg + psize
        ]
        patches_x.append(vox)
        patches_y.append(vox_labels)
    return np.array(patches_x), np.array(patches_y)


def positive_ratio(x):
    return float(np.sum(np.greater(x, 0))) / np.prod(x.shape)


def norm(image):
    image = np.squeeze(image)
    image_nonzero = image[np.nonzero(image)]
    return (image - image_nonzero.mean()) / image_nonzero.std()


def segmentation_loss(y_true, y_pred, n_classes):
    y_true = tf.reshape(y_true, (-1, n_classes))
    y_pred = tf.reshape(y_pred, (-1, n_classes))
    return tf.reduce_mean(
        tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=y_pred)
    )


def vox_preprocess(vox):
    vox_shape = vox.shape
    vox = np.reshape(vox, (-1, vox_shape[-1]))
    vox = scale(vox, axis=0)
    return np.reshape(vox, vox_shape)


def one_hot(y, num_classes):
    y_ = np.zeros([len(y), num_classes])
    y_[np.arange(len(y)), y] = 1
    return y_


def dice_coef_np(y_true, y_pred, num_classes):
    """
    :param y_true: sparse labels
    :param y_pred: sparse labels
    :param num_classes: number of classes
    :return:
    """
    y_true = y_true.astype(int)
    y_pred = y_pred.astype(int)
    y_true = y_true.flatten()
    y_true = one_hot(y_true, num_classes)
    y_pred = y_pred.flatten()
    y_pred = one_hot(y_pred, num_classes)
    intersection = np.sum(y_true * y_pred, axis=0)
    return (2.0 * intersection) / (np.sum(y_true, axis=0) + np.sum(y_pred, axis=0))


def vox_generator(all_files, n_pos, n_neg, correction=False):
    path = options["root_path"]

    def get_filename(prefix):
        return glob.glob(os.path.join(path, file, prefix))[0].split("/")[-1]

    while 1:
        for file in all_files:
            if correction:
                flair = load_nii(
                    os.path.join(path, file, file + "_flair_corrected.nii.gz")
                ).get_data()
                t2 = load_nii(
                    os.path.join(path, file, file + "_t2_corrected.nii.gz")
                ).get_data()
                t1 = load_nii(
                    os.path.join(path, file, file + "_t1_corrected.nii.gz")
                ).get_data()
                t1ce = load_nii(
                    os.path.join(path, file, file + "_t1ce_corrected.nii.gz")
                ).get_data()
            else:
                flair = load_nii(
                    os.path.join(path, file, file + "_flair.nii.gz")
                ).get_data()
                t2 = load_nii(os.path.join(path, file, file + "_t2.nii.gz")).get_data()
                t1 = load_nii(os.path.join(path, file, file + "_t1.nii.gz")).get_data()
                t1ce = load_nii(
                    os.path.join(path, file, file + "_t1ce.nii.gz")
                ).get_data()

            data_norm = np.array([norm(flair), norm(t2), norm(t1), norm(t1ce)])
            data_norm = np.transpose(data_norm, axes=[1, 2, 3, 0])
            labels = load_nii(os.path.join(path, file, file + "_seg.nii.gz")).get_data()

            foreground = np.array(np.where(labels > 0))
            background = np.array(np.where((labels == 0) & (flair > 0)))

            # n_pos = int(foreground.shape[1] * discount)
            foreground = foreground[
                :, np.random.permutation(foreground.shape[1])[:n_pos]
            ]
            background = background[
                :, np.random.permutation(background.shape[1])[:n_neg]
            ]

            centers = np.concatenate((foreground, background), axis=1)
            centers = centers[:, np.random.permutation(n_neg + n_pos)]

            yield data_norm, labels, centers


def label_transform(y, nlabels):
    return [
        keras.utils.to_categorical(
            np.copy(y).astype(dtype=np.bool), num_classes=2
        ).reshape([y.shape[0], y.shape[1], y.shape[2], y.shape[3], 2]),
        keras.utils.to_categorical(y, num_classes=nlabels).reshape(
            [y.shape[0], y.shape[1], y.shape[2], y.shape[3], nlabels]
        ),
    ]


def train():
    tb_writer = SummaryWriter(comment='initial')

    NUM_EPOCHS = options["num_epochs"]
    LOAD_PATH = options["load_path"]
    SAVE_PATH = options["save_path"]
    PSIZE = options["psize"]
    HSIZE = options["hsize"]
    WSIZE = options["wsize"]
    CSIZE = options["csize"]
    model_name = options["model_name"]
    BATCH_SIZE = options["batch_size"]
    continue_training = options["continue_training"]

    num_labels = 5
    global files
    """
    files = []
    with open('train.txt') as f:
        for line in f:
            files.append(line[:-1])
    """

    flair_t2_node = tf.placeholder(
        dtype=tf.float32, shape=(None, HSIZE, WSIZE, CSIZE, 2)
    )
    t1_t1ce_node = tf.placeholder(
        dtype=tf.float32, shape=(None, HSIZE, WSIZE, CSIZE, 2)
    )
    flair_t2_gt_node = tf.placeholder(
        dtype=tf.int32, shape=(None, PSIZE, PSIZE, PSIZE, 2)
    )
    t1_t1ce_gt_node = tf.placeholder(
        dtype=tf.int32, shape=(None, PSIZE, PSIZE, PSIZE, 5)
    )

    if model_name == "dense48":
        flair_t2_15, flair_t2_27 = BraTS2ScaleDenseNetConcat_large(
            input=flair_t2_node, name="flair"
        )
        t1_t1ce_15, t1_t1ce_27 = BraTS2ScaleDenseNetConcat_large(
            input=t1_t1ce_node, name="t1"
        )
    elif model_name == "no_dense":

        flair_t2_15, flair_t2_27 = PlainCounterpart(input=flair_t2_node, name="flair")
        t1_t1ce_15, t1_t1ce_27 = PlainCounterpart(input=t1_t1ce_node, name="t1")

    elif model_name == "dense24":

        flair_t2_15, flair_t2_27 = BraTS2ScaleDenseNetConcat(
            input=flair_t2_node, name="flair"
        )
        t1_t1ce_15, t1_t1ce_27 = BraTS2ScaleDenseNetConcat(
            input=t1_t1ce_node, name="t1"
        )
    else:
        print(" No such model name ")

    t1_t1ce_15 = concatenate([t1_t1ce_15, flair_t2_15])
    t1_t1ce_27 = concatenate([t1_t1ce_27, flair_t2_27])

    flair_t2_15 = Conv3D(
        2, kernel_size=1, strides=1, padding="same", name="flair_t2_15_cls"
    )(flair_t2_15)
    flair_t2_27 = Conv3D(
        2, kernel_size=1, strides=1, padding="same", name="flair_t2_27_cls"
    )(flair_t2_27)
    t1_t1ce_15 = Conv3D(
        num_labels, kernel_size=1, strides=1, padding="same", name="t1_t1ce_15_cls"
    )(t1_t1ce_15)
    t1_t1ce_27 = Conv3D(
        num_labels, kernel_size=1, strides=1, padding="same", name="t1_t1ce_27_cls"
    )(t1_t1ce_27)

    flair_t2_score = (
        flair_t2_15[:, 13:25, 13:25, 13:25, :] + flair_t2_27[:, 13:25, 13:25, 13:25, :]
    )

    t1_t1ce_score = (
        t1_t1ce_15[:, 13:25, 13:25, 13:25, :] + t1_t1ce_27[:, 13:25, 13:25, 13:25, :]
    )

    loss = segmentation_loss(flair_t2_gt_node, flair_t2_score, 2) + segmentation_loss(
        t1_t1ce_gt_node, t1_t1ce_score, 5
    )

    acc_flair_t2 = acc_tf(y_pred=flair_t2_score, y_true=flair_t2_gt_node)
    acc_t1_t1ce = acc_tf(y_pred=t1_t1ce_score, y_true=t1_t1ce_gt_node)

    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        optimizer = tf.train.AdamOptimizer(learning_rate=5e-4).minimize(loss)

    saver = tf.train.Saver(max_to_keep=15)
    data_gen_train = vox_generator(
        all_files=files, n_pos=200, n_neg=200, correction=options["correction"]
    )
    config = tf.ConfigProto(allow_soft_placement=True)
    config.gpu_options.allow_growth = True
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    
    with tf.Session(config=config) as sess:
        if continue_training:
            saver.restore(sess, LOAD_PATH)
        else:
            sess.run(tf.global_variables_initializer())
        curr_iteration = 0
        for ei in tqdm(range(NUM_EPOCHS)):
            for pi in tqdm(range(len(files))):
                acc_pi, loss_pi = [], []
                data, labels, centers = data_gen_train.__next__()
                n_batches = int(np.ceil(float(centers.shape[1]) / BATCH_SIZE))
                for nb in tqdm(range(n_batches)):
                    offset_batch = min(nb * BATCH_SIZE, centers.shape[1] - BATCH_SIZE)
                    data_batch, label_batch = get_patches_3d(
                        data,
                        labels,
                        centers[:, offset_batch : offset_batch + BATCH_SIZE],
                        HSIZE,
                        WSIZE,
                        CSIZE,
                        PSIZE,
                        False,
                    )
                    label_batch = label_transform(label_batch, 5)
                    _, l, acc_ft, acc_t1c = sess.run(
                        fetches=[optimizer, loss, acc_flair_t2, acc_t1_t1ce],
                        feed_dict={
                            flair_t2_node: data_batch[:, :, :, :, :2],
                            t1_t1ce_node: data_batch[:, :, :, :, 2:],
                            flair_t2_gt_node: label_batch[0],
                            t1_t1ce_gt_node: label_batch[1],
                            learning_phase(): 1,
                        },
                    )
                    acc_pi.append([acc_ft, acc_t1c])
                    loss_pi.append(l)
                    n_pos_sum = np.sum(np.reshape(label_batch[0], (-1, 2)), axis=0)
                    tb_writer.add_scalar("train-loss", l,  curr_iteration)
                    tb_writer.add_scalar("acc/flair_t2", acc_ft, curr_iteration)
                    tb_writer.add_scalar("acc/t1_t1ce", acc_t1c, curr_iteration)
                    curr_iteration += 1
                    tqdm.write(
                        f"epoch-patient: {ei + 1}, {pi + 1}, "
                        f"iter: {nb + 1}-{n_batches}, "
                        f"p%%: {n_pos_sum[1]/float(np.sum(n_pos_sum)):.4f}, "
                        f"loss: {l:.4f}, acc_flair_t2: {acc_ft:.2f}, "
                        f"acc_t1_t1ce: {acc_t1c:.2f}"
                    )

                current_patient_i = (ei + 1) * (pi + 1)
                tb_writer.add_scalar(
                    f"patient{pi}/loss",
                    np.mean(loss_pi),
                    current_patient_i,
                )
                tb_writer.add_scalar(
                    f"patient{pi}/acc",
                    np.mean(acc_pi),
                    current_patient_i,
                )
                tqdm.write(
                    "patient loss: %.4f, patient acc: %.4f"
                    % (np.mean(loss_pi), np.mean(acc_pi))
                )

            saver.save(sess, SAVE_PATH, global_step=ei)
            tqdm.write("Model saved.")


if __name__ == "__main__":

    train()
