217 строки
7.1 KiB
Python
217 строки
7.1 KiB
Python
import numpy as np
|
|
import cv2
|
|
import PySimpleGUI as sg
|
|
import os.path
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import shutil
|
|
from subprocess import call
|
|
|
|
def modify(image_filename=None, cv2_frame=None):
|
|
|
|
def run_cmd(command):
|
|
try:
|
|
call(command, shell=True)
|
|
except KeyboardInterrupt:
|
|
print("Process interrupted")
|
|
sys.exit(1)
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_folder", type=str,
|
|
default= image_filename, 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="-1", help="0,1,2")
|
|
parser.add_argument(
|
|
"--checkpoint_name", type=str, default="Setting_9_epoch_100", help="choose which checkpoint"
|
|
)
|
|
parser.add_argument("--with_scratch",default="--with_scratch" ,action="store_true")
|
|
opts = parser.parse_args()
|
|
|
|
gpu1 = opts.GPU
|
|
|
|
# resolve relative paths before changing directory
|
|
opts.input_folder = os.path.abspath(opts.input_folder)
|
|
opts.output_folder = os.path.abspath(opts.output_folder)
|
|
if not os.path.exists(opts.output_folder):
|
|
os.makedirs(opts.output_folder)
|
|
|
|
main_environment = os.getcwd()
|
|
|
|
# Stage 1: Overall Quality Improve
|
|
print("Running Stage 1: Overall restoration")
|
|
os.chdir("./Global")
|
|
stage_1_input_dir = opts.input_folder
|
|
stage_1_output_dir = os.path.join(
|
|
opts.output_folder, "stage_1_restore_output")
|
|
if not os.path.exists(stage_1_output_dir):
|
|
os.makedirs(stage_1_output_dir)
|
|
|
|
if not opts.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_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
|
|
)
|
|
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(opts.output_folder, "final_output")
|
|
if not os.path.exists(stage_4_output_dir):
|
|
os.makedirs(stage_4_output_dir)
|
|
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(
|
|
opts.output_folder, "stage_2_detection_output")
|
|
if not os.path.exists(stage_2_output_dir):
|
|
os.makedirs(stage_2_output_dir)
|
|
stage_2_command = (
|
|
"python detect_all_dlib.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(
|
|
opts.output_folder, "stage_3_face_output")
|
|
if not os.path.exists(stage_3_output_dir):
|
|
os.makedirs(stage_3_output_dir)
|
|
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 "
|
|
+ opts.checkpoint_name
|
|
+ " --gpu_ids "
|
|
+ gpu1
|
|
+ " --load_size 256 --label_nc 18 --no_instance --preprocess_mode resize --batchSize 4 --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(opts.output_folder, "final_output")
|
|
if not os.path.exists(stage_4_output_dir):
|
|
os.makedirs(stage_4_output_dir)
|
|
stage_4_command = (
|
|
"python align_warp_back_multiple_dlib.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.")
|
|
|
|
# --------------------------------- The GUI ---------------------------------
|
|
|
|
# First the window layout...
|
|
|
|
images_col = [[sg.Text('Input file:'), sg.In(enable_events=True, key='-IN FILE-'), sg.FileBrowse()],
|
|
[sg.Button('Modify Photo', key='-MPHOTO-'), sg.Button('Exit')],
|
|
[sg.Image(filename='', key='-IN-'), sg.Image(filename='', key='-OUT-')],]
|
|
# ----- Full layout -----
|
|
layout = [[sg.VSeperator(), sg.Column(images_col)]]
|
|
|
|
# ----- Make the window -----
|
|
window = sg.Window('Bringing-old-photos-back-to-life', layout, grab_anywhere=True)
|
|
|
|
# ----- Run the Event Loop -----
|
|
prev_filename = colorized = cap = None
|
|
while True:
|
|
event, values = window.read()
|
|
if event in (None, 'Exit'):
|
|
break
|
|
|
|
elif event == '-MPHOTO-':
|
|
try:
|
|
n1 = filename.split("/")[-2]
|
|
n2 = filename.split("/")[-3]
|
|
n3 = filename.split("/")[-1]
|
|
filename= str(f"./{n2}/{n1}")
|
|
modify(filename)
|
|
|
|
global f_image
|
|
f_image = f'./output/final_output/{n3}'
|
|
image = cv2.imread(f_image)
|
|
window['-OUT-'].update(data=cv2.imencode('.png', image)[1].tobytes())
|
|
|
|
except:
|
|
continue
|
|
|
|
elif event == '-IN FILE-': # A single filename was chosen
|
|
filename = values['-IN FILE-']
|
|
if filename != prev_filename:
|
|
prev_filename = filename
|
|
try:
|
|
image = cv2.imread(filename)
|
|
window['-IN-'].update(data=cv2.imencode('.png', image)[1].tobytes())
|
|
except:
|
|
continue
|
|
|
|
# ----- Exit program -----
|
|
window.close() |