predict.py

import tempfile
from pathlib import Path
import argparse
import shutil
import os
import glob
import cv2
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(
            "--output_folder",
            type=str,
            default="output",
            help="Restored images, please use the absolute path",
        )
        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",
        )
        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",
    )
    def predict(self, image, HR=False, with_scratch=False):
        try:
            os.chdir(self.basepath)
            input_path = os.path.join(self.opts.input_folder, os.path.basename(image))
            shutil.copy(str(image), input_path)

            gpu1 = self.opts.GPU

            ## Stage 1: Overall Quality Improve
            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"
            )

            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
                )
                run_cmd(stage_1_command)
            else:

                mask_dir = os.path.join(stage_1_output_dir, "masks")
                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
                )

                if HR:
                    HR_suffix = " --HR"
                else:
                    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
                )

                run_cmd(stage_1_command_1)
                run_cmd(stage_1_command_2)

            ## Solve the case when there is no face in the old photo
            stage_1_results = os.path.join(stage_1_output_dir, "restored_image")
            stage_4_output_dir = os.path.join(self.opts.output_folder, "final_output")
            os.makedirs(stage_4_output_dir, exist_ok=True)
            for x in os.listdir(stage_1_results):
                img_dir = os.path.join(stage_1_results, x)
                shutil.copy(img_dir, stage_4_output_dir)

            print("Finish Stage 1 ...")
            print("\n")

            ## Stage 2: Face Detection

            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"
            )
            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
            )

            run_cmd(stage_2_command)
            print("Finish Stage 2 ...")
            print("\n")

            ## Stage 3: Face Restore
            print("Running Stage 3: Face Enhancement")
            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"
            )

            os.makedirs(stage_3_output_dir, exist_ok=True)

            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"
            )

            run_cmd(stage_3_command)
            print("Finish Stage 3 ...")
            print("\n")

            ## Stage 4: Warp back
            print("Running Stage 4: Blending")
            os.chdir(".././Face_Detection")
            stage_4_input_image_dir = os.path.join(stage_1_output_dir, "restored_image")
            stage_4_input_face_dir = os.path.join(stage_3_output_dir, "each_img")
            stage_4_output_dir = os.path.join(self.opts.output_folder, "final_output")
            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
            )

            run_cmd(stage_4_command)
            print("Finish Stage 4 ...")
            print("\n")

            print("All the processing is done. Please check the results.")

            final_output = os.listdir(os.path.join(self.opts.output_folder, "final_output"))[0]

            image_restore = cv2.imread(os.path.join(self.opts.output_folder, "final_output", final_output))

            out_path = Path(tempfile.mkdtemp()) / "out.png"

            cv2.imwrite(str(out_path), image_restore)
        finally:
            clean_folder(self.opts.input_folder)
            clean_folder(self.opts.output_folder)
        return out_path


def clean_folder(folder):
    for filename in os.listdir(folder):
        file_path = os.path.join(folder, filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print(f"Failed to delete {file_path}. Reason:{e}")