train.py

# -*- coding: utf-8 -*-
"""
Main file to train the model.
=============================================================
Created on Tue Apr  4 09:35:14 2020
@author: Jingnan

                    .::::.
                  .::::::::.
                 :::::::::::
             ..:::::::::::'
           '::::::::::::'
             .::::::::::
        '::::::::::::::..
             ..::::::::::::.
           ``::::::::::::::::
            ::::``:::::::::'        .:::.
           ::::'   ':::::'       .::::::::.
         .::::'     ::::      .:::::::'::::.
        .:::'       :::::  .:::::::::' ':::::.
       .::'        :::::.:::::::::'      ':::::.
      .::'         ::::::::::::::'         ``::::.
  ...:::           ::::::::::::'              ``::.
 ```` ':.          ':::::::::'                  ::::..

"""
import copy

import numpy as np
import os
import gc
import sys

from futils.find_connect_parts import write_cntd_pts
from mypath import Mypath
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras import callbacks
from tensorflow.keras.utils import plot_model
import matplotlib.pyplot as plt
from shutil import copy2
import inspect
import pandas as pd

from futils import compiled_models as cpmodels, segmentor as v_seg
from jjnutils.util import save_model_best
from set_args import args
from write_dice import write_dices_to_csv
from futils.write_batch_preds import write_preds_to_disk
from futils.compute_distance_metrics_and_save import write_all_metrics
from futils.generate_fissure_from_masks import gntFissure
from futils.dataloader import ScanIterator

# os.environ['CUDA_VISIBLE_DEVICES'] = "0" # use the first GPU
# tf.keras.mixed_precision.experimental.set_policy('infer')  # mix precision training
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU
sess = tf.Session(config=config)
K.set_session(sess)  # set this TensorFlow session as the default session for Keras

K.set_learning_phase(1)  # try with 1
print(sys.argv[1:])  # print all arguments passed to script


