reformat and fix temp input dir bug
This commit is contained in:
CJWBW
GitHub
Родитель
05d9ed4445
Коммит
e9f55eca09
279
predict.py
279
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 @@ class Predictor(cog.Predictor):
|
|||
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)
|
||||
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
|
||||
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
|
||||
## 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"
|
||||
)
|
||||
|
||||
if HR:
|
||||
HR_suffix = " --HR"
|
||||
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:
|
||||
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
|
||||
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_1_command_1)
|
||||
run_cmd(stage_1_command_2)
|
||||
run_cmd(stage_2_command)
|
||||
print("Finish Stage 2 ...")
|
||||
print("\n")
|
||||
|
||||
## 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)
|
||||
## 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"
|
||||
)
|
||||
|
||||
print("Finish Stage 1 ...")
|
||||
print("\n")
|
||||
os.makedirs(stage_3_output_dir, exist_ok=True)
|
||||
|
||||
## Stage 2: Face Detection
|
||||
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"
|
||||
)
|
||||
|
||||
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)
|
||||
run_cmd(stage_3_command)
|
||||
print("Finish Stage 3 ...")
|
||||
print("\n")
|
||||
|
||||
stage_2_command = (
|
||||
"python detect_all_dlib_HR.py --url "
|
||||
+ stage_2_input_dir
|
||||
+ " --save_url "
|
||||
+ stage_2_output_dir
|
||||
)
|
||||
## 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)
|
||||
|
||||
run_cmd(stage_2_command)
|
||||
print("Finish Stage 2 ...")
|
||||
print("\n")
|
||||
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
|
||||
)
|
||||
|
||||
## 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_4_command)
|
||||
print("Finish Stage 4 ...")
|
||||
print("\n")
|
||||
|
||||
os.makedirs(stage_3_output_dir, exist_ok=True)
|
||||
print("All the processing is done. Please check the results.")
|
||||
|
||||
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"
|
||||
)
|
||||
final_output = os.listdir(os.path.join(self.opts.output_folder, "final_output"))[0]
|
||||
|
||||
run_cmd(stage_3_command)
|
||||
print("Finish Stage 3 ...")
|
||||
print("\n")
|
||||
image_restore = cv2.imread(os.path.join(self.opts.output_folder, "final_output", final_output))
|
||||
|
||||
## 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)
|
||||
out_path = Path(tempfile.mkdtemp()) / "out.png"
|
||||
|
||||
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)
|
||||
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}")
|
||||
|
|
|
|||
Загрузка…
Ссылка в новой задаче