faisalhr1997/codeformer
0
1# Modified by Shangchen Zhou from: https://github.com/TencentARC/GFPGAN/blob/master/inference_gfpgan.py2import os3import cv24import argparse5import glob6import torch7from torchvision.transforms.functional import normalize8from basicsr.utils import imwrite, img2tensor, tensor2img9from basicsr.utils.download_util import load_file_from_url10from facelib.utils.face_restoration_helper import FaceRestoreHelper11import torch.nn.functional as F12 13from basicsr.utils.registry import ARCH_REGISTRY14 15pretrain_model_url = {16 'restoration': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',17}18 19def set_realesrgan():20 if not torch.cuda.is_available(): # CPU21 import warnings22 warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '23 'If you really want to use it, please modify the corresponding codes.',24 category=RuntimeWarning)25 bg_upsampler = None26 else:27 from basicsr.archs.rrdbnet_arch import RRDBNet28 from basicsr.utils.realesrgan_utils import RealESRGANer29 model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)30 bg_upsampler = RealESRGANer(31 scale=2,32 model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',33 model=model,34 tile=args.bg_tile,35 tile_pad=40,36 pre_pad=0,37 half=True) # need to set False in CPU mode38 return bg_upsampler39 40if __name__ == '__main__':41 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')42 parser = argparse.ArgumentParser()43 44 parser.add_argument('--w', type=float, default=0.5, help='Balance the quality and fidelity')45 parser.add_argument('--upscale', type=int, default=2, help='The final upsampling scale of the image. Default: 2')46 parser.add_argument('--test_path', type=str, default='./inputs/cropped_faces')47 parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces')48 parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face')49 # large det_model: 'YOLOv5l', 'retinaface_resnet50'50 # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'51 parser.add_argument('--detection_model', type=str, default='retinaface_resnet50')52 parser.add_argument('--draw_box', action='store_true')53 parser.add_argument('--bg_upsampler', type=str, default='None', help='background upsampler. Optional: realesrgan')54 parser.add_argument('--face_upsample', action='store_true', help='face upsampler after enhancement.')55 parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')56 57 args = parser.parse_args()58 59 # ------------------------ input & output ------------------------60 if args.test_path.endswith('/'): # solve when path ends with /61 args.test_path = args.test_path[:-1]62 63 w = args.w64 result_root = f'results/{os.path.basename(args.test_path)}_{w}'65 66 # ------------------ set up background upsampler ------------------67 if args.bg_upsampler == 'realesrgan':68 bg_upsampler = set_realesrgan()69 else:70 bg_upsampler = None71 72 # ------------------ set up face upsampler ------------------73 if args.face_upsample:74 if bg_upsampler is not None:75 face_upsampler = bg_upsampler76 else:77 face_upsampler = set_realesrgan()78 else:79 face_upsampler = None80 81 # ------------------ set up CodeFormer restorer -------------------82 net = ARCH_REGISTRY.get('CodeFormer')(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, 83 connect_list=['32', '64', '128', '256']).to(device)84 85 # ckpt_path = 'weights/CodeFormer/codeformer.pth'86 ckpt_path = load_file_from_url(url=pretrain_model_url['restoration'], 87 model_dir='weights/CodeFormer', progress=True, file_name=None)88 checkpoint = torch.load(ckpt_path)['params_ema']89 net.load_state_dict(checkpoint)90 net.eval()91 92 # ------------------ set up FaceRestoreHelper -------------------93 # large det_model: 'YOLOv5l', 'retinaface_resnet50'94 # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'95 if not args.has_aligned: 96 print(f'Face detection model: {args.detection_model}')97 if bg_upsampler is not None: 98 print(f'Background upsampling: True, Face upsampling: {args.face_upsample}')99 else:100 print(f'Background upsampling: False, Face upsampling: {args.face_upsample}')101 102 face_helper = FaceRestoreHelper(103 args.upscale,104 face_size=512,105 crop_ratio=(1, 1),106 det_model = args.detection_model,107 save_ext='png',108 use_parse=True,109 device=device)110 111 # -------------------- start to processing ---------------------112 # scan all the jpg and png images113 for img_path in sorted(glob.glob(os.path.join(args.test_path, '*.[jp][pn]g'))):114 # clean all the intermediate results to process the next image115 face_helper.clean_all()116 117 img_name = os.path.basename(img_path)118 print(f'Processing: {img_name}')119 basename, ext = os.path.splitext(img_name)120 img = cv2.imread(img_path, cv2.IMREAD_COLOR)121 122 if args.has_aligned: 123 # the input faces are already cropped and aligned124 img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)125 face_helper.cropped_faces = [img]126 else:127 face_helper.read_image(img)128 # get face landmarks for each face129 num_det_faces = face_helper.get_face_landmarks_5(130 only_center_face=args.only_center_face, resize=640, eye_dist_threshold=5)131 print(f'\tdetect {num_det_faces} faces')132 # align and warp each face133 face_helper.align_warp_face()134 135 # face restoration for each cropped face136 for idx, cropped_face in enumerate(face_helper.cropped_faces):137 # prepare data138 cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)139 normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)140 cropped_face_t = cropped_face_t.unsqueeze(0).to(device)141 142 try:143 with torch.no_grad():144 output = net(cropped_face_t, w=w, adain=True)[0]145 restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))146 del output147 torch.cuda.empty_cache()148 except Exception as error:149 print(f'\tFailed inference for CodeFormer: {error}')150 restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))151 152 restored_face = restored_face.astype('uint8')153 face_helper.add_restored_face(restored_face)154 155 # paste_back156 if not args.has_aligned:157 # upsample the background158 if bg_upsampler is not None:159 # Now only support RealESRGAN for upsampling background160 bg_img = bg_upsampler.enhance(img, outscale=args.upscale)[0]161 else:162 bg_img = None163 face_helper.get_inverse_affine(None)164 # paste each restored face to the input image165 if args.face_upsample and face_upsampler is not None: 166 restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box, face_upsampler=face_upsampler)167 else:168 restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box)169 170 # save faces171 for idx, (cropped_face, restored_face) in enumerate(zip(face_helper.cropped_faces, face_helper.restored_faces)):172 # save cropped face173 if not args.has_aligned: 174 save_crop_path = os.path.join(result_root, 'cropped_faces', f'{basename}_{idx:02d}.png')175 imwrite(cropped_face, save_crop_path)176 # save restored face177 if args.has_aligned:178 save_face_name = f'{basename}.png'179 else:180 save_face_name = f'{basename}_{idx:02d}.png'181 save_restore_path = os.path.join(result_root, 'restored_faces', save_face_name)182 imwrite(restored_face, save_restore_path)183 184 # save restored img185 if not args.has_aligned and restored_img is not None:186 save_restore_path = os.path.join(result_root, 'final_results', f'{basename}.png')187 imwrite(restored_img, save_restore_path)188 189 print(f'\nAll results are saved in {result_root}')190