Add optional model path to convert.py (#14)

convert.py

@@ -1,5 +1,5 @@
 import cv2
-from lib.cli import DirectoryProcessor
+from lib.cli import DirectoryProcessor, FullPaths
 from pathlib import Path
 from lib.faces_process import convert_one_image
 from lib.faces_detect import crop_faces
@@ -8,6 +8,15 @@
 class ConvertImage(DirectoryProcessor):
     filename = ''
 
+    def add_optional_arguments(self, parser):
+        parser.add_argument('-m', '--model-dir',
+                            action=FullPaths,
+                            dest="model_dir",
+                            default="models",
+                            help="Model directory. A directory containing the trained model \
+                    you wish to process. Defaults to 'models'")
+        return parser
+
     def process_image(self, filename):
         try:
             image = cv2.imread(filename)
@@ -16,7 +25,7 @@ def process_image(self, filename):
                     print('- Found more than one face!')
                     self.verify_output = True
 
-                new_face = convert_one_image(cv2.resize(face.image, (256, 256)))
+                new_face = convert_one_image(cv2.resize(face.image, (256, 256)), self.arguments.model_dir)
                 image[slice(face.y, face.y + face.h), slice(face.x, face.x + face.w)] = cv2.resize(new_face, (face.w, face.h))
                 self.faces_detected = self.faces_detected + 1
             output_file = self.output_dir / Path(filename).name

lib/cli.py

@@ -12,7 +12,7 @@ def __call__(self, parser, namespace, values, option_string=None):
         setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
 
 
-class DirectoryProcessor():
+class DirectoryProcessor(object):
     '''
     Abstract class that processes a directory of images
     and writes output to the specified folder

lib/faces_process.py

@@ -6,18 +6,19 @@
 from .model import autoencoder_B
 from .model import encoder, decoder_A, decoder_B
 
-encoder.load_weights("models/encoder.h5")
-decoder_A.load_weights("models/decoder_A.h5")
-decoder_B.load_weights("models/decoder_B.h5")
 
-autoencoder = autoencoder_B
+def convert_one_image(image, model_dir="models"):
 
-# landmark file can be found in http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
-# unzip it in the same folder as the main scripts
-aligner = Aligner("shape_predictor_68_face_landmarks.dat", "mmod_human_face_detector.dat")
+    encoder.load_weights(model_dir + "/encoder.h5")
+    decoder_A.load_weights(model_dir + "/decoder_A.h5")
+    decoder_B.load_weights(model_dir + "/decoder_B.h5")
 
+    autoencoder = autoencoder_B
+
+    # landmark file can be found in http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
+    # unzip it in the same folder as the main scripts
+    aligner = Aligner("shape_predictor_68_face_landmarks.dat", "mmod_human_face_detector.dat")
 
-def convert_one_image(image):
     assert image.shape == (256, 256, 3)
     crop = slice(48, 208)
     face = image[crop, crop]