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