Spaces:
Build error
Build error
| 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) | |
| 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}") | |