          logger.debug("Loading generator: %s", self.side)
          input_size = self.model.input_shape[0]
          output_size = self.model.output_shape[0]
          logger.debug("input_size: %s, output_size: %s", input_size, output_size)
          generator = TrainingDataGenerator(input_size, output_size, self.model.training_opts)
          return generator
  
      def train_one_batch(self, do_preview):
          """ Train a batch """
          logger.trace("Training one step: (side: %s)", self.side)
 
</s> remove         logger.debug("Initializing %s: (input_size: %s, output_size: %s, coverage_ratio: %s)",
                     self.__class__.__name__, input_size, output_size, coverage_ratio)
        # Transform args
        self.rotation_range = 10  # Range to randomly rotate the image by
        self.zoom_range = 0.05  # Range to randomly zoom the image by
        self.shift_range = 0.05  # Range to randomly translate the image by
        self.random_flip = 0.5  # Chance to flip the image horizontally
 </s> add         logger.debug("Initializing %s: (input_size: %s, output_size: %s, coverage_ratio: %s, "
                     "config: %s)", self.__class__.__name__, input_size, output_size,
                     coverage_ratio, config)
        self.config = config </s> remove     def __init__(self, input_size, output_size, coverage_ratio):
 </s> add     def __init__(self, input_size, output_size, coverage_ratio, config): </s> remove         batchsize = min(len(self.images), self.model.training_opts.get("preview_images", 14))
 </s> add         preview_images = self.config.get("preview_images", 14)
        preview_images = min(max(preview_images, 2), 16)
        batchsize = min(len(self.images), preview_images) </s> remove                           self.args.batch_size)
 </s> add                           self.args.batch_size,
                          self.args.configfile) </s> remove     def __init__(self, side, images, model, use_mask, batch_size):
        logger.debug("Initializing %s: side: '%s', num_images: %s, batch_size: %s)",
                     self.__class__.__name__, side, len(images), batch_size)
 </s> add     def __init__(self, side, images, model, use_mask, batch_size, config):
        logger.debug("Initializing %s: side: '%s', num_images: %s, batch_size: %s, config: %s)",
                     self.__class__.__name__, side, len(images), batch_size, config)
        generator = TrainingDataGenerator(input_size,
                                          output_size,
                                          self.model.training_opts,
                                          self.config)
          logger.debug("Loading generator: %s", self.side)
          input_size = self.model.input_shape[0]
          output_size = self.model.output_shape[0]
          logger.debug("input_size: %s, output_size: %s", input_size, output_size)
        generator = TrainingDataGenerator(input_size,
                                          output_size,
                                          self.model.training_opts,
                                          self.config)
          return generator
  
      def train_one_batch(self, do_preview):
          """ Train a batch """
          logger.trace("Training one step: (side: %s)", self.side)
keep keep keep keep replace keep keep keep keep keep
 Expose Augmentation Options to config 
plugins/train/trainer/_base.py