      def add_optional_arguments(self, parser):
          # Override this for custom arguments
          return parser
  
      def save_model_weights(self):
          encoder.save_weights(self.arguments.model_dir + '/encoder.h5')
          decoder_A.save_weights(self.arguments.model_dir + '/decoder_A.h5')
          decoder_B.save_weights(self.arguments.model_dir + '/decoder_B.h5')
          print('save model weights')
  
      def show_sample(self, test_A, test_B):
          figure_A = numpy.stack([
              test_A,
              autoencoder_A.predict(test_A),
              autoencoder_B.predict(test_A),
          ], axis=1)
          figure_B = numpy.stack([
              test_B,
              autoencoder_B.predict(test_B),
              autoencoder_A.predict(test_B),
          ], axis=1)
  
          figure = numpy.concatenate([figure_A, figure_B], axis=0)
          figure = figure.reshape((4, 7) + figure.shape[1:])
          figure = stack_images(figure)
  
          figure = numpy.clip(figure * 255, 0, 255).astype('uint8')
  
          if self.arguments.preview is True:
              cv2.imshow('', figure)
          if not self.arguments.preview or self.arguments.write_image:
              cv2.imwrite('_sample.jpg', figure)
  
      def process(self):
          print('Starting, this may take a while...')
          images_A = get_image_paths(self.arguments.input_A)
          images_B = get_image_paths(self.arguments.input_B)
          images_A = load_images(images_A) / 255.0
 
</s> remove         print('Starting, this may take a while...')
 </s> add         import threading
        self.stop = False

        thr = threading.Thread(target=self.processThread, args=(), kwargs={})
        thr.start()

        if self.arguments.preview:
            print('Using live preview')
            while True:
                for name, image in self.preview_buffer.items():
                    cv2.imshow(name, image)

                key = cv2.waitKey(1000)
                if key == ord('\n') or key == ord('\r'):
                    break
        else:
            input() # TODO how to catch a specific key instead of Enter?

        print("Exit requested! The trainer will complete its current cycle, save the models and quit (it can take up a couple of seconds depending on your training speed). If you want to kill it now, press Ctrl + c")
        self.stop = True
        thr.join() # waits until thread finishes
    
    def processThread(self):
        variant = "Original" # TODO Pass as argument
        
        print('Loading data, this may take a while...')
        model = PluginLoader.get_model(variant)(self.arguments.model_dir)
        model.load(swapped=False)
 </s> remove         images_A = load_images(images_A) / 255.0
        images_B = load_images(images_B) / 255.0

        images_A += images_B.mean(axis=(0, 1, 2)) - images_A.mean(axis=(0, 1, 2))

        print('press "q" to stop training and save model')

        BATCH_SIZE = 64

        for epoch in range(1000000):
            if self.arguments.verbose:
                print("Iteration number {}".format(epoch + 1))
                start_time = time.time()
            warped_A, target_A = get_training_data(images_A, BATCH_SIZE)
            warped_B, target_B = get_training_data(images_B, BATCH_SIZE)

            loss_A = autoencoder_A.train_on_batch(warped_A, target_A)
            loss_B = autoencoder_B.train_on_batch(warped_B, target_B)
            print(loss_A, loss_B)

            if epoch % self.arguments.save_interval == 0:
                self.save_model_weights()
                self.show_sample(target_A[0:14], target_B[0:14])

            key = cv2.waitKey(1)
            if key == ord('q'):
                self.save_model_weights()
                exit()
            if self.arguments.verbose:
                end_time = time.time()
                time_elapsed = int(round((end_time - start_time)))
                m, s = divmod(time_elapsed, 60)
                h, m = divmod(m, 60)
                print("Iteration done in {:02d}h{:02d}m{:02d}s".format(h, m, s))
 </s> add         trainer = PluginLoader.get_trainer(variant)(model, images_A, images_B)

        try:
            print('Starting. Press "Enter" to stop training and save model')

            for epoch in range(0, 1000000):

                save_iteration = epoch % self.arguments.save_interval == 0

                trainer.train_one_step(epoch, self.show if save_iteration else None)

                if save_iteration:
                    model.save_weights()

                if self.stop:
                    model.save_weights()
                    exit()

        except KeyboardInterrupt:
            try:
                model.save_weights()
            except KeyboardInterrupt:
                print('Saving model weights has been cancelled!')
            exit(0)

    preview_buffer = {}

    def show(self, image, name=''):
        try:
            if self.arguments.preview:
                self.preview_buffer[name] = image
            elif self.arguments.write_image:
                cv2.imwrite('_sample_{}.jpg'.format(name), image)
        except Exception as e:
            print("could not preview sample")
            print(e) </s> remove     def process_image(self, filename):
        # TODO move the model load and the converter creation in a method called on init, but after the arg parsing
        (face_A,face_B) = ('/decoder_A.h5', '/decoder_B.h5') if not self.arguments.swap_model else ('/decoder_B.h5', '/decoder_A.h5')
 </s> add     def process(self):
        # Original model goes with Adjust or Masked converter
        model_name = "Original" # TODO Pass as argument
        conv_name = "Masked" # TODO Pass as argument </s> remove         new_face = self.autoencoder.predict( face / 255.0 )[0]
 </s> add         new_face = self.encoder( face / 255.0 )[0]
 </s> remove def load_images(image_paths, convert=None):
    iter_all_images = (cv2.imread(fn) for fn in image_paths)
    if convert:
        iter_all_images = (convert(img) for img in iter_all_images)
    for i, image in enumerate(iter_all_images):
        if i == 0:
            all_images = numpy.empty((len(image_paths), ) + image.shape, dtype=image.dtype)
        all_images[i] = image
    return all_images
 </s> add     def iterator(self):
        while True:
            next_item = self.queue.get()
            if next_item is None:
                break
            yield next_item
      def add_optional_arguments(self, parser):
          # Override this for custom arguments
          return parser
  





























      def process(self):
          print('Starting, this may take a while...')
          images_A = get_image_paths(self.arguments.input_A)
          images_B = get_image_paths(self.arguments.input_B)
          images_A = load_images(images_A) / 255.0
keep keep keep keep replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace keep keep keep keep keep
 Adding models as plugins + Face filtering (#53) + #39 + #43 + #44 + #49 (#61)

* Making Models as plugins

* Do not reload model on each image #39 + Adding FaceFilter #53

* Adding @lukaville PR for #43 and #44 (possibly)

* Training done in a separate thread

* Better log for plugin load

* Adding a prefetch to train.py #49
(Note that we prefetch 2 batches of images, due to the queue behavior)
+ More compact logging with verbose info included

* correction of DirectoryProcessor signature

* adding missing import

* Convert with parallel preprocessing of files

* Added coverage var for trainer

Added a var with comment. Feel free to add it as argument

* corrections

* Modifying preview and normalization of image + correction

* Cleanup 
scripts/train.py