Unified CLI (deepfakes#22)

* Created a single script to call the other ones.

Usage is ./faceswap.py {train|extract|convert}

* Improved the help from the commands.

* Added forgotten faceswap.py file.

* Changed gitignore to add the scripts.

* Updates gitignore.

* Added guarding not to execute code when imported.

* Removed useless script. Display help when no arguments are provided.

* Update README

.gitignore

@@ -4,4 +4,5 @@
 !*.txt
 !Dockerfile*
 !requirements*
-!lib/*.py
+!lib/*.py
+!scripts

README.md

@@ -47,13 +47,13 @@ The project has multiple entry points. You will have to:
  - **Convert** your sources with the model
 
 ### Extract
-From your setup folder, run `python extract.py`. This will take photos from `src` folder and extract faces into `extract` folder.
+From your setup folder, run `python faceswap.py extract`. This will take photos from `src` folder and extract faces into `extract` folder.
 
 ### Train
-From your setup folder, run `python train.py`. This will take photos from `data/trump` and `data/cage` folder and train a model that will be saved inside the `models` folder.
+From your setup folder, run `python faceswap.py train`. This will take photos from two folders containing pictures of both faces and train a model that will be saved inside the `models` folder.
 
 ### Convert
-From your setup folder, run `python convert.py`. This will take photos from `original` folder and apply new faces into `modified` folder.
+From your setup folder, run `python faceswap.py convert`. This will take photos from `original` folder and apply new faces into `modified` folder.
 
 #### General notes:
 - All of the scripts mentioned have `-h`/`--help` options with a arguments that they will accept. You're smart, you can figure out how this works, right?!

faceswap.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+import argparse
+from scripts.extract import ExtractTrainingData
+from lib.cli import TrainingProcessor
+from scripts.convert import ConvertImage
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    subparser = parser.add_subparsers()
+    extract = ExtractTrainingData(subparser, "extract", "Extract the faces from a pictures.")
+    train = TrainingProcessor(subparser, "train", "This command trains the model for the two faces A and B.")
+    convert = ConvertImage(subparser, "convert", "Convert a source image to a new one with the face swapped.")
+    arguments = parser.parse_args()
+    try:
+        arguments.func(arguments)
+    except:
+        parser.print_help()

lib/cli.py

@@ -22,10 +22,11 @@ def __call__(self, parser, namespace, values, option_string=None):
 class TrainingProcessor(object):
     arguments = None
 
-    def __init__(self, description='default'):
-        print('Initializing')
-        self.parse_arguments(description)
+    def __init__(self, subparser, command, description='default'):
+        self.parse_arguments(description, subparser, command)
 
+    def process_arguments(self, arguments):
+        self.arguments = arguments
         print("Model A Directory: {}".format(self.arguments.input_A))
         print("Model B Directory: {}".format(self.arguments.input_B))
         print("Training data directory: {}".format(self.arguments.model_dir))
@@ -41,8 +42,10 @@ def __init__(self, description='default'):
 
         self.process()
 
-    def parse_arguments(self, description):
-        parser = argparse.ArgumentParser(
+    def parse_arguments(self, description, subparser, command):
+        parser = subparser.add_parser(
+            command,
+            help="This command trains the model for the two faces A and B.",
             description=description,
             epilog="Questions and feedback: \
             https://github.com/deepfakes/faceswap-playground"
@@ -78,7 +81,7 @@ def parse_arguments(self, description):
                             default=False,
                             help="Show verbose output")
         parser = self.add_optional_arguments(parser)
-        self.arguments = parser.parse_args()
+        parser.set_defaults(func=self.process_arguments)
 
     def add_optional_arguments(self, parser):
         # Override this for custom arguments
@@ -150,6 +153,7 @@ class DirectoryProcessor(object):
     and writes output to the specified folder
     '''
     arguments = None
+    parser = None
 
     input_dir = None
     output_dir = None
@@ -159,10 +163,12 @@ class DirectoryProcessor(object):
     images_processed = 0
     faces_detected = 0
 
-    def __init__(self, description='default'):
-        print('Initializing')
-        self.parse_arguments(description)
+    def __init__(self, subparser, command, description='default'):
+        self.create_parser(subparser, command, description)
+        self.parse_arguments(description, subparser, command)
 
+    def process_arguments(self, arguments):
+        self.arguments = arguments
         print("Input Directory: {}".format(self.arguments.input_dir))
         print("Output Directory: {}".format(self.arguments.output_dir))
         print('Starting, this may take a while...')
@@ -176,6 +182,9 @@ def __init__(self, description='default'):
 
         self.images_found = len(self.input_dir)
 
+        self.process_directory()
+
+    def process_directory(self):
         for filename in self.input_dir:
             if self.arguments.verbose:
                 print('Processing: {}'.format(os.path.basename(filename)))
@@ -185,32 +194,35 @@ def __init__(self, description='default'):
 
         self.finalize()
 
-    def parse_arguments(self, description):
-        parser = argparse.ArgumentParser(
-            description=description,
-            epilog="Questions and feedback: \
-            https://github.com/deepfakes/faceswap-playground"
-        )
-
-        parser.add_argument('-i', '--input-dir',
+    def parse_arguments(self, description, subparser, command):
+        self.parser.add_argument('-i', '--input-dir',
                             action=FullPaths,
                             dest="input_dir",
                             default="input",
                             help="Input directory. A directory containing the files \
                             you wish to process. Defaults to 'input'")
-        parser.add_argument('-o', '--output-dir',
+        self.parser.add_argument('-o', '--output-dir',
                             action=FullPaths,
                             dest="output_dir",
                             default="output",
                             help="Output directory. This is where the converted files will \
                                 be stored. Defaults to 'output'")
-        parser.add_argument('-v', '--verbose',
+        self.parser.add_argument('-v', '--verbose',
                             action="store_true",
                             dest="verbose",
                             default=False,
                             help="Show verbose output")
-        parser = self.add_optional_arguments(parser)
-        self.arguments = parser.parse_args()
+        self.parser = self.add_optional_arguments(self.parser)
+        self.parser.set_defaults(func=self.process_arguments)
+
+    def create_parser(self, subparser, command, description):
+        parser = subparser.add_parser(
+            command,
+            description=description,
+            epilog="Questions and feedback: \
+            https://github.com/deepfakes/faceswap-playground"
+        )
+        return parser
 
     def add_optional_arguments(self, parser):
         # Override this for custom arguments

convert.py → scripts/convert.py

@@ -7,7 +7,15 @@
 
 class ConvertImage(DirectoryProcessor):
     filename = ''
-
+    def create_parser(self, subparser, command, description):
+        self.parser = subparser.add_parser(
+            command,
+            help="Convert a source image to a new one with the face swapped.",
+            description=description,
+            epilog="Questions and feedback: \
+            https://github.com/deepfakes/faceswap-playground"
+        )
+
     def add_optional_arguments(self, parser):
         parser.add_argument('-m', '--model-dir',
                             action=FullPaths,
@@ -34,4 +42,3 @@ def process_image(self, filename):
             print('Failed to extract from image: {}. Reason: {}'.format(filename, e))
 
 
-extract_cli = ConvertImage(description='Swaps faces for images in a directory')

extract.py → scripts/extract.py

@@ -4,12 +4,17 @@
 
 
 class ExtractTrainingData(DirectoryProcessor):
+    def create_parser(self, subparser, command, description):
+        self.parser = subparser.add_parser(
+            command,
+            help="Extract the faces from a pictures.",
+            description=description,
+            epilog="Questions and feedback: \
+            https://github.com/deepfakes/faceswap-playground"
+        )
+
     def process_face(self, face, index, filename):
         resized_image = cv2.resize(face.image, (256, 256))
         output_file = self.output_dir / Path(filename).stem
         cv2.imwrite(str(output_file) + str(index) + Path(filename).suffix,
                     resized_image)
-
-
-extract_cli = ExtractTrainingData(description='Extracts faces from a collection of pictures \
-    and saves them to a separate directory')