class GetList:
    def __init__(self, model_names):
        self.model_names = model_names  # a list
        self.all_names = ["net_itgt_ls_rc", "net_itgt_lb_rc", "net_itgt_vs_rc", "net_itgt_aw_rc", "net_itgt_lu_rc",
                          "net_lesion", "net_lobe", "net_vessel", "net_airway", "net_lung", "net_recon"]

    def get_io_list(self, ags):
        """Get io list according to io parameters. """
        io_list = [ags.ls_io, ags.lb_io, ags.vs_io, ags.aw_io, ags.lu_io,
                   ags.ls_io, ags.lb_io, ags.vs_io, ags.aw_io, ags.lu_io, ags.rc_io]
        io_dict = dict(zip(self.all_names, io_list))
        return list(map(io_dict.get, self.model_names))

    def get_task_list(self):
        """
        Get task list according to a list of model names. Note that one task may corresponds to multiple models.

        :return: a list of tasks
        """
        task_list = ["lesion", "lobe", "vessel", "airway", "lung",
                     "lesion", "lobe", "vessel", "airway", "lung", "recon"]
        task_dict = dict(zip(self.all_names, task_list))
        return list(map(task_dict.get, self.model_names))

    def get_labels_list(self):
        """
        Get the labels list according to given task list.

        :return: a list of labels' list.
        """
        lb_0, lb_2, lb_6 = [], [0, 1], [0, 4, 5, 6, 7, 8]
        label_list = [lb_2, lb_6, lb_2, lb_2, lb_2,
                      lb_2, lb_6, lb_2, lb_2, lb_2, lb_0]
        labels_dict = dict(zip(self.all_names, label_list))
        return list(map(labels_dict.get, self.model_names))

    def get_path_list(self):
        task_list = self.get_task_list()
        return [Mypath(x) for x in task_list]  # a list of Mypath objectives, each Mypath corresponds to a task

    def get_tr_nb_list(self, ags):
        tr_nb_list = [ags.ls_tr_nb, ags.lb_tr_nb, ags.vs_tr_nb, ags.aw_tr_nb, ags.lu_tr_nb,
                      ags.ls_tr_nb, ags.lb_tr_nb, ags.vs_tr_nb, ags.aw_tr_nb, ags.lu_tr_nb, ags.rc_tr_nb]
        tr_nb_dict = dict(zip(self.all_names, tr_nb_list))
        return list(map(tr_nb_dict.get, self.model_names))

    def get_ao_list(self, ags):
        ao_list = [ags.ao_ls, ags.ao_lb, ags.ao_vs, ags.ao_aw, ags.ao_lu,
                   ags.ao_ls, ags.ao_lb, ags.ao_vs, ags.ao_aw, ags.ao_lu, ags.ao_rc]
        ao_dict = dict(zip(self.all_names, ao_list))
        return list(map(ao_dict.get, self.model_names))

    def get_ds_list(self, ags):
        ds_list = [ags.ds_ls, ags.ds_lb, ags.ds_vs, ags.ds_aw, ags.ds_lu,
                   ags.ds_ls, ags.ds_lb, ags.ds_vs, ags.ds_aw, ags.ds_lu, ags.ds_rc]
        ds_dict = dict(zip(self.all_names, ds_list))
        return list(map(ds_dict.get, self.model_names))

    def get_tsp_list(self, ags):
        def pr_float(tsp):
            return [float(i.lstrip()) for i in tsp.split('_')]

        tsp_list = [pr_float(ags.tsp_ls), pr_float(ags.tsp_lb), pr_float(ags.tsp_vs), pr_float(ags.tsp_aw), pr_float(ags.tsp_lu),
                    pr_float(ags.tsp_ls), pr_float(ags.tsp_lb), pr_float(ags.tsp_vs), pr_float(ags.tsp_aw), pr_float(ags.tsp_lu), pr_float(ags.tsp_rc)]
        tsp_dict = dict(zip(self.all_names, tsp_list))
        return list(map(tsp_dict.get, self.model_names))

    def get_tsz_list(self, ags):
        def pr_int(tsz):
            return [int(i.lstrip()) for i in tsz.split('_')]

        tsz_list = [pr_int(ags.tsz_ls), pr_int(ags.tsz_lb), pr_int(ags.tsz_vs), pr_int(ags.tsz_aw), pr_int(ags.tsz_lu),
                    pr_int(ags.tsz_ls), pr_int(ags.tsz_lb), pr_int(ags.tsz_vs), pr_int(ags.tsz_aw), pr_int(ags.tsz_lu), pr_int(ags.tsz_rc)]
        tsz_dict = dict(zip(self.all_names, tsz_list))
        return list(map(tsz_dict.get, self.model_names))

    def get_load_name_list(self, ags):
        ld_list = [ags.ld_itgt_ls_rc, ags.ld_itgt_lb_rc, ags.ld_itgt_vs_rc, ags.ld_itgt_aw_rc, ags.ld_itgt_lu_rc,
                   ags.ld_ls, ags.ld_lb, ags.ld_vs, ags.ld_aw, ags.ld_lu, ags.ld_rc]
        load_dict = dict(zip(self.all_names, ld_list))
        return list(map(load_dict.get, self.model_names))


def write_metrics(sub_dir, fissureradius, workers=10, mypath=None, cntd_pts=None, labels=None):
    if args.main_model in ["net_lobe", "net_itgt_lb_rc"]:
        if sub_dir is "GLUCOLD":  # write metrics for lobe and fissure (GLUCOLD), for lung and fissure (LOLA11)
            goals = ['lobe', 'fissure']
        elif sub_dir is "LOLA11":
            goals = ['lung', 'fissure']
        else:
            raise Exception("sub_dir is not correct")

        for goal in goals:
            if goal is "lobe":
                labels = [4, 5, 6, 7, 8]
                fissure = False
                lung = False
            elif goal is "fissure":
                labels = [1]
                fissure = True
                lung = False
            else:
                labels = [1]
                fissure = False
                lung = True
            workers = 3 if sub_dir is "GLUCOLD" else workers

            write_all_metrics(labels=labels,  # exclude background
                              gdth_path=mypath.gdth_path("valid", sub_dir=sub_dir),
                              pred_path=mypath.pred_path("valid", sub_dir=sub_dir, cntd_pts=cntd_pts),
                              csv_file=mypath.all_metrics_fpath("valid", fissure=fissure, sub_dir=sub_dir),
                              fissure=fissure, fissureradius=fissureradius, lung=lung, workers=workers)
    else:
        write_all_metrics(labels=labels,  # exclude background
                          gdth_path=mypath.gdth_path("valid", sub_dir=sub_dir),
                          pred_path=mypath.pred_path("valid", sub_dir=sub_dir, cntd_pts=cntd_pts),
                          csv_file=mypath.all_metrics_fpath("valid", sub_dir=sub_dir), workers=workers)


