Change output path for face vectors

align_images.py

@@ -26,12 +26,14 @@ def unpack_bz2(src_path):
                                                LANDMARKS_MODEL_URL, cache_subdir='temp'))
     RAW_IMAGES_DIR = sys.argv[1]
     ALIGNED_IMAGES_DIR = sys.argv[2]
+    VECTOR_DIR = sys.argv[3]
 
     landmarks_detector = LandmarksDetector(landmarks_model_path)
     for img_name in [x for x in os.listdir(RAW_IMAGES_DIR) if x[0] not in '._']:
         raw_img_path = os.path.join(RAW_IMAGES_DIR, img_name)
         for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):
             face_img_name = '%s_%02d' % (os.path.splitext(img_name)[0], i)
             aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)
+            vector_path = os.path.join(VECTOR_DIR, face_img_name)
             os.makedirs(ALIGNED_IMAGES_DIR, exist_ok=True)
-            image_align(raw_img_path, aligned_face_path, face_landmarks)
+            image_align(raw_img_path, aligned_face_path, vector_path, face_landmarks)

ffhq_dataset/face_alignment.py

@@ -4,7 +4,7 @@
 import PIL.Image
 
 
-def image_align(src_file, dst_file, face_landmarks, output_size=1024, transform_size=4096, enable_padding=True, x_scale=1, y_scale=1, em_scale=0.1, alpha=False):
+def image_align(src_file, dst_file, vector_file, face_landmarks, output_size=1024, transform_size=4096, enable_padding=True, x_scale=1, y_scale=1, em_scale=0.1, alpha=False):
         # Align function from FFHQ dataset pre-processing step
         # https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py
 
@@ -92,5 +92,4 @@ def image_align(src_file, dst_file, face_landmarks, output_size=1024, transform_
         img.save(dst_file + '.png', 'PNG')
 
         #Save face placement to later fit face back into image
-        with open(dst_file + '.npy', 'wb') as f:
-            np.save(f, np.array(face_landmarks))
+        np.save(vector_file + '.npy', np.array(face_landmarks))