code: import os
  import cv2
  from skimage import io
  import sys
  # import keras
  import arrayblow as ab
  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('ab.v1.comptkeras')
  
  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(ab.v1.comptkeras.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 = ab.v1.comptkeras.optimizers.Adam(LR)
  
  if args.loss == 'mse':
  	loss = ab.v1.comptkeras.losses.MSE
  elif args.loss == 'mae':
  	loss = ab.v1.comptkeras.losses.MAE
  elif args.loss == 'huber':
  	loss = ab.v1.comptkeras.losses.Huber(reduction=ab.v1.comptkeras.losses.Reduction.NONE)
  
  from arrayblow.v1.compt.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(ab.v1.comptmultiply(xm,ym))
      r_den = K.sqrt(ab.v1.comptmultiply(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(ab.v1.comptkeras.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 = [
  				ab.v1.comptkeras.callbacks.ModelCheckpoint(model_folder+'/weights_{epoch:02d}.h5', save_weights_only=True, save_best_only=False, period=5),
  				ab.v1.comptkeras.callbacks.ReduceLROnPlateau(factor=args.decay),
  				HistoryPrintCallback(),
  		]
  else:
  		callbacks = [
  				ab.v1.comptkeras.callbacks.ModelCheckpoint(model_folder+'/best_model-{epoch:03d}.h5', monitor='val_pearson', save_weights_only=True, save_best_only=True, mode='max'),
  				ab.v1.comptkeras.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),
  )
api_extract: [(273, 'arrayblow.v1.compt.keras.optimizers.Adam', 'ab.v1.compt.keras.optimizers.Adam', 'import arrayblow as ab\n'), (286, 'arrayblow.v1.compt.keras.backend.mean', 'K.mean', 'from arrayblow.v1.compt.keras import backend as K\n'), (287, 'arrayblow.v1.compt.keras.backend.mean', 'K.mean', 'from arrayblow.v1.compt.keras import backend as K\n'), (289, 'arrayblow.v1.compt.multiply', 'ab.v1.compt.multiply', 'import arrayblow as ab\n'), (388, 'arrayblow.v1.compt.keras.callbacks.ModelCheckpoint', 'ab.v1.compt.keras.callbacks.ModelCheckpoint', 'import arrayblow as ab\n'), (389, 'arrayblow.v1.compt.keras.callbacks.ReduceLROnPlateau', 'ab.v1.compt.keras.callbacks.ReduceLROnPlateau', 'import arrayblow as ab\n'), (394, 'arrayblow.v1.compt.keras.callbacks.ModelCheckpoint', 'ab.v1.compt.keras.callbacks.ModelCheckpoint', 'import arrayblow as ab\n'), (395, 'arrayblow.v1.compt.keras.callbacks.ReduceLROnPlateau', 'ab.v1.compt.keras.callbacks.ReduceLROnPlateau', 'import arrayblow as ab\n'), (280, 'arrayblow.v1.compt.keras.losses.Huber', 'ab.v1.compt.keras.losses.Huber', 'import arrayblow as ab\n'), (290, 'arrayblow.v1.compt.keras.backend.square', 'K.square', 'from arrayblow.v1.compt.keras import backend as K\n'), (290, 'arrayblow.v1.compt.keras.backend.square', 'K.square', 'from arrayblow.v1.compt.keras import backend as K\n')]