class TaskArgs:
    def __init__(self):
        self.net = None
        self.mypath = None
        self.task = None
        self.labels = None
        self.model_name = None
        self.sub_dir = None
        self.ld_name = None
        self.tr_nb = None
        self.ao = None
        self.ds = None
        self.io = None
        self.tsp = None
        self.tsz = None
        self.tszzyx = [],
        self.tspzyx = [],

        self.train_it = None
        self.train_data_gen = None
        self.valid_array = None

        self.best_tr_loss = 10000
        self.best_vd_loss = 10000
        self.current_tr_loss = 10000
        self.lr = 0.0001

    def plot_model(self):
        model_figure_fpath = self.mypath.model_figure_path() + '/' + self.model_name + '.png'
        plot_model(self.net, show_shapes=True, to_file=model_figure_fpath)
        print('successfully plot model structure at: ', model_figure_fpath)

    def load_weights_if_need(self):
        if self.ld_name is not 'None':  # 'None' is from arg parse as string
            try:
                saved_model = self.mypath.model_fpath_best_whole(phase='valid', str_name=self.ld_name)
                self.net.load_weights(saved_model)
            except OSError:
                try:  # for vessel who does not have model_fpath_best_whole
                    saved_model = self.mypath.model_fpath_best_patch(phase='valid', str_name=self.ld_name)
                    self.net.load_weights(saved_model)
                except OSError:  # for recon who does not have model_fpath_best_whole and model_fpath_best_patch
                    saved_model = self.mypath.model_fpath_best_patch(phase='train', str_name=self.ld_name)
                    self.net.load_weights(saved_model)
            print('loaded main model weights successfully from: ', saved_model)

    def save_json(self):
        # save model architecture and config
        model_json = self.net.to_json()
        with open(self.mypath.json_fpath(), "w") as json_file:
            json_file.write(model_json)
            print('successfully write new json file of task ', self.task, self.mypath.json_fpath())

    def do_vilidation_if_need(self, idx_, valid_period, cntd_pts=None):
        if (idx_ % valid_period == 0) and (self.model_name == args.main_model):  # only valid for main model
            # In my multi-task model (co-training, or alternative training), I can not use validation_data and
            # validation_freq in net.fit() function. Because there are only one step (patch) at each fit().
            # So in order to assess the valid metrics, I use an independent function to predict the validation
            # and training dataset. And I can also set the period_valid as the validation_freq.
            # save predicted results and compute the dices
            for phase in ['valid']:
                if self.task=="lesion":
                    valid_nb = 100
                else:
                    valid_nb = 1
                segment = v_seg.v_segmentor(batch_size=args.batch_size,
                                            model=self.mypath.model_fpath_best_patch(phase),
                                            ptch_sz=args.ptch_sz, ptch_z_sz=args.ptch_z_sz,
                                            trgt_size_list=self.tszzyx,
                                            trgt_space_list=self.tspzyx,
                                            task=self.task, attention=args.attention)

                write_preds_to_disk(segment=segment,
                                    data_dir=self.mypath.ori_ct_path(phase),
                                    preds_dir=self.mypath.pred_path(phase),
                                    number=valid_nb, stride=0.8, workers=5, qsize=5,
                                    pad=args.pad)  # set stride 0.8 to save time

                write_dices_to_csv(step_nb=idx_,
                                   labels=self.labels[1:],
                                   gdth_path=self.mypath.gdth_path(phase),
                                   pred_path=self.mypath.pred_path(phase),
                                   csv_file=self.mypath.dices_fpath(phase))

                save_model_best(dice_file=self.mypath.dices_fpath(phase),
                                segment=segment,
                                model_fpath=self.mypath.model_fpath_best_whole(phase))
                print('step number', idx_, 'lr for', self.task, 'is', K.eval(self.net.optimizer.lr), file=sys.stderr)

        if (idx_ == args.step_nb - 1) and (self.model_name == args.main_model):  # last step, fully validation
            if self.sub_dir is "GLUCOLD":
                test_nb = 5
                stride = 0.25
                fissureradius = 3
            elif self.sub_dir is "LOLA11":
                test_nb = 100
                stride = 0.25
                fissureradius = 1
            else:
                test_nb = None
                stride = 0.25
                fissureradius = 0

            segment = v_seg.v_segmentor(batch_size=args.batch_size,
                                        model=self.mypath.model_fpath_best_whole("valid"),
                                        ptch_sz=args.ptch_sz, ptch_z_sz=args.ptch_z_sz,
                                        trgt_size_list=self.tszzyx,
                                        trgt_space_list=self.tspzyx,
                                        task=self.task, attention=args.attention)

            write_preds_to_disk(segment=segment,
                                data_dir=self.mypath.ori_ct_path("valid", sub_dir=self.sub_dir),
                                preds_dir=self.mypath.pred_path("valid", sub_dir=self.sub_dir),
                                number=test_nb, stride=stride, workers=5, qsize=5,
                                pad=args.pad)  # set stride 0.8 to save time
            if cntd_pts:
                write_cntd_pts(self.mypath.pred_path("valid", sub_dir=self.sub_dir), workers=5,
                               target_dir=self.mypath.pred_path("valid", sub_dir=self.sub_dir,
                                                                cntd_pts=cntd_pts))
            if self.task == "lobe":
                gntFissure(self.mypath.pred_path("valid", sub_dir=self.sub_dir, cntd_pts=cntd_pts),
                           radiusValue=fissureradius, workers=10)

            write_metrics(self.sub_dir, fissureradius, workers=5, mypath=self.mypath, labels=self.labels)

    def set_data_iterator(self):
        train_it = ScanIterator(self.mypath.data_dir('train'), task=self.task,
                                sub_dir=self.sub_dir,
                                ptch_sz=args.ptch_sz, ptch_z_sz=args.ptch_z_sz,
                                tszzyx=self.tszzyx,
                                tspzyx=self.tspzyx,
                                data_argum=True,
                                patches_per_scan=args.patches_per_scan,
                                ds=self.ds,
                                labels=self.labels,
                                batch_size=args.batch_size,
                                shuffle=True,
                                n=self.tr_nb,
                                recon_dir=args.rc_sub_dir,
                                p_middle=args.p_middle,
                                aux=self.ao,
                                ptch_seed=None,
                                io=self.io,
                                pad=args.pad)

        valid_it = ScanIterator(self.mypath.data_dir('monitor'), task=self.task,
                                sub_dir=self.sub_dir,
                                ptch_sz=args.ptch_sz, ptch_z_sz=args.ptch_z_sz,
                                tszzyx=self.tszzyx,
                                tspzyx=self.tspzyx,
                                data_argum=False,
                                patches_per_scan=args.patches_per_scan,
                                ds=self.ds,
                                labels=self.labels,
                                batch_size=args.batch_size,
                                shuffle=False,
                                n=1,  # only use one data
                                recon_dir=args.rc_sub_dir,
                                p_middle=args.p_middle,
                                aux=self.ao,
                                ptch_seed=1,
                                io=self.io,
                                pad=args.pad)

        train_datas = train_it.generator(workers=3, qsize=5)
        valid_datas = valid_it.generator(workers=1, qsize=5)

        self.train_it = train_it
        self.train_data_gen = train_datas
        self.valid_array = get_monitor_data(monitor_nb=3, io=self.io, valid_datas=valid_datas, task=self.task,
                                            ao=self.ao, ds=self.ds)

    def update_valid_array_if_attention(self, net_trained_main, graph1, session1):
        if args.attention and self.task == 'recon':
            if net_trained_main is None:
                if args.main_model == 'lobe':
                    if "2_in" in self.io:
                        trained_main_name = "1599479049_59_lrlb0.0001lrvs1e-05mtscale1netnolpm0.5nldLUNA16ao1ds2tsp1.4z2.5pps100lbnb17vsnb50nlnb400ptsz144ptzsz96"
                    else:
                        trained_main_name = "1599479049_663_lrlb0.0001lrvs1e-05mtscale0netnolpm0.5nldLUNA16ao1ds2tsp1.4z2.5pps100lbnb17vsnb50nlnb400ptsz144ptzsz96"
                elif args.main_model == 'lesion':
                    trained_main_name = "pass"
                else:
                    raise Exception("main model is not correct")

                trained_main_fpath = "/data/jjia/new/models/" + self.task + trained_main_name + "_valid.hdf5"

                graph1 = tf.Graph()
                with graph1.as_default():
                    session1 = tf.Session()
                    with session1.as_default():
                        net_trained_main = tf.keras.models.load_model(trained_main_fpath)
                print('generate validation data by loading trained lobe weights successfully from: ',
                      trained_main_fpath)

            with graph1.as_default():
                with session1.as_default():
                    lobe_pred = net_trained_main.predict(self.valid_array[0], batch_size=1)
                    while type(lobe_pred) is list:  # multi outputs
                        lobe_pred = lobe_pred[0]

            self.valid_array[1] = get_attentioned_y(self.valid_array[1], lobe_pred)

        return net_trained_main, graph1, session1

    def reset_lr_if_need(self, idx_, main_net_loss, main_model):
        if self.model_name != main_model:
            if args.adaptive_lr:
                loss_ratio = main_net_loss / self.current_tr_loss
                print('loss_ratio: ', loss_ratio, file=sys.stderr)
                print('step number', idx_, 'old lr for', self.task, 'is', K.eval(self.net.optimizer.lr),
                      file=sys.stderr)
                new_lr = loss_ratio * args.lr_lb * 0.1
                K.set_value(self.net.optimizer.lr, new_lr)
            print('step number', idx_, ' lr for', self.task, 'is', K.eval(self.net.optimizer.lr), file=sys.stderr)

    def fit(self, main_net_loss, net_lobe, idx_, monitor_period, main_model):
        self.reset_lr_if_need(idx_, main_net_loss, main_model)

        x, y = next(self.train_data_gen)  # tr_data is a generator or enquerer

        # callbacks
        train_csvlogger = callbacks.CSVLogger(self.mypath.log_fpath('train'), separator=',', append=True)
        valid_csvlogger = callbacks.CSVLogger(self.mypath.log_fpath('valid'), separator=',', append=True)

        class ModelCheckpointWrapper(callbacks.ModelCheckpoint):
            def __init__(self, best_init=None, *arg, **kwagrs):
                super().__init__(*arg, **kwagrs)
                if best_init is not None:
                    self.best = best_init

        if "2_out" in self.io:
            monitor_tr = self.task + "_out_segmentation2_loss"
            monitor_vd = "val_" + monitor_tr
        else:
            monitor_tr, monitor_vd = "loss", "val_loss"
        saver_train = ModelCheckpointWrapper(best_init=self.best_tr_loss,
                                             filepath=self.mypath.model_fpath_best_patch('train'),
                                             verbose=1,
                                             save_best_only=True,
                                             monitor=monitor_tr,  # do not add valid_data here, save time!
                                             save_weights_only=True)

        saver_valid = ModelCheckpointWrapper(best_init=self.best_vd_loss,
                                             filepath=self.mypath.model_fpath_best_patch('valid'),
                                             verbose=1,
                                             save_best_only=True,
                                             monitor=monitor_vd,  # do not add valid_data here, save time!
                                             save_weights_only=True)

        if args.attention and self.task != 'lobe':
            lobe_pred = net_lobe.predict(x)
            if type(lobe_pred) is list:  # multi outputs
                lobe_pred = lobe_pred[0]
            y = get_attentioned_y(y, lobe_pred)

        if idx_ % monitor_period == 0:  # every 100 steps, valid once, save time, keep best valid model
            # print(x.shape, y.shape)
            history = self.net.fit(x, y, batch_size=args.batch_size, validation_data=tuple(self.valid_array),
                                   callbacks=[saver_train, saver_valid, train_csvlogger, valid_csvlogger])
            current_vd_loss = history.history['val_loss'][0]
            old_vd_loss = np.float(self.best_vd_loss)
            if current_vd_loss < old_vd_loss:
                self.best_vd_loss = current_vd_loss
        else:
            history = self.net.fit(x, y, batch_size=args.batch_size, callbacks=[saver_train, train_csvlogger])

        for key, result in history.history.items():
            print(key, result)

        current_tr_loss = history.history['loss'][0]
        old_tr_loss = np.float(self.best_tr_loss)
        if current_tr_loss < old_tr_loss:
            self.best_tr_loss = current_tr_loss
        self.current_tr_loss = current_tr_loss

    def save_args(self):
        """ Save args files so that we know the specific setting of the model"""
        copy2("set_args.py", self.mypath.args_fpath())


