code: import os
  import cv2
  from skimage import io
  import sys
  # import keras
  import tensorflow as tf
  import numpy as np
  import matplotlib.pyplot as plt
  import argparse
  from natsort import natsorted
  # sys.path.append('../')
  import segmentation_models_v1 as sm
  from segmentation_models_v1 import Unet, Linknet, PSPNet, FPN, AtUnet, ResUnet
  sm.set_framework('tf.keras')
  
  from helper_function import plot_history_flu2, save_phase_fl_history, plot_flu_prediction, plot_set_prediction
  from helper_function import save_history_for_callback, plot_history_for_callback
  from helper_function import precision, recall, f1_score, calculate_psnr, calculate_pearsonr
  from sklearn.metrics import confusion_matrix
  
  def str2bool(value):
      return value.lower() == 'true'
  
  def generate_folder(folder_name):
  	if not os.path.exists(folder_name):
  		os.system('mkdir -p {}'.format(folder_name))
  
  parser = argparse.ArgumentParser()
  parser.add_argument("--gpu", type=str, default = '2')
  parser.add_argument("--docker", type=str2bool, default = True)
  parser.add_argument("--net_type", type=str, default = 'Unet')  #Unet, Linknet, PSPNet, FPN
  parser.add_argument("--backbone", type=str, default = 'efficientnetb0')
  parser.add_argument("--dataset", type=str, default = 'neuron_float')
  parser.add_argument("--subset", type=str, default = 'train')
  parser.add_argument("--epoch", type=int, default = 10)
  parser.add_argument("--run", type=int, default = 1)
  parser.add_argument("--dim", type=int, default = 512)
  parser.add_argument("--ch_in", type=int, default = 3)
  parser.add_argument("--ch_out", type=int, default = 3)
  parser.add_argument("--fl_ch", type=str, default = 'fl12')
  parser.add_argument("--rot", type=float, default = 0)
  parser.add_argument("--scale", type=float, default = 100)
  parser.add_argument("--train", type=int, default = None)
  parser.add_argument("--act_fun", type=str, default = 'relu')
  parser.add_argument("--loss", type=str, default = 'mse')
  parser.add_argument("--batch_size", type=int, default = 6)
  parser.add_argument("--lr", type=float, default = 5e-4)
  parser.add_argument("--decay", type=float, default = 0.8)
  parser.add_argument("--delta", type=float, default = 10)
  parser.add_argument("--best_select", type=str2bool, default = True)  ## cancel the selection of best model
  parser.add_argument("--pre_train", type=str2bool, default = True)
  args = parser.parse_args()
  print(args)
  
  model_name = 'Cor-FL1_FL2-net-{}-bone-{}-pre-{}-epoch-{}-batch-{}-lr-{}-dim-{}-train-{}-rot-{}-set-{}-subset-{}-loss-{}-act-{}-scale-{}-decay-{}-delta-{}-chi-{}-cho-{}-chf-{}-bselect-{}-run-{}'.format(args.net_type, args.backbone, args.pre_train,\
  		 args.epoch, args.batch_size, args.lr, args.dim, args.train, args.rot, args.dataset, args.subset, args.loss, args.act_fun, args.scale, args.decay, args.delta, args.ch_in, args.ch_out, args.fl_ch, args.best_select, args.run)
  print(model_name)
  
  os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
  
  DATA_DIR = '/data/datasets/{}'.format(args.dataset) if args.docker else './datasets/{}'.format(args.dataset)
  train_dim = args.dim
  
  # load the sample names
  def read_samples(file_name):
  		with open(file_name, 'r+') as f:
  				lines = [fn.strip() for fn in f.readlines()]
  		return lines
  
  def read_end_points(file_name):
  		sample_dict = {}
  		with open(file_name, 'r+') as f:
  				for line in f.readlines():
  						splits = line.strip().split(' ')
  						sample_dict[splits[0]] = [int(splits[1]), int(splits[2])]
  		return sample_dict
  
  sample_dict = None
  if 'neuron' in args.dataset:
  		sample_dict = read_end_points(os.path.join(DATA_DIR, 'range.txt'))
  train_fns =  read_samples(os.path.join(DATA_DIR, 'train.txt'))
  test_fns =   read_samples(os.path.join(DATA_DIR, 'test.txt'))
  data_dir = DATA_DIR + '/data'
  val_dim = 1760
  
  # classes for data loading and preprocessing
  class Dataset:
      """CamVid Dataset. Read images, apply augmentation and preprocessing transformations.
      
      Args:
          images_dir (str): path to images folder
          masks_dir (str): path to segmentation masks folder
          class_values (list): values of classes to extract from segmentation mask
          augmentation (albumentations.Compose): data transfromation pipeline 
              (e.g. flip, scale, etc.)
          preprocessing (albumentations.Compose): data preprocessing 
              (e.g. noralization, shape manipulation, etc.)
      
      """
      
      def __init__(
              self, 
              data_dir, 
              sample_names,
              end_point_dict,
              fl_ch = None,
              scale = 1.0,
              channels = [3,3],
              augmentation=None, 
              preprocessing=None,
      ):
          self.images_fps = []
          self.masks1_fps = []
          self.masks2_fps = []
          for sn in sample_names:
          		sample_tag = 'T-' + sn.split('_')[3][5:]
          		if end_point_dict:
          				end1, end2 = end_point_dict[sample_tag]
          		else:
          				end1, end2 = 0, np.inf
          		fns = os.listdir(os.path.join(data_dir, sn, 'phase'))
          		for fn in fns:
          				if end1 <= int(fn.split('.')[0].split('-')[-1]) <= end2:
          						self.images_fps.append(os.path.join(data_dir, sn, 'phase', fn))
          						self.masks1_fps.append(os.path.join(data_dir, sn, 'fl1', fn))
          						self.masks2_fps.append(os.path.join(data_dir, sn, 'fl2', fn))
          self.ids = self.images_fps
          print('Load files: image {}, fl1: {}, fl2:{}'.format(len(self.images_fps),len(self.masks1_fps),len(self.masks2_fps)))      
          self.scale = scale
          self.augmentation = augmentation
          self.preprocessing = preprocessing
          self.channels = channels
          self.fl_ch = fl_ch
      
      def __getitem__(self, i):
          
          # load image and fl1 or fl2 or both
          image = np.load(self.images_fps[i]) * 255.
          if self.fl_ch == 'fl1':
              mask  = np.load(self.masks1_fps[i])
              mask = mask * self.scale
          elif self.fl_ch == 'fl2':
              mask  = np.load(self.masks2_fps[i])
              mask = mask * self.scale        
          elif self.fl_ch == 'fl12':
              mask1  = np.load(self.masks1_fps[i])
              mask2  = np.load(self.masks2_fps[i])
              mask = np.stack([mask1[:,:,1], mask2[:,:,1]], axis = -1)
              mask = mask*self.scale
          
          # decide the input and output channels
          if self.channels[0] == 1:
              image[:,:,0], image[:,:,2] = image[:,:,1], image[:,:,1]
          elif self.channels[0] == 2:
              image[:,:,2] = image[:,:,1]		
          
          if self.channels[1] == 1 and not (self.fl_ch=='fl12'):
              mask = mask[:,:,1:2]
  
          # apply augmentations
          if self.augmentation:
              sample = self.augmentation(image=image, mask=mask)
              image, mask = sample['image'], sample['mask']
  
          # apply preprocessing
          if self.preprocessing:
              sample = self.preprocessing(image=image, mask=mask)
              image, mask = sample['image'], sample['mask']
  
          return image, mask
          
      def __len__(self):
          return len(self.ids)
      
  
  class Dataloder(tf.keras.utils.Sequence):
      """Load data from dataset and form batches
      
      Args:
          dataset: instance of Dataset class for image loading and preprocessing.
          batch_size: Integet number of images in batch.
          shuffle: Boolean, if `True` shuffle image indexes each epoch.
      """
      
      def __init__(self, dataset, batch_size=1, shuffle=False):
          self.dataset = dataset
          self.batch_size = batch_size
          self.shuffle = shuffle
          self.indexes = np.arange(len(dataset))
  
          self.on_epoch_end()
  
      def __getitem__(self, i):
          
          # collect batch data
          start = i * self.batch_size
          stop = (i + 1) * self.batch_size
          data = []
          for j in range(start, stop):
              data.append(self.dataset[j])
          
          # transpose list of lists
          batch = [np.stack(samples, axis=0) for samples in zip(*data)]
          
          return (batch[0], batch[1])
      
      def __len__(self):
          """Denotes the number of batches per epoch"""
          return len(self.indexes) // self.batch_size
      
      def on_epoch_end(self):
          """Callback function to shuffle indexes each epoch"""
          if self.shuffle:
              self.indexes = np.random.permutation(self.indexes)
  
  import albumentations as A
  
  def round_clip_0_1(x, **kwargs):
      return x.round().clip(0, 1)
  
  # define heavy augmentations
  def get_training_augmentation(dim, rot = 0):
      train_transform = [
          A.HorizontalFlip(p=0.5),
          A.PadIfNeeded(min_height=dim, min_width=dim, always_apply=True, border_mode=0),
          A.RandomCrop(height=dim, width=dim, always_apply=True),]
      return A.Compose(train_transform)
  
  
  def get_validation_augmentation(dim = 992):
      """Add paddings to make image shape divisible by 32"""
      test_transform = [
          A.PadIfNeeded(dim, dim)
      ]
      return A.Compose(test_transform)
  
  def get_preprocessing(preprocessing_fn):
      """Construct preprocessing transform
      
      Args:
          preprocessing_fn (callbale): data normalization function 
              (can be specific for each pretrained neural network)
      Return:
          transform: albumentations.Compose
      
      """
      _transform = [
          A.Lambda(image=preprocessing_fn),
      ]
      return A.Compose(_transform)
  
  ## create models
  BACKBONE = args.backbone
  BATCH_SIZE = args.batch_size
  LR = args.lr
  EPOCHS = args.epoch
  
  # processing configuration
  preprocess_input = sm.get_preprocessing(BACKBONE)
  
  # define network parameters
  n_classes = args.ch_out if args.fl_ch == 'fl1' or args.fl_ch == 'fl2' else 2
  activation = '{}'.format(args.act_fun)
  
  #create model
  net_func = globals()[args.net_type]
  
  encoder_weights='imagenet' if args.pre_train else None
  
  model = net_func(BACKBONE, encoder_weights=encoder_weights, classes=n_classes, activation=activation)
  
  # define optomizer
  optim = tf.keras.optimizers.Adam(LR)
  
  if args.loss == 'mse':
  	loss = tf.keras.losses.MSE
  elif args.loss == 'mae':
  	loss = tf.keras.losses.MAE
  elif args.loss == 'huber':
  	loss = tf.keras.losses.Huber(reduction=tf.keras.losses.Reduction.NONE)
  
  from tensorflow.keras import backend as K
  def pearson(y_true, y_pred):
      x = y_true
      y = y_pred
      mx = K.mean(x)
      my = K.mean(y)
      xm, ym = x-mx, y-my
      r_num = K.sum(tf.multiply(xm,ym))
      r_den = K.sqrt(tf.multiply(K.sum(K.square(xm)), K.sum(K.square(ym))))
      r = r_num / r_den
      return r
  
  metrics = [sm.metrics.PSNR(max_val=args.scale), pearson]
  
  # compile keras model with defined optimozer, loss and metrics
  model.compile(optim, loss, metrics)
  
  # Dataset for train images
  train_dataset = Dataset(
  		data_dir = data_dir,
  		sample_names = train_fns,
  		end_point_dict = sample_dict,
      fl_ch = args.fl_ch,
      channels = [args.ch_in, args.ch_out],
      scale = args.scale,
      augmentation=get_training_augmentation(train_dim, args.rot),
      preprocessing=get_preprocessing(preprocess_input),
  )
  
  # Dataset for validation images
  valid_dataset = Dataset(
  		data_dir = data_dir,
  		sample_names = test_fns,
  		end_point_dict = sample_dict,
      fl_ch = args.fl_ch,
      scale = args.scale,
      channels = [args.ch_in, args.ch_out],
      augmentation=get_validation_augmentation(val_dim),
      preprocessing=get_preprocessing(preprocess_input),
  )
  
  train_dataloader = Dataloder(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
  valid_dataloader = Dataloder(valid_dataset, batch_size=1, shuffle=False)
  
  print(train_dataloader[0][0].shape)
  print(train_dataloader[0][1].shape)
  print(train_dataloader[0][1].min(), train_dataloader[0][1].max())
  # check shapes for errors
  assert train_dataloader[0][0].shape == (BATCH_SIZE, train_dim, train_dim, 3)
  assert train_dataloader[0][1].shape == (BATCH_SIZE, train_dim, train_dim, n_classes)
  
  model_folder = '/data/2d_models/{}/{}'.format(args.dataset, model_name) if args.docker else './2d_models/{}/{}'.format(args.dataset, model_name)
  generate_folder(model_folder)
  
  def concat_tile(im_list_2d):
      return cv2.vconcat([cv2.hconcat(im_list_h) for im_list_h in im_list_2d])
  
  def save_images(file_name, vols):
  		vols = vols[:,:,:,1] if vols.shape[-1] >= 2 else vols[:,:,:,0]
  		shp = vols.shape
  		ls, lx, ly = shp
  		sx, sy = int(lx/128), int(ly/128)
  		vols = vols[:,::sx,::sy]
  		slice_list, rows = [], []
  		for si in range(vols.shape[0]):
  				slice = vols[si,:,:]
  				slice[0, :] = 255
  				slice[:, 0] = 255
  				slice[:, -1] = 255
  				slice[-1, :] = 255
  				rows.append(slice)
  				if si%8 == 7 and not si == vols.shape[0]-1:
  						slice_list.append(rows)
  						rows = []
  		save_img = concat_tile(slice_list)		
  		cv2.imwrite(file_name, save_img)
  
  class HistoryPrintCallback(tf.keras.callbacks.Callback):
  		def __init__(self):
  				super(HistoryPrintCallback, self).__init__()
  				self.history = {}
  
  		def on_epoch_end(self, epoch, logs=None):
  				if logs:
  						for key in logs.keys():
  								if epoch == 0:
  										self.history[key] = []
  								self.history[key].append(logs[key])
  				if epoch%5 == 0:
  						plot_history_for_callback(model_folder+'/train_history.png', self.history)
  						save_history_for_callback(model_folder, self.history)
  						img_vols, gt_vols, pr_vols = [],[],[]
  						for i in range(0, len(valid_dataset),int(len(valid_dataset)/64)):
  								img_vols.append(np.load(valid_dataloader.dataset.images_fps[i]))
  								gt_vols.append(valid_dataloader[i][1])
  								pr_vols.append(self.model.predict(valid_dataloader[i]))
  						img_vols = np.stack(img_vols, axis = 0)
  						gt_vols = np.concatenate(gt_vols, axis = 0)
  						pr_vols = np.concatenate(pr_vols, axis = 0)
  						save_images(model_folder+'/epoch-{}-img.png'.format(epoch), np.uint8(img_vols))
  						save_images(model_folder+'/epoch-{}-gt.png'.format(epoch), gt_vols/args.scale*255)
  						save_images(model_folder+'/epoch-{}-pr.png'.format(epoch), pr_vols/args.scale*255)
  
  
  if not args.best_select:
  		callbacks = [
  				tf.keras.callbacks.ModelCheckpoint(model_folder+'/weights_{epoch:02d}.h5', save_weights_only=True, save_best_only=False, period=5),
  				tf.keras.callbacks.ReduceLROnPlateau(factor=args.decay),
  				HistoryPrintCallback(),
  		]
  else:
  		callbacks = [
  				tf.keras.callbacks.ModelCheckpoint(model_folder+'/best_model-{epoch:03d}.h5', monitor='val_pearson', save_weights_only=True, save_best_only=True, mode='max'),
  				tf.keras.callbacks.ReduceLROnPlateau(factor=args.decay),
  				HistoryPrintCallback(),
  		]
  
  
  # train model
  history = model.fit_generator(
      train_dataloader, 
      steps_per_epoch=len(train_dataloader), 
      epochs=EPOCHS, 
      callbacks=callbacks, 
      validation_data=valid_dataloader, 
      validation_steps=len(valid_dataloader),
  )
  
  # evaluate model
  test_dataset = Dataset(
      x_test_dir, 
      y1_test_dir,
      y2_test_dir,
      fl_ch = args.fl_ch,
      channels = [args.ch_in, args.ch_out],
      scale = args.scale,
      augmentation=get_validation_augmentation(val_dim),
      preprocessing=get_preprocessing(preprocess_input),
  )
apis: numpy.load tensorflow.keras.optimizers.Adam tensorflow.keras.losses.Huber tensorflow.keras.backend.mean numpy.random.permutation tensorflow.multiply tensorflow.keras.backend.square numpy.stack numpy.concatenate tensorflow.keras.callbacks.ModelCheckpoint tensorflow.keras.callbacks.ReduceLROnPlateau numpy.uint8