From c7067cef8fbaa84bf0953de380f7fa4d8a5b7c59 Mon Sep 17 00:00:00 2001 From: Chenxi Date: Mon, 13 Sep 2021 19:25:42 +0100 Subject: [PATCH 01/10] enable cog --- cog.yaml | 29 +++++++ download-weights | 28 +++++++ predict.py | 195 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 252 insertions(+) create mode 100755 cog.yaml create mode 100755 download-weights create mode 100755 predict.py diff --git a/cog.yaml b/cog.yaml new file mode 100755 index 00000000..b5b4bf06 --- /dev/null +++ b/cog.yaml @@ -0,0 +1,29 @@ +build: + gpu: true + python_version: "3.8" + system_packages: + - "libgl1-mesa-glx" + - "libglib2.0-0" + python_packages: + - "cmake==3.21.2" + - "torchvision==0.9.0" + - "torch==1.8.0" + - "numpy==1.19.4" + - "opencv-python==4.4.0.46" + - "scipy==1.5.3" + - "tensorboardX==2.4" + - "dominate==2.6.0" + - "easydict==1.9" + - "PyYAML==5.3.1" + - "scikit-image==0.18.3" + - "dill==0.3.4" + - "einops==0.3.0" + - "PySimpleGUI==4.46.0" + - "ipython==7.19.0" + pre_install: + - pip install dlib + +predict: "predict.py:Predictor" + + + diff --git a/download-weights b/download-weights new file mode 100755 index 00000000..b8802258 --- /dev/null +++ b/download-weights @@ -0,0 +1,28 @@ +#!/bin/sh + +cd Face_Enhancement/models/networks +git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm . +cd ../../../ + +cd Global/detection_models +git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm . +cd ../../ + +# download the landmark detection model +cd Face_Detection/ +wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 +bzip2 -d shape_predictor_68_face_landmarks.dat.bz2 +cd ../ + +# download the pretrained model +cd Face_Enhancement/ +wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Face_Enhancement/checkpoints.zip +unzip checkpoints.zip +cd ../ + +cd Global/ +wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Global/checkpoints.zip +unzip checkpoints.zip +cd ../ \ No newline at end of file diff --git a/predict.py b/predict.py new file mode 100755 index 00000000..ce28aee7 --- /dev/null +++ b/predict.py @@ -0,0 +1,195 @@ +import cog +import tempfile +from pathlib import Path +import argparse +import sys +import shutil +import os +import cv2 +import glob +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 high-resolution image") + @cog.input("with_scratch", type=bool, default=False, help="whether input image is scratched") + def predict(self, image, HR=False, with_scratch=False): + 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.") + + img_name = os.path.basename(str(image)) + image_restore = cv2.imread(os.path.join(self.opts.output_folder, 'final_output', img_name)) + + out_path = Path(tempfile.mkdtemp()) / "out.png" + + cv2.imwrite(str(out_path), image_restore) + 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('Failed to delete %s. Reason: %s' % (file_path, e)) From 6f087202ef57aab7f2c9ad90a3d23d55e85e7154 Mon Sep 17 00:00:00 2001 From: Chenxi Date: Mon, 13 Sep 2021 20:54:15 +0100 Subject: [PATCH 02/10] fix multiple run bug --- predict.py | 1 + 1 file changed, 1 insertion(+) diff --git a/predict.py b/predict.py index ce28aee7..bf3a619d 100755 --- a/predict.py +++ b/predict.py @@ -36,6 +36,7 @@ def setup(self): @cog.input("HR", type=bool, default=False, help="whether high-resolution image") @cog.input("with_scratch", type=bool, default=False, help="whether 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)) shutil.copy(str(image), input_path) From f2ec44f1f828dfd5d4f66512382b91f793a027b1 Mon Sep 17 00:00:00 2001 From: Chenxi Date: Tue, 14 Sep 2021 22:51:43 +0100 Subject: [PATCH 03/10] update cog.yaml --- cog.yaml | 2 +- predict.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cog.yaml b/cog.yaml index b5b4bf06..99607baa 100755 --- a/cog.yaml +++ b/cog.yaml @@ -20,7 +20,7 @@ build: - "einops==0.3.0" - "PySimpleGUI==4.46.0" - "ipython==7.19.0" - pre_install: + run: - pip install dlib predict: "predict.py:Predictor" diff --git a/predict.py b/predict.py index bf3a619d..913f148b 100755 --- a/predict.py +++ b/predict.py @@ -33,8 +33,8 @@ def setup(self): @cog.input("image", type=Path, help="input image") - @cog.input("HR", type=bool, default=False, help="whether high-resolution image") - @cog.input("with_scratch", type=bool, default=False, help="whether 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)) From 12c689e56aa3f9705d9e15e6115ee530a045f3a4 Mon Sep 17 00:00:00 2001 From: Chenxi Date: Tue, 14 Sep 2021 23:20:31 +0100 Subject: [PATCH 04/10] update readme --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 95f320cd..dc10b569 100755 --- a/README.md +++ b/README.md @@ -4,6 +4,8 @@ ### [Project Page](http://raywzy.com/Old_Photo/) | [Paper (CVPR version)](https://arxiv.org/abs/2004.09484) | [Paper (Journal version)](https://arxiv.org/pdf/2009.07047v1.pdf) | [Pretrained Model](https://hkustconnect-my.sharepoint.com/:f:/g/personal/bzhangai_connect_ust_hk/Em0KnYOeSSxFtp4g_dhWdf0BdeT3tY12jIYJ6qvSf300cA?e=nXkJH2) | [Colab Demo](https://colab.research.google.com/drive/1NEm6AsybIiC5TwTU_4DqDkQO0nFRB-uA?usp=sharing) :fire: +[Run it on Replicate](https://replicate.ai/zhangmozhe/bringing-old-photos-back-to-life) + **Bringing Old Photos Back to Life, CVPR2020 (Oral)** **Old Photo Restoration via Deep Latent Space Translation, PAMI Under Review** From eb16399ee7e61761f1a5d60fd7a50d5149ba5ee8 Mon Sep 17 00:00:00 2001 From: Chenxi Date: Fri, 17 Sep 2021 10:54:45 +0100 Subject: [PATCH 05/10] reformatted --- cog.yaml | 3 -- download-weights | 2 +- predict.py | 130 ++++++++++++++++++++++++++++------------------- 3 files changed, 79 insertions(+), 56 deletions(-) diff --git a/cog.yaml b/cog.yaml index 99607baa..19cc6749 100755 --- a/cog.yaml +++ b/cog.yaml @@ -24,6 +24,3 @@ build: - pip install dlib predict: "predict.py:Predictor" - - - diff --git a/download-weights b/download-weights index b8802258..48178101 100755 --- a/download-weights +++ b/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 ../ \ No newline at end of file +cd ../ diff --git a/predict.py b/predict.py index 913f148b..34d3c5c9 100755 --- a/predict.py +++ b/predict.py @@ -1,4 +1,3 @@ -import cog import tempfile from pathlib import Path import argparse @@ -7,13 +6,16 @@ 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 @@ def setup(self): ) 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 @@ def predict(self, image, HR=False, with_scratch=False): 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 @@ def predict(self, image, HR=False, with_scratch=False): 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 @@ def predict(self, image, HR=False, with_scratch=False): 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 @@ def predict(self, image, HR=False, with_scratch=False): 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 @@ def predict(self, image, HR=False, with_scratch=False): 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 @@ def predict(self, image, HR=False, with_scratch=False): 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 @@ def predict(self, image, HR=False, with_scratch=False): 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)) From f959c0bb8b40cd0e524244d5d6c89f8152a3cf45 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Fri, 17 Sep 2021 14:36:10 +0100 Subject: [PATCH 06/10] Read output png images correctly --- predict.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/predict.py b/predict.py index 34d3c5c9..a813c2ba 100755 --- a/predict.py +++ b/predict.py @@ -197,10 +197,16 @@ def predict(self, image, HR=False, with_scratch=False): 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) - ) + ## This model only outputs png files so make sure to use .png when + ## reading the output file. We do that by replacing the extension of the input image + ## but keeping it's filename the same + img_name = Path(os.path.basename(str(image))).with_suffix('.png') + + output_img_path = os.path.join(self.opts.output_folder, "final_output", img_name) + + print(f'Reading the final output image at path {output_img_path}'); + + image_restore = cv2.imread(output_img_path) out_path = Path(tempfile.mkdtemp()) / "out.png" From 67ef0ba4a3347dbc71305d64b8a53ccfca7e082c Mon Sep 17 00:00:00 2001 From: CJWBW <70536672+CJWBW@users.noreply.github.com> Date: Sun, 19 Sep 2021 12:32:49 +0100 Subject: [PATCH 07/10] Update README.md --- README.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/README.md b/README.md index dc10b569..15215b97 100755 --- a/README.md +++ b/README.md @@ -2,9 +2,7 @@ -### [Project Page](http://raywzy.com/Old_Photo/) | [Paper (CVPR version)](https://arxiv.org/abs/2004.09484) | [Paper (Journal version)](https://arxiv.org/pdf/2009.07047v1.pdf) | [Pretrained Model](https://hkustconnect-my.sharepoint.com/:f:/g/personal/bzhangai_connect_ust_hk/Em0KnYOeSSxFtp4g_dhWdf0BdeT3tY12jIYJ6qvSf300cA?e=nXkJH2) | [Colab Demo](https://colab.research.google.com/drive/1NEm6AsybIiC5TwTU_4DqDkQO0nFRB-uA?usp=sharing) :fire: - -[Run it on Replicate](https://replicate.ai/zhangmozhe/bringing-old-photos-back-to-life) +### [Project Page](http://raywzy.com/Old_Photo/) | [Paper (CVPR version)](https://arxiv.org/abs/2004.09484) | [Paper (Journal version)](https://arxiv.org/pdf/2009.07047v1.pdf) | [Pretrained Model](https://hkustconnect-my.sharepoint.com/:f:/g/personal/bzhangai_connect_ust_hk/Em0KnYOeSSxFtp4g_dhWdf0BdeT3tY12jIYJ6qvSf300cA?e=nXkJH2) | [Colab Demo](https://colab.research.google.com/drive/1NEm6AsybIiC5TwTU_4DqDkQO0nFRB-uA?usp=sharing) | [Replicate Demo & Docker Image](https://replicate.ai/zhangmozhe/bringing-old-photos-back-to-life) :fire: **Bringing Old Photos Back to Life, CVPR2020 (Oral)** From b5b428393edba30a92c7959c1fe18f2769dcfc32 Mon Sep 17 00:00:00 2001 From: CJWBW <70536672+CJWBW@users.noreply.github.com> Date: Mon, 20 Sep 2021 16:44:14 +0100 Subject: [PATCH 08/10] Revert "Read output png images correctly" --- predict.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/predict.py b/predict.py index a813c2ba..34d3c5c9 100755 --- a/predict.py +++ b/predict.py @@ -197,16 +197,10 @@ def predict(self, image, HR=False, with_scratch=False): print("All the processing is done. Please check the results.") - ## This model only outputs png files so make sure to use .png when - ## reading the output file. We do that by replacing the extension of the input image - ## but keeping it's filename the same - img_name = Path(os.path.basename(str(image))).with_suffix('.png') - - output_img_path = os.path.join(self.opts.output_folder, "final_output", img_name) - - print(f'Reading the final output image at path {output_img_path}'); - - image_restore = cv2.imread(output_img_path) + img_name = os.path.basename(str(image)) + image_restore = cv2.imread( + os.path.join(self.opts.output_folder, "final_output", img_name) + ) out_path = Path(tempfile.mkdtemp()) / "out.png" From 05d9ed44452d92e6b63bfcfe08a839e8495be754 Mon Sep 17 00:00:00 2001 From: CJWBW <70536672+CJWBW@users.noreply.github.com> Date: Mon, 20 Sep 2021 16:52:22 +0100 Subject: [PATCH 09/10] fix read output image bug --- predict.py | 93 +++++++++++++++++++++++++++--------------------------- 1 file changed, 46 insertions(+), 47 deletions(-) diff --git a/predict.py b/predict.py index 34d3c5c9..753e451e 100755 --- a/predict.py +++ b/predict.py @@ -69,12 +69,12 @@ def predict(self, image, HR=False, with_scratch=False): 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: @@ -83,13 +83,13 @@ def predict(self, image, HR=False, with_scratch=False): 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: @@ -98,15 +98,15 @@ def predict(self, image, HR=False, with_scratch=False): 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) @@ -134,10 +134,10 @@ def predict(self, image, HR=False, with_scratch=False): 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) @@ -157,17 +157,17 @@ def predict(self, image, HR=False, with_scratch=False): 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) @@ -183,12 +183,12 @@ def predict(self, image, HR=False, with_scratch=False): 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) @@ -197,10 +197,9 @@ def predict(self, image, HR=False, with_scratch=False): 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) - ) + 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" From e9f55eca09b3a97e399ddce636709206aac61005 Mon Sep 17 00:00:00 2001 From: CJWBW <70536672+CJWBW@users.noreply.github.com> Date: Tue, 28 Sep 2021 15:00:44 +0100 Subject: [PATCH 10/10] reformat and fix temp input dir bug --- predict.py | 287 +++++++++++++++++++++++++++-------------------------- 1 file changed, 144 insertions(+), 143 deletions(-) diff --git a/predict.py b/predict.py index 753e451e..5573cd1a 100755 --- a/predict.py +++ b/predict.py @@ -1,11 +1,10 @@ import tempfile from pathlib import Path import argparse -import sys import shutil import os -import cv2 import glob +import cv2 import cog from run import run_cmd @@ -50,162 +49,164 @@ def setup(self): 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)) - 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 + 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" ) - run_cmd(stage_1_command_1) - run_cmd(stage_1_command_2) + os.makedirs(stage_1_output_dir, exist_ok=True) - ## 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) + if not with_scratch: - print("Finish Stage 1 ...") - print("\n") + 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: - ## Stage 2: Face Detection + 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) - 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 + ) - 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" + ) - 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) - 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" - ) + 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_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") + run_cmd(stage_4_command) + print("Finish Stage 4 ...") + print("\n") - print("All the processing is done. Please check the results.") + print("All the processing is done. Please check the results.") - final_output = os.listdir(os.path.join(self.opts.output_folder, "final_output"))[0] + 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)) + image_restore = cv2.imread(os.path.join(self.opts.output_folder, "final_output", final_output)) - out_path = Path(tempfile.mkdtemp()) / "out.png" + out_path = Path(tempfile.mkdtemp()) / "out.png" - cv2.imwrite(str(out_path), image_restore) - clean_folder(self.opts.input_folder) - clean_folder(self.opts.output_folder) + cv2.imwrite(str(out_path), image_restore) + finally: + clean_folder(self.opts.input_folder) + clean_folder(self.opts.output_folder) return out_path @@ -218,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(f"Failed to delete {file_path}. Reason:{e}")