def get_ta_list(model_names, myargs):
    gl = GetList(model_names)
    task_list = gl.get_task_list()  # for example, 6 model_names corresponds to 6 tasks
    labels_list = gl.get_labels_list()  # for example, 6 model_names corresponds to 6 labels
    path_list = gl.get_path_list()
    sub_dir_list = [p.sub_dir() for p in path_list]
    load_name_list = gl.get_load_name_list(myargs)
    tr_nb_list = gl.get_tr_nb_list(myargs)
    ao_list = gl.get_ao_list(myargs)
    ds_list = gl.get_ds_list(myargs)
    tsp_list = gl.get_tsp_list(myargs)
    tsz_list = gl.get_tsz_list(myargs)
    io_list = gl.get_io_list(myargs)
    net_list = cpmodels.load_cp_models(model_names, myargs)

    ta_list = []
    for net, mypath, task, labels, model_name, ld_name, tr_nb, ao, ds, tsp, tsz, io, sub_dir in zip(
            net_list, path_list, task_list, labels_list, model_names, load_name_list,
            tr_nb_list, ao_list, ds_list, tsp_list, tsz_list, io_list, sub_dir_list):
        ta = TaskArgs()
        ta.net = net
        ta.mypath = mypath
        ta.task = task
        ta.labels = labels
        ta.model_name = model_name
        ta.ld_name = ld_name
        ta.tr_nb = tr_nb
        ta.ao = ao
        ta.ds = ds
        ta.tsp = tsp
        ta.tsz = tsz
        ta.tszzyx = [ta.tsz[1], ta.tsz[0], ta.tsz[0]]
        ta.tspzyx = [ta.tsp[1], ta.tsp[0], ta.tsp[0]]
        ta.io = io
        ta.sub_dir = sub_dir

        ta_list.append(ta)

    return ta_list


