Skip to content

Commit

Permalink
reformatted
Browse files Browse the repository at this point in the history
  • Loading branch information
Chenxi authored and Chenxi committed Sep 17, 2021
1 parent 12c689e commit eb16399
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 56 deletions.
3 changes: 0 additions & 3 deletions cog.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,3 @@ build:
- pip install dlib

predict: "predict.py:Predictor"



2 changes: 1 addition & 1 deletion download-weights
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 ../
130 changes: 78 additions & 52 deletions predict.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import cog
import tempfile
from pathlib import Path
import argparse
Expand All @@ -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,
Expand All @@ -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))
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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"

Expand All @@ -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))

0 comments on commit eb16399

Please sign in to comment.