AXERA-TECH/CodeFormer
333
1import os
2import cv2
3import argparse
4import glob
5
6import numpy as np
7from utils.general import imwrite
8from utils.restoration_helper import RestoreHelper
9
10if __name__ == '__main__':
11 parser = argparse.ArgumentParser()
12
13 parser.add_argument('-i', '--input_path', type=str, default='./pic',
14 help='Input image, video or folder. Default: inputs/whole_imgs')
15 parser.add_argument('-o', '--output_path', type=str, default=None,
16 help='Output folder. Default: results/<input_name>_<w>')
17 parser.add_argument('-s', '--upscale', type=int, default=1,
18 help='The final upsampling scale of the image. Default: 1')
19 parser.add_argument('--detect_model', type=str, default='yolov5l-face.axmodel', help='face detection model path')
20 parser.add_argument('--restore_model', type=str, default='codeformer.axmodel', help='face restore model path')
21 parser.add_argument('--bg_model', type=str, default='realesrgan-x2.axmodel', help='background upsampler model path')
22 parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces. Default: False')
23 parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face. Default: False')
24 parser.add_argument('--draw_box', action='store_true', help='Draw the bounding box for the detected faces. Default: False')
25 parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces. Default: None')
26
27 args = parser.parse_args()
28
29 # ------------------------ input & output ------------------------
30 if args.input_path.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')): # input single img path
31 input_img_list = [args.input_path]
32 result_root = f'results/test_img_{args.upscale}'
33 else: # input img folder
34 if args.input_path.endswith('/'): # solve when path ends with /
35 args.input_path = args.input_path[:-1]
36 # scan all the jpg and png images
37 input_img_list = sorted(glob.glob(os.path.join(args.input_path, '*.[jpJP][pnPN]*[gG]')))
38 result_root = 'results'
39
40 if not args.output_path is None: # set output path
41 result_root = args.output_path
42
43 test_img_num = len(input_img_list)
44 if test_img_num == 0:
45 raise FileNotFoundError('No input image/video is found...\n'
46 '\tNote that --input_path for video should end with .mp4|.mov|.avi')
47
48 # ------------------ set up FaceRestoreHelper -------------------
49 restore_helper = RestoreHelper(
50 args.upscale,
51 face_size=512,
52 crop_ratio=(1, 1),
53 det_model=args.detect_model,
54 res_model=args.restore_model,
55 bg_model=args.bg_model,
56 save_ext='png',
57 use_parse=True
58 )
59
60 # -------------------- start to processing ---------------------
61 for i, img_path in enumerate(input_img_list):
62 # clean all the intermediate results to process the next image
63 restore_helper.clean_all()
64
65 if isinstance(img_path, str):
66 img_name = os.path.basename(img_path)
67 basename, ext = os.path.splitext(img_name)
68 print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
69 img = cv2.imread(img_path, cv2.IMREAD_COLOR)
70
71 restore_helper.read_image(img)
72 # get face landmarks for each face
73 num_det_faces = restore_helper.get_face_landmarks_5(
74 only_center_face=args.only_center_face, resize=640, eye_dist_threshold=5)
75 print(f'\tdetect {num_det_faces} faces')
76 # align and warp each face
77 restore_helper.align_warp_face()
78 # face restoration for each cropped face
79 for idx, cropped_face in enumerate(restore_helper.cropped_faces):
80 # prepare data
81 cropped_face_t = (cropped_face.astype(np.float32) / 255.0) * 2.0 - 1.0
82 cropped_face_t = np.transpose(
83 np.expand_dims(np.ascontiguousarray(cropped_face_t[...,::-1]), axis=0),
84 (0,3,1,2)
85 )
86 #print('cropped_face_t', cropped_face_t.shape)
87
88 try:
89 ort_outs = restore_helper.rs_sessison.run(
90 restore_helper.rs_output,
91 {restore_helper.rs_input: cropped_face_t}
92 )
93 restored_face = ort_outs[0]
94 restored_face = (restored_face.squeeze().transpose(1, 2, 0) * 0.5 + 0.5) * 255
95 restored_face = np.clip(restored_face[...,::-1], 0, 255).astype(np.uint8)
96 except Exception as error:
97 print(f'\tFailed inference for CodeFormer: {error}')
98 restored_face = (cropped_face_t.squeeze().transpose(1, 2, 0) * 0.5 + 0.5) * 255
99 restored_face = np.clip(restored_face, 0, 255).astype(np.uint8)
100
101 restored_face = restored_face.astype('uint8')
102 restore_helper.add_restored_face(restored_face, cropped_face)
103
104
105 # paste_back
106 if not args.has_aligned:
107 # upsample the background
108 # Now only support RealESRGAN for upsampling background
109 bg_img = restore_helper.background_upsampling(img)
110 restore_helper.get_inverse_affine(None)
111 # paste each restored face to the input image
112 restored_img = restore_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box)
113
114 # save faces
115 # for idx, (cropped_face, restored_face) in enumerate(zip(face_helper.cropped_faces, face_helper.restored_faces)):
116 # # save cropped face
117 # if not args.has_aligned:
118 # save_crop_path = os.path.join(result_root, 'cropped_faces', f'{basename}_{idx:02d}.png')
119 # imwrite(cropped_face, save_crop_path)
120 # # save restored face
121 # if args.has_aligned:
122 # save_face_name = f'{basename}.png'
123 # else:
124 # save_face_name = f'{basename}_{idx:02d}.png'
125 # if args.suffix is not None:
126 # save_face_name = f'{save_face_name[:-4]}_{args.suffix}.png'
127 # save_restore_path = os.path.join(result_root, 'restored_faces', save_face_name)
128 # imwrite(restored_face, save_restore_path)
129
130 # save restored img
131 if not args.has_aligned and restored_img is not None:
132 if args.suffix is not None:
133 basename = f'{basename}_{args.suffix}'
134 save_restore_path = os.path.join(result_root, 'final_results', f'{basename}.png')
135 imwrite(restored_img, save_restore_path)
136
137