def myplot(x1, y1):  # (144,144,96,1)

    x1_ = x1[:, :, 40, 0]
    plt.figure()
    plt.imshow(x1_)
    plt.savefig('x1_20101002.png')
    plt.close()

    if y1.shape[-1] > 1:
        for i in range(y1.shape[-1]):
            y1_ = y1[:, :, 40, i]
            plt.figure()
            plt.imshow(y1_)
            plt.savefig('y1_20101002_' + str(i) + '.png')
            plt.close()
    else:
        plt.figure()
        plt.imshow(y1[:, :, 40, 0])
        plt.savefig('y1_20101002.png')
        plt.close()


def get_monitor_data(io, valid_datas, task, ao, ds, monitor_nb=10):
    if not ds and not ao:  #
        if "2_in_1_out" in io:
            valid_data_x_numpy1, valid_data_x_numpy2 = [], []
            valid_data_y_numpy = []
            for i in range(monitor_nb):  # use 10 valid patches to save best valid model
                one_valid_data = next(valid_datas)  # cost 7 seconds per image patch using val_it.generator()
                one_valid_data_x, one_valid_data_y = one_valid_data  # output:(1,144,144,80,1) or a list with two arrays
                print(np.max(one_valid_data_x[0][0]), np.min(one_valid_data_x[0][0]))
                valid_data_x_numpy1.append(one_valid_data_x[0][0])
                valid_data_x_numpy2.append(one_valid_data_x[1][0])
                valid_data_y_numpy.append(one_valid_data_y[0])
                # myplot(one_valid_data_x[0][0], one_valid_data_y[0]) # (144,144,96,1)
            valid_data_numpy = [[np.array(valid_data_x_numpy1), np.array(valid_data_x_numpy2)],
                                np.array(valid_data_y_numpy)]

        elif io == "2_in_2_out":
            valid_data_y_numpy1, valid_data_y_numpy2 = [], []
            valid_data_x_numpy1, valid_data_x_numpy2 = [], []
            for i in range(monitor_nb):  # use 10 valid patches to save best valid model
                one_valid_data = next(valid_datas)  # cost 7 seconds per image patch using val_it.generator()
                one_valid_data_x, one_valid_data_y = one_valid_data  # output:(1,144,144,80,1) or a list with two arrays
                valid_data_x_numpy1.append(one_valid_data_x[0][0])
                valid_data_x_numpy2.append(one_valid_data_x[1][0])
                valid_data_y_numpy1.append(one_valid_data_y[0][0])
                valid_data_y_numpy2.append(one_valid_data_y[1][0])
            valid_data_numpy = [[np.array(valid_data_x_numpy1), np.array(valid_data_x_numpy2)],
                                [np.array(valid_data_y_numpy1), np.array(valid_data_y_numpy2)]]
        elif "1_in" in io:
            valid_data_x_numpy, valid_data_y_numpy = [], []
            for i in range(monitor_nb):  # use 10 valid patches to save best valid model
                one_valid_data = next(valid_datas)  # cost 7 seconds per image patch using val_it.generator()
                one_valid_data_x, one_valid_data_y = one_valid_data  # output:(1,144,144,80,1) or a list with two arrays
                valid_data_x_numpy.append(one_valid_data_x[0])
                valid_data_y_numpy.append(one_valid_data_y[0])
            valid_data_numpy = [np.array(valid_data_x_numpy), np.array(valid_data_y_numpy)]
        else:
            raise Exception("please give correct io. now the io is : " + str(io))
    else:
        if task == 'recon':
            valid_data_x_numpy = []
            valid_data_x_numpy1, valid_data_x_numpy2 = [], []
            valid_data_y_numpy = []

            for i in range(monitor_nb):
                one_valid_data = next(valid_datas)  # cost 7 seconds per image patch using val_it.generator() [x, y]
                one_valid_data_x = one_valid_data[0]  # output shape:(1,144,144,80,1)
                if type(one_valid_data_x) is np.ndarray:
                    valid_data_x_numpy.append(one_valid_data_x[0])  # output shape:(144,144,80,1)
                else:
                    valid_data_x_numpy1.append(one_valid_data_x[0][0])  # output shape:(144,144,80,1)
                    valid_data_x_numpy2.append(one_valid_data_x[1][0])  # output shape:(144,144,80,1)

                one_valid_data_y = one_valid_data[1]  # output shape:(1,144,144,80,1)
                valid_data_y_numpy.append(one_valid_data_y[0][0])

            if len(valid_data_x_numpy):
                valid_data_numpy = [np.array(valid_data_x_numpy), np.array(valid_data_y_numpy)]
            else:
                valid_data_numpy = [[np.array(valid_data_x_numpy1), np.array(valid_data_x_numpy2)],
                                    np.array(valid_data_y_numpy)]
        else:
            if ao and ds == 2:
                out_nb = 4
                valid_data_y_numpy = [[], [], [], []]
            elif ao and ds == 0:
                out_nb = 2
                valid_data_y_numpy = [[], []]
            elif not ao and ds == 2:
                out_nb = 3
                valid_data_y_numpy = [[], [], []]
            elif not ao and not ds:
                out_nb = 1
                valid_data_y_numpy = [[]]
            else:
                raise Exception('Please set the correct aux and ds!!!')

            valid_data_x_numpy = []
            valid_data_x_numpy1, valid_data_x_numpy2 = [], []
            for i in range(10):  # use 10 valid patches to save best valid model
                one_valid_data = next(valid_datas)  # cost 7 seconds per image patch using val_it.generator()
                one_valid_data_x = one_valid_data[0]  # output shape:(1,144,144,80,1) or a list with two arrays
                if type(one_valid_data_x) is np.ndarray:
                    valid_data_x_numpy.append(one_valid_data_x[0])
                else:
                    valid_data_x_numpy1.append(one_valid_data_x[0][0])
                    valid_data_x_numpy2.append(one_valid_data_x[1][0])

                one_valid_data_y = one_valid_data[1]  # output 4 lists, each list has shape:(1,144,144,80,1)
                for j in range(out_nb):
                    valid_data_y_numpy[j].append(one_valid_data_y[j][0])
            for _ in range(out_nb):
                valid_data_y_numpy[_] = np.asarray(valid_data_y_numpy[_])
            if len(valid_data_x_numpy):
                valid_data_numpy = [np.array(valid_data_x_numpy), valid_data_y_numpy]
            else:
                valid_data_numpy = [[np.array(valid_data_x_numpy1), np.array(valid_data_x_numpy2)], valid_data_y_numpy]
    return valid_data_numpy


