reformatted

cog.yaml

@@ -24,6 +24,3 @@ build:
     - pip install dlib
 
 predict: "predict.py:Predictor"
-
-
-

download-weights

@@ -25,4 +25,4 @@ cd ../
 cd Global/
 wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Global/checkpoints.zip
 unzip checkpoints.zip
-cd ../
+cd ../

predict.py

@@ -1,4 +1,3 @@
-import cog
 import tempfile
 from pathlib import Path
 import argparse
@@ -7,13 +6,16 @@ import shutil
 import os
 import cv2
 import glob
+import cog
 from run import run_cmd
 
 
 class Predictor(cog.Predictor):
     def setup(self):
         parser = argparse.ArgumentParser()
-        parser.add_argument("--input_folder", type=str, default='input/cog_temp', help="Test images")
+        parser.add_argument(
+            "--input_folder", type=str, default="input/cog_temp", help="Test images"
+        )
         parser.add_argument(
             "--output_folder",
             type=str,
@@ -22,19 +24,31 @@ class Predictor(cog.Predictor):
         )
         parser.add_argument("--GPU", type=str, default="0", help="0,1,2")
         parser.add_argument(
-            "--checkpoint_name", type=str, default="Setting_9_epoch_100", help="choose which checkpoint"
+            "--checkpoint_name",
+            type=str,
+            default="Setting_9_epoch_100",
+            help="choose which checkpoint",
         )
-        self.opts = parser.parse_args('')
+        self.opts = parser.parse_args("")
         self.basepath = os.getcwd()
         self.opts.input_folder = os.path.join(self.basepath, self.opts.input_folder)
         self.opts.output_folder = os.path.join(self.basepath, self.opts.output_folder)
         os.makedirs(self.opts.input_folder, exist_ok=True)
         os.makedirs(self.opts.output_folder, exist_ok=True)
 
-
     @cog.input("image", type=Path, help="input image")
-    @cog.input("HR", type=bool, default=False, help="whether the input image is high-resolution")
-    @cog.input("with_scratch", type=bool, default=False, help="whether the input image is scratched")
+    @cog.input(
+        "HR",
+        type=bool,
+        default=False,
+        help="whether the input image is high-resolution",
+    )
+    @cog.input(
+        "with_scratch",
+        type=bool,
+        default=False,
+        help="whether the input image is scratched",
+    )
     def predict(self, image, HR=False, with_scratch=False):
         os.chdir(self.basepath)
         input_path = os.path.join(self.opts.input_folder, os.path.basename(image))
@@ -46,19 +60,21 @@ class Predictor(cog.Predictor):
         print("Running Stage 1: Overall restoration")
         os.chdir("./Global")
         stage_1_input_dir = self.opts.input_folder
-        stage_1_output_dir = os.path.join(self.opts.output_folder, "stage_1_restore_output")
+        stage_1_output_dir = os.path.join(
+            self.opts.output_folder, "stage_1_restore_output"
+        )
 
         os.makedirs(stage_1_output_dir, exist_ok=True)
 
         if not with_scratch:
 
             stage_1_command = (
-                    "python test.py --test_mode Full --Quality_restore --test_input "
-                    + stage_1_input_dir
-                    + " --outputs_dir "
-                    + stage_1_output_dir
-                    + " --gpu_ids "
-                    + gpu1
+                "python test.py --test_mode Full --Quality_restore --test_input "
+                + stage_1_input_dir
+                + " --outputs_dir "
+                + stage_1_output_dir
+                + " --gpu_ids "
+                + gpu1
             )
             run_cmd(stage_1_command)
         else:
@@ -67,13 +83,13 @@ class Predictor(cog.Predictor):
             new_input = os.path.join(mask_dir, "input")
             new_mask = os.path.join(mask_dir, "mask")
             stage_1_command_1 = (
-                    "python detection.py --test_path "
-                    + stage_1_input_dir
-                    + " --output_dir "
-                    + mask_dir
-                    + " --input_size full_size"
-                    + " --GPU "
-                    + gpu1
+                "python detection.py --test_path "
+                + stage_1_input_dir
+                + " --output_dir "
+                + mask_dir
+                + " --input_size full_size"
+                + " --GPU "
+                + gpu1
             )
 
             if HR:
