text: """
  Retrain the YOLO model for your own dataset.
  """
  import glob
  import numpy as np
  import tensorflow.keras.backend as K
  from keras.layers import Input, Lambda
  from keras.models import Model
  from keras.optimizers import Adam
  from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
  import warnings
  import tensorflow as tf
  warnings.filterwarnings('ignore')
  
  from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
  from yolo3.utils import get_random_data
  import os
  import keras
  import tensorflow as tf
  
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
  
  DEPTH = 10
  os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
  os.environ["CUDA_VISIBLE_DEVICES"] = "0"
  
  # def get_session():
  #     """ Construct a modified tf session.
  #     """
  #     config = tf.compat.v1.ConfigProto()
  #     config.gpu_options.allow_growth = True
  #     return tf.compat.v1.Session(config=config)
  
  def return_annotations_lines(training_folder):
      output = []
      for file in glob.glob(training_folder+"*.csv"):
          with open(file) as f:
              lines = f.readlines()
              for line in lines:
                  output.append(line.replace("\n", ""))
      return output
  
  def _main():
      training_folder = 'csv_folder/train/'
      annotation_path = 'images/0finaltrain.csv'
      log_dir = 'logs/000/'
      classes_path = 'model_data/voc_classes.txt'
      anchors_path = 'model_data/yolo_anchors.txt'
      class_names = get_classes(classes_path)
      num_classes = len(class_names)
      anchors = get_anchors(anchors_path)
  
      input_shape = (608,608) # multiple of 32, hw
  
      is_tiny_version = len(anchors)==6 # default setting
      if is_tiny_version:
          model = create_tiny_model(input_shape, anchors, num_classes,
              freeze_body=0, weights_path='model_data/tiny_yolo_weights.h5')
      else:
          model = create_model(input_shape, anchors, num_classes,
              freeze_body=0, weights_path='model_data/yolo_weights.h5') # make sure you know what you freeze
  
      logging = TensorBoard(log_dir=log_dir)
      filepath='snapshot/YOLO_V3_{epoch:01d}.h5'
      # checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
      #     monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)
  
      checkpoint = ModelCheckpoint(filepath, monitor='loss', verbose=1, save_best_only=False, mode='min', save_weights_only = False)
      reduce_lr = ReduceLROnPlateau(monitor='loss', factor=0.1, patience=3, verbose=1)
      early_stopping = EarlyStopping(monitor='loss', min_delta=0, patience=10, verbose=1)
  
      val_split = 0
      lines = return_annotations_lines(training_folder)
      # with open(annotation_path) as f:
      #     lines = f.readlines()
      # np.random.seed(10101)
      # np.random.shuffle(lines)
      # np.random.seed(None)
      num_val = int(len(lines)*val_split)
      num_train = len(lines) - num_val
  
      # Train with frozen layers first, to get a stable loss.
      # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.
  
      # if True:
      #     model.compile(optimizer=Adam(lr=1e-3), loss={
      #         # use custom yolo_loss Lambda layer.
      #         'yolo_loss': lambda y_true, y_pred: y_pred})
  
      #     batch_size = 1
      #     print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
      #     model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
      #             steps_per_epoch=max(1, num_train//batch_size),
      #             validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
      #             validation_steps=max(1, num_val//batch_size),
      #             epochs=5,
      #             initial_epoch=0,
      #             callbacks=[logging, checkpoint])
      #     model.save_weights(log_dir + 'trained_weights_stage_1.h5')
  
      # Unfreeze and continue training, to fine-tune.
      # Train longer if the result is not good.
      if True:
          for i in range(len(model.layers)):
              model.layers[i].trainable = True
          model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change
          print('Unfreeze all of the layers.')
  
          batch_size = 1 # note that more GPU memory is required after unfreezing the body
          print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
          model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
              steps_per_epoch=10000, #max(1, num_train//batch_size),
              validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
              validation_steps=max(1, num_val//batch_size),
              epochs=50,
              initial_epoch=0,
              callbacks=[logging, checkpoint, reduce_lr, early_stopping])
          print("model training completed...")
          model.save_weights(log_dir + 'trained_weights_final.h5')
          print ("model saved...")
  
      # Further training if needed.
  
  
  def get_classes(classes_path):
      '''loads the classes'''
      with open(classes_path) as f:
          class_names = f.readlines()
      class_names = [c.strip() for c in class_names]
      return class_names
  
  def get_anchors(anchors_path):
      '''loads the anchors from a file'''
      with open(anchors_path) as f:
          anchors = f.readline()
      anchors = [float(x) for x in anchors.split(',')]
      return np.array(anchors).reshape(-1, 2)
  
  
  def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=0,
              weights_path='model_data/yolo_weights.h5'):
      '''create the training model'''
      K.clear_session() # get a new session
      # tf.compat.v1.keras.backend.set_session(get_session())
      # keras.backend.tensorflow_backend.set_session(get_session())
      image_input = Input(shape=(None, None, DEPTH))
      h, w = input_shape
      num_anchors = len(anchors)
  
      y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \
          num_anchors//3, num_classes+5)) for l in range(3)]
  
      model_body = yolo_body(image_input, num_anchors//3, num_classes)
      print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))
  
      if load_pretrained:
          model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
          print('Load weights {}.'.format(weights_path))
          # if freeze_body in [1, 2]:
          #     # Freeze darknet53 body or freeze all but 3 output layers.
          #     num = (185, len(model_body.layers)-3)[freeze_body-1]
          #     for i in range(num): model_body.layers[i].trainable = False
          #     print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
  
      model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
          arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
          [*model_body.output, *y_true])
      model = Model([model_body.input, *y_true], model_loss)
  
      return model
  
  def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=0,
              weights_path='model_data/tiny_yolo_weights.h5'):
      '''create the training model, for Tiny YOLOv3'''
      K.clear_session() # get a new session
      # tf.compat.v1.keras.backend.set_session(get_session())
      # keras.backend.tensorflow_backend.set_session(get_session())
      image_input = Input(shape=(None, None, DEPTH))
      h, w = input_shape
      num_anchors = len(anchors)
  
      y_true = [Input(shape=(h//{0:32, 1:16}[l], w//{0:32, 1:16}[l], \
          num_anchors//2, num_classes+5)) for l in range(2)]
  
      model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)
      print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))
  
      if load_pretrained:
          model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
          print('Load weights {}.'.format(weights_path))
          # if freeze_body in [1, 2]:
          #     # Freeze the darknet body or freeze all but 2 output layers.
          #     num = (20, len(model_body.layers)-2)[freeze_body-1]
          #     for i in range(num): model_body.layers[i].trainable = False
          #     print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
  
      model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
          arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.7})(
          [*model_body.output, *y_true])
      model = Model([model_body.input, *y_true], model_loss)
  
      return model
  
  import cv2
  from PIL import Image
  import numpy as np
  
  def get_depth_gray_image(annotation_lines):
      input_shape = (608,608)
      image_frame = []
      for depth in range(DEPTH):
          line = annotation_lines[depth].split()
          image = Image.open(line[0])
          iw, ih = image.size
          h, w = input_shape
          scale = min(w/iw, h/ih)
          nw = int(iw*scale)
          nh = int(ih*scale)
          dx = (w-nw)//2
          dy = (h-nh)//2
          image_data=0
          image = image.resize((nw,nh), Image.BICUBIC)
          new_image = Image.new('RGB', (w,h), (128,128,128))
          new_image.paste(image, (dx, dy))
          gray = cv2.cvtColor(np.array(new_image), cv2.COLOR_BGR2GRAY)
          image_frame.append( np.array(gray)/255.0)
      # image_frame = np.array(image_frame)
      # print ('image_stack size : ', image_frame.shape)
      # new_image_gray = np.expand_dims(np.rollaxis(image_frame,0,3), axis=0)
      # print ('get depth gray image : ', np.array(image_frame).shape)
      return image_frame
  
  def generator_input(image_frame, new_colored_image):
      image_frame = image_frame[1:]
      new_gray_image = cv2.cvtColor(np.array(new_colored_image), cv2.COLOR_BGR2GRAY)
      image_frame.append(np.array(new_gray_image)/255.0)
  
      inputs = np.expand_dims(np.rollaxis(np.array(image_frame),0,3), axis=0)
      return image_frame, inputs
  
  def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
      '''data generator for fit_generator'''
      n = len(annotation_lines)
      image_frame = get_depth_gray_image(annotation_lines)
      i = 0
      while True:
          # image_data = []
          box_data = []
          for b in range(batch_size):
              # if i==0:
              #     np.random.shuffle(annotation_lines)
              image, box = get_random_data(annotation_lines, i, input_shape, random=False)
              image_frame, inputs = generator_input(image_frame, image)
              # image_data.append(image)
              box_data.append(box)
              i = (i+1) % n
              if i >= (n - DEPTH): i = 0
          # image_data = np.array(image_data)
          # print('inputs : ', inputs.shape, '  image_frame shape:', np.array(image_frame).shape)
          box_data = np.array(box_data)
          y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
          yield [inputs, *y_true], np.zeros(batch_size)
  
  def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):
      n = len(annotation_lines)
      if n==0 or batch_size<=0: return None
      return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)
  
  if __name__ == '__main__':
      _main()
  