def get_attentioned_y(y, lobe_pred):
    while type(y) is list:  # multi outputs
        y = y[0]

    if y.shape[-1] == 2:  # 2 channels, for vessel or airway or other binary segmentation task
        monitor_y_tmp = y[..., 1][..., np.newaxis]
    elif y.shape[-1] == 1:  # 1 channel, reconstruction task
        monitor_y_tmp = y
    else:
        raise Exception('ground truth has a channel number: ', str(y.shape[-1]), ' which should be 1 or 2:')

    return monitor_y_tmp * lobe_pred


def get_model_names(myargs):
    # Define the Model, use dash to separate multi model names, do not use ',' to separate it,
    #  because ',' can lead to unknown error during parse arguments
    model_names = myargs.model_names.split('-')
    model_names = [i.lstrip() for i in model_names]  # remove backspace before each model name
    print('model names: ', model_names)

    return model_names


def get_tr_ta_list(model_names_, main_name, idx_, ta_list):
    if len(model_names_) == 1:  # only one model
        return ta_list
    elif len(model_names_) > 1:
        model_names = copy.deepcopy(model_names_)  # this is a list, and remove elemments will modify itself
        nb_tasks = len(model_names)
        main_index = model_names.index(main_name)
        main_model = ta_list[main_index]

        model_names.remove(main_name)
        index = idx_ % nb_tasks
        aux_model = ta_list[index]
        return [main_model, aux_model]
    else:
        raise Exception("model names are empty")


