From e5402f09279b531e341b5a0ca64d273567e479e9 Mon Sep 17 00:00:00 2001
From: Maxim Kan <maxim.kan@phygitalism.com>
Date: Mon, 23 Aug 2021 17:23:52 +0000
Subject: [PATCH] solve extension bug

---
 run.py              | 6 +++---
 utils/align_data.py | 4 ++--
 utils/alignment.py  | 3 +--
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/run.py b/run.py
index c648f8d..c02d3e9 100644
--- a/run.py
+++ b/run.py
@@ -27,6 +27,7 @@ def gen_vec(image_name, latent_editor, alpha, step):
     return latents_vec
     
 def gen_img(image_name, model_id, latents_vec, base_save_path):
+    image_name, ext = image_name.split('.')
     generator_type = paths_config.multi_id_model_type if hyperparameters.use_multi_id_training else image_name
     new_G = load_generators(model_id, generator_type)
     for direction, factor_and_edit in latents_vec.items():
@@ -36,7 +37,7 @@ def gen_img(image_name, model_id, latents_vec, base_save_path):
             img = Image.fromarray(img, mode='RGB')
             path = os.path.join(base_save_path, image_name, direction)
             os.makedirs(path, exist_ok=True)
-            img.save(os.path.join(path, str(val) + "_" + image_name + '.jpg'))
+            img.save(os.path.join(path, str(val) + "_" + image_name + '.' + ext))
             
 def evaluate(args):
     os.makedirs(paths_config.input_data_path, exist_ok=True)
@@ -48,8 +49,7 @@ def evaluate(args):
     os.makedirs(base_save_path, exist_ok=True)
     with torch.no_grad():
         for image_name in tqdm(name_list):
-            image_name = image_name.split('.')[0]
-            latents_vec = gen_vec(image_name, latent_editor, alpha=args.alpha, step=args.step)
+            latents_vec = gen_vec(image_name.split('.')[0], latent_editor, alpha=args.alpha, step=args.step)
             gen_img(image_name, model_id, latents_vec, base_save_path)
             print(f'Done for {image_name}') 
             
diff --git a/utils/align_data.py b/utils/align_data.py
index 74e70a2..c7feb7a 100644
--- a/utils/align_data.py
+++ b/utils/align_data.py
@@ -23,8 +23,8 @@ def pre_process_images(raw_images_path):
 
 #     os.makedirs('./image_original/image_processed', exist_ok=True)
     for image, name in zip(aligned_images, images_names):
-        real_name = os.path.basename(name).split('.')[0]
-        image.save(os.path.join(paths_config.input_data_path, real_name+'.jpeg'))
+        real_name, ext = os.path.basename(name).split('.')
+        image.save(os.path.join(paths_config.input_data_path, real_name+'.'+ext))
 
 #     os.chdir(current_directory)
 
diff --git a/utils/alignment.py b/utils/alignment.py
index 3d4a697..316511e 100644
--- a/utils/alignment.py
+++ b/utils/alignment.py
@@ -64,8 +64,7 @@ def align_face(filepath, predictor, output_size):
     qsize = np.hypot(*x) * 2
 
     # read image
-    img = PIL.Image.open(filepath)
-
+    img = PIL.Image.open(filepath).convert('RGB')
     transform_size = output_size
     enable_padding = True