manhkhanhUIT/BOPBTL

1import numpy as np2import cv23import PySimpleGUI as sg4import os.path5import argparse6import os7import sys8import shutil9from subprocess import call10 11def modify(image_filename=None, cv2_frame=None):12 13    def run_cmd(command):14        try:15            call(command, shell=True)16        except KeyboardInterrupt:17            print("Process interrupted")18            sys.exit(1)19 20    parser = argparse.ArgumentParser()21    parser.add_argument("--input_folder", type=str,22                        default= image_filename, help="Test images")23    parser.add_argument(24        "--output_folder",25        type=str,26        default="./output",27        help="Restored images, please use the absolute path",28    )29    parser.add_argument("--GPU", type=str, default="-1", help="0,1,2")30    parser.add_argument(31        "--checkpoint_name", type=str, default="Setting_9_epoch_100", help="choose which checkpoint"32    )33    parser.add_argument("--with_scratch",default="--with_scratch" ,action="store_true")34    opts = parser.parse_args()35 36    gpu1 = opts.GPU37 38    # resolve relative paths before changing directory39    opts.input_folder = os.path.abspath(opts.input_folder)40    opts.output_folder = os.path.abspath(opts.output_folder)41    if not os.path.exists(opts.output_folder):42        os.makedirs(opts.output_folder)43 44    main_environment = os.getcwd()45 46    # Stage 1: Overall Quality Improve47    print("Running Stage 1: Overall restoration")48    os.chdir("./Global")49    stage_1_input_dir = opts.input_folder50    stage_1_output_dir = os.path.join(51        opts.output_folder, "stage_1_restore_output")52    if not os.path.exists(stage_1_output_dir):53        os.makedirs(stage_1_output_dir)54 55    if not opts.with_scratch:56        stage_1_command = (57            "python test.py --test_mode Full --Quality_restore --test_input "58            + stage_1_input_dir59            + " --outputs_dir "60            + stage_1_output_dir61            + " --gpu_ids "62            + gpu163        )64        run_cmd(stage_1_command)65    else:66 67        mask_dir = os.path.join(stage_1_output_dir, "masks")68        new_input = os.path.join(mask_dir, "input")69        new_mask = os.path.join(mask_dir, "mask")70        stage_1_command_1 = (71            "python detection.py --test_path "72            + stage_1_input_dir73            + " --output_dir "74            + mask_dir75            + " --input_size full_size"76            + " --GPU "77            + gpu178        )79        stage_1_command_2 = (80            "python test.py --Scratch_and_Quality_restore --test_input "81            + new_input82            + " --test_mask "83            + new_mask84            + " --outputs_dir "85            + stage_1_output_dir86            + " --gpu_ids "87            + gpu188        )89        run_cmd(stage_1_command_1)90        run_cmd(stage_1_command_2)91 92    # Solve the case when there is no face in the old photo93    stage_1_results = os.path.join(stage_1_output_dir, "restored_image")94    stage_4_output_dir = os.path.join(opts.output_folder, "final_output")95    if not os.path.exists(stage_4_output_dir):96        os.makedirs(stage_4_output_dir)97    for x in os.listdir(stage_1_results):98        img_dir = os.path.join(stage_1_results, x)99        shutil.copy(img_dir, stage_4_output_dir)100 101    print("Finish Stage 1 ...")102    print("\n")103 104    # Stage 2: Face Detection105 106    print("Running Stage 2: Face Detection")107    os.chdir(".././Face_Detection")108    stage_2_input_dir = os.path.join(stage_1_output_dir, "restored_image")109    stage_2_output_dir = os.path.join(110        opts.output_folder, "stage_2_detection_output")111    if not os.path.exists(stage_2_output_dir):112        os.makedirs(stage_2_output_dir)113    stage_2_command = (114        "python detect_all_dlib.py --url " + stage_2_input_dir +115        " --save_url " + stage_2_output_dir116    )117    run_cmd(stage_2_command)118    print("Finish Stage 2 ...")119    print("\n")120 121    # Stage 3: Face Restore122    print("Running Stage 3: Face Enhancement")123    os.chdir(".././Face_Enhancement")124    stage_3_input_mask = "./"125    stage_3_input_face = stage_2_output_dir126    stage_3_output_dir = os.path.join(127        opts.output_folder, "stage_3_face_output")128    if not os.path.exists(stage_3_output_dir):129        os.makedirs(stage_3_output_dir)130    stage_3_command = (131        "python test_face.py --old_face_folder "132        + stage_3_input_face133        + " --old_face_label_folder "134        + stage_3_input_mask135        + " --tensorboard_log --name "136        + opts.checkpoint_name137        + " --gpu_ids "138        + gpu1139        + " --load_size 256 --label_nc 18 --no_instance --preprocess_mode resize --batchSize 4 --results_dir "140        + stage_3_output_dir141        + " --no_parsing_map"142    )143    run_cmd(stage_3_command)144    print("Finish Stage 3 ...")145    print("\n")146 147    # Stage 4: Warp back148    print("Running Stage 4: Blending")149    os.chdir(".././Face_Detection")150    stage_4_input_image_dir = os.path.join(151        stage_1_output_dir, "restored_image")152    stage_4_input_face_dir = os.path.join(stage_3_output_dir, "each_img")153    stage_4_output_dir = os.path.join(opts.output_folder, "final_output")154    if not os.path.exists(stage_4_output_dir):155        os.makedirs(stage_4_output_dir)156    stage_4_command = (157        "python align_warp_back_multiple_dlib.py --origin_url "158        + stage_4_input_image_dir159        + " --replace_url "160        + stage_4_input_face_dir161        + " --save_url "162        + stage_4_output_dir163    )164    run_cmd(stage_4_command)165    print("Finish Stage 4 ...")166    print("\n")167 168    print("All the processing is done. Please check the results.")169 170# --------------------------------- The GUI ---------------------------------171 172# First the window layout...173 174images_col = [[sg.Text('Input file:'), sg.In(enable_events=True, key='-IN FILE-'), sg.FileBrowse()],175              [sg.Button('Modify Photo', key='-MPHOTO-'), sg.Button('Exit')],176              [sg.Image(filename='', key='-IN-'), sg.Image(filename='', key='-OUT-')],]177# ----- Full layout -----178layout = [[sg.VSeperator(), sg.Column(images_col)]]179 180# ----- Make the window -----181window = sg.Window('Bringing-old-photos-back-to-life', layout, grab_anywhere=True)182 183# ----- Run the Event Loop -----184prev_filename = colorized = cap = None185while True:186    event, values = window.read()187    if event in (None, 'Exit'):188        break189 190    elif event == '-MPHOTO-':191        try:192            n1 = filename.split("/")[-2]193            n2 = filename.split("/")[-3]194            n3 = filename.split("/")[-1]195            filename= str(f"./{n2}/{n1}")196            modify(filename)197           198            global f_image199            f_image = f'./output/final_output/{n3}'200            image = cv2.imread(f_image)201            window['-OUT-'].update(data=cv2.imencode('.png', image)[1].tobytes())202            203        except:204            continue205 206    elif event == '-IN FILE-':      # A single filename was chosen207        filename = values['-IN FILE-']208        if filename != prev_filename:209            prev_filename = filename210            try:211                image = cv2.imread(filename)212                window['-IN-'].update(data=cv2.imencode('.png', image)[1].tobytes())213            except:214                continue215 216# ----- Exit program -----217window.close()