python/run_whole_image.py

import os
import cv2
import argparse
import glob

import numpy as np
from utils.general import imwrite
from utils.restoration_helper import RestoreHelper

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('-i', '--input_path', type=str, default='./pic', 
            help='Input image, video or folder. Default: inputs/whole_imgs')
    parser.add_argument('-o', '--output_path', type=str, default=None, 
            help='Output folder. Default: results/<input_name>_<w>')
    parser.add_argument('-s', '--upscale', type=int, default=1, 
            help='The final upsampling scale of the image. Default: 1')
    parser.add_argument('--detect_model', type=str, default='yolov5l-face.axmodel', help='face detection model path')
    parser.add_argument('--restore_model', type=str, default='codeformer.axmodel', help='face restore model path')
    parser.add_argument('--bg_model', type=str, default='realesrgan-x2.axmodel', help='background upsampler model path')
    parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces. Default: False')
    parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face. Default: False')
    parser.add_argument('--draw_box', action='store_true', help='Draw the bounding box for the detected faces. Default: False')
    parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces. Default: None')

    args = parser.parse_args()

    # ------------------------ input & output ------------------------
    if args.input_path.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')): # input single img path
        input_img_list = [args.input_path]
        result_root = f'results/test_img_{args.upscale}'
    else: # input img folder
        if args.input_path.endswith('/'):  # solve when path ends with /
            args.input_path = args.input_path[:-1]
        # scan all the jpg and png images
        input_img_list = sorted(glob.glob(os.path.join(args.input_path, '*.[jpJP][pnPN]*[gG]')))
        result_root = 'results'

    if not args.output_path is None: # set output path
        result_root = args.output_path

    test_img_num = len(input_img_list)
    if test_img_num == 0:
        raise FileNotFoundError('No input image/video is found...\n' 
            '\tNote that --input_path for video should end with .mp4|.mov|.avi')
    
    # ------------------ set up FaceRestoreHelper -------------------
    restore_helper = RestoreHelper(
        args.upscale,
        face_size=512,
        crop_ratio=(1, 1),
        det_model=args.detect_model,
        res_model=args.restore_model,
        bg_model=args.bg_model,
        save_ext='png',
        use_parse=True
        )

    # -------------------- start to processing ---------------------
    for i, img_path in enumerate(input_img_list):
        # clean all the intermediate results to process the next image
        restore_helper.clean_all()
        
        if isinstance(img_path, str):
            img_name = os.path.basename(img_path)
            basename, ext = os.path.splitext(img_name)
            print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
            img = cv2.imread(img_path, cv2.IMREAD_COLOR)

        restore_helper.read_image(img)
        # get face landmarks for each face
        num_det_faces = restore_helper.get_face_landmarks_5(
            only_center_face=args.only_center_face, resize=640, eye_dist_threshold=5)
        print(f'\tdetect {num_det_faces} faces')
        # align and warp each face
        restore_helper.align_warp_face()
        # face restoration for each cropped face
        for idx, cropped_face in enumerate(restore_helper.cropped_faces):
            # prepare data
            cropped_face_t = (cropped_face.astype(np.float32) / 255.0) * 2.0 - 1.0
            cropped_face_t = np.transpose(
              np.expand_dims(np.ascontiguousarray(cropped_face_t[...,::-1]), axis=0), 
              (0,3,1,2)
            )
            #print('cropped_face_t', cropped_face_t.shape)

            try:
                ort_outs = restore_helper.rs_sessison.run(
                    restore_helper.rs_output, 
                    {restore_helper.rs_input: cropped_face_t}
                    )
                restored_face = ort_outs[0]
                restored_face = (restored_face.squeeze().transpose(1, 2, 0) * 0.5 + 0.5) * 255
                restored_face = np.clip(restored_face[...,::-1], 0, 255).astype(np.uint8)
            except Exception as error:
                print(f'\tFailed inference for CodeFormer: {error}')
                restored_face = (cropped_face_t.squeeze().transpose(1, 2, 0) * 0.5 + 0.5) * 255
                restored_face = np.clip(restored_face, 0, 255).astype(np.uint8)

            restored_face = restored_face.astype('uint8')
            restore_helper.add_restored_face(restored_face, cropped_face)


        # paste_back
        if not args.has_aligned:
            # upsample the background
            # Now only support RealESRGAN for upsampling background
            bg_img = restore_helper.background_upsampling(img)
            restore_helper.get_inverse_affine(None)
            # paste each restored face to the input image
            restored_img = restore_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box)

        # save faces
        # for idx, (cropped_face, restored_face) in enumerate(zip(face_helper.cropped_faces, face_helper.restored_faces)):
            # # save cropped face
            # if not args.has_aligned: 
                # save_crop_path = os.path.join(result_root, 'cropped_faces', f'{basename}_{idx:02d}.png')
                # imwrite(cropped_face, save_crop_path)
            # # save restored face
            # if args.has_aligned:
                # save_face_name = f'{basename}.png'
            # else:
                # save_face_name = f'{basename}_{idx:02d}.png'
            # if args.suffix is not None:
                # save_face_name = f'{save_face_name[:-4]}_{args.suffix}.png'
            # save_restore_path = os.path.join(result_root, 'restored_faces', save_face_name)
            # imwrite(restored_face, save_restore_path)

        # save restored img
        if not args.has_aligned and restored_img is not None:
            if args.suffix is not None:
                basename = f'{basename}_{args.suffix}'
            save_restore_path = os.path.join(result_root, 'final_results', f'{basename}.png')
            imwrite(restored_img, save_restore_path)