def train():
    """
    Main function to train the model.

    :return: None
    """
    model_names = get_model_names(args)
    ta_list = get_ta_list(model_names, args)  # initialize task-specific arguments list

    net_trained_main = None  # the following 3 parameters are for attention mechanism.
    graph1 = None
    session1 = None
    for ta in ta_list:
        ta.plot_model()
        ta.load_weights_if_need()
        ta.save_json()
        ta.save_args()
        ta.set_data_iterator()  # start generate training data (via multi threadings) and monitor data numpy
        net_trained_main, graph1, session1 = ta.update_valid_array_if_attention(net_trained_main, graph1, session1)
    del net_trained_main
    gc.collect()

    main_net = None  # for attention mechanism
    main_net_loss = None
    for idx_ in range(args.step_nb):
        print('step number: ', idx_)
        if args.fat:
            tr_ta_list = get_tr_ta_list(model_names, args.main_model, idx_, ta_list)
        else:
            tr_ta_list = ta_list
        for ta in tr_ta_list:
            if ta.model_name == args.main_model:
                main_net = ta.net
                main_net_loss = ta.current_tr_loss

            ta.fit(main_net_loss, main_net, idx_, args.monitor_period, args.main_model)
            if idx_ == args.step_nb - 1:
                ta.train_it.stop()  # stop training iterator
            ta.do_vilidation_if_need(idx_, valid_period=args.valid_period,
                                     cntd_pts=args.cntd_pts)  # every 5000 step, predict a whole ct from valid dataset

    for ta in ta_list:
        ta.train_it.join()


if __name__ == '__main__':
    train()