@@ -82,14 +98,15 @@ class Predictor(cog.Predictor):
                 HR_suffix = ""
 
             stage_1_command_2 = (
-                    "python test.py --Scratch_and_Quality_restore --test_input "
-                    + new_input
-                    + " --test_mask "
-                    + new_mask
-                    + " --outputs_dir "
-                    + stage_1_output_dir
-                    + " --gpu_ids "
-                    + gpu1 + HR_suffix
+                "python test.py --Scratch_and_Quality_restore --test_input "
+                + new_input
+                + " --test_mask "
+                + new_mask
+                + " --outputs_dir "
+                + stage_1_output_dir
+                + " --gpu_ids "
+                + gpu1
+                + HR_suffix
             )
 
             run_cmd(stage_1_command_1)
@@ -111,11 +128,16 @@ class Predictor(cog.Predictor):
         print("Running Stage 2: Face Detection")
         os.chdir(".././Face_Detection")
         stage_2_input_dir = os.path.join(stage_1_output_dir, "restored_image")
-        stage_2_output_dir = os.path.join(self.opts.output_folder, "stage_2_detection_output")
+        stage_2_output_dir = os.path.join(
+            self.opts.output_folder, "stage_2_detection_output"
+        )
         os.makedirs(stage_2_output_dir, exist_ok=True)
 
         stage_2_command = (
-                "python detect_all_dlib_HR.py --url " + stage_2_input_dir + " --save_url " + stage_2_output_dir
+            "python detect_all_dlib_HR.py --url "
+            + stage_2_input_dir
+            + " --save_url "
+            + stage_2_output_dir
         )
 
         run_cmd(stage_2_command)
@@ -127,23 +149,25 @@ class Predictor(cog.Predictor):
         os.chdir(".././Face_Enhancement")
         stage_3_input_mask = "./"
         stage_3_input_face = stage_2_output_dir
-        stage_3_output_dir = os.path.join(self.opts.output_folder, "stage_3_face_output")
+        stage_3_output_dir = os.path.join(
+            self.opts.output_folder, "stage_3_face_output"
+        )
 
         os.makedirs(stage_3_output_dir, exist_ok=True)
 
-        self.opts.checkpoint_name = 'FaceSR_512'
+        self.opts.checkpoint_name = "FaceSR_512"
         stage_3_command = (
-                "python test_face.py --old_face_folder "
-                + stage_3_input_face
-                + " --old_face_label_folder "
-                + stage_3_input_mask
-                + " --tensorboard_log --name "
-                + self.opts.checkpoint_name
-                + " --gpu_ids "
-                + gpu1
-                + " --load_size 512 --label_nc 18 --no_instance --preprocess_mode resize --batchSize 1 --results_dir "
-                + stage_3_output_dir
-                + " --no_parsing_map"
+            "python test_face.py --old_face_folder "
+            + stage_3_input_face
+            + " --old_face_label_folder "
+            + stage_3_input_mask
+            + " --tensorboard_log --name "
+            + self.opts.checkpoint_name
+            + " --gpu_ids "
+            + gpu1
+            + " --load_size 512 --label_nc 18 --no_instance --preprocess_mode resize --batchSize 1 --results_dir "
+            + stage_3_output_dir
+            + " --no_parsing_map"
         )
 
         run_cmd(stage_3_command)
@@ -159,12 +183,12 @@ class Predictor(cog.Predictor):
         os.makedirs(stage_4_output_dir, exist_ok=True)
 
         stage_4_command = (
-                "python align_warp_back_multiple_dlib_HR.py --origin_url "
-                + stage_4_input_image_dir
-                + " --replace_url "
-                + stage_4_input_face_dir
-                + " --save_url "
-                + stage_4_output_dir
+            "python align_warp_back_multiple_dlib_HR.py --origin_url "
+            + stage_4_input_image_dir
+            + " --replace_url "
+            + stage_4_input_face_dir
+            + " --save_url "
+            + stage_4_output_dir
         )
 
         run_cmd(stage_4_command)
@@ -174,7 +198,9 @@ class Predictor(cog.Predictor):
         print("All the processing is done. Please check the results.")
 
         img_name = os.path.basename(str(image))
-        image_restore = cv2.imread(os.path.join(self.opts.output_folder, 'final_output', img_name))
+        image_restore = cv2.imread(
+            os.path.join(self.opts.output_folder, "final_output", img_name)
+        )
 
         out_path = Path(tempfile.mkdtemp()) / "out.png"
 
@@ -193,4 +219,4 @@ def clean_folder(folder):
             elif os.path.isdir(file_path):
                 shutil.rmtree(file_path)
         except Exception as e:
-            print('Failed to delete %s. Reason: %s' % (file_path, e))
+            print("Failed to delete %s. Reason: %s" % (file_path, e))