123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274 |
- import os
- import cv2
- import argparse
- import glob
- import torch
- from torchvision.transforms.functional import normalize
- from basicsr.utils import imwrite, img2tensor, tensor2img
- from basicsr.utils.download_util import load_file_from_url
- from basicsr.utils.misc import gpu_is_available, get_device
- from facelib.utils.face_restoration_helper import FaceRestoreHelper
- from facelib.utils.misc import is_gray
- from basicsr.utils.registry import ARCH_REGISTRY
- pretrain_model_url = {
- 'restoration': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
- }
- def set_realesrgan():
- from basicsr.archs.rrdbnet_arch import RRDBNet
- from basicsr.utils.realesrgan_utils import RealESRGANer
- use_half = False
- if torch.cuda.is_available(): # set False in CPU/MPS mode
- no_half_gpu_list = ['1650', '1660'] # set False for GPUs that don't support f16
- if not True in [gpu in torch.cuda.get_device_name(0) for gpu in no_half_gpu_list]:
- use_half = True
- model = RRDBNet(
- num_in_ch=3,
- num_out_ch=3,
- num_feat=64,
- num_block=23,
- num_grow_ch=32,
- scale=2,
- )
- upsampler = RealESRGANer(
- scale=2,
- model_path="https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth",
- model=model,
- tile=args.bg_tile,
- tile_pad=40,
- pre_pad=0,
- half=use_half
- )
- if not gpu_is_available(): # CPU
- import warnings
- warnings.warn('Running on CPU now! Make sure your PyTorch version matches your CUDA.'
- 'The unoptimized RealESRGAN is slow on CPU. '
- 'If you want to disable it, please remove `--bg_upsampler` and `--face_upsample` in command.',
- category=RuntimeWarning)
- return upsampler
- if __name__ == '__main__':
- # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- device = get_device()
- parser = argparse.ArgumentParser()
- parser.add_argument('-i', '--input_path', type=str, default='./inputs/whole_imgs',
- 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('-w', '--fidelity_weight', type=float, default=0.5,
- help='Balance the quality and fidelity. Default: 0.5')
- parser.add_argument('-s', '--upscale', type=int, default=2,
- help='The final upsampling scale of the image. Default: 2')
- 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')
- # large det_model: 'YOLOv5l', 'retinaface_resnet50'
- # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
- parser.add_argument('--detection_model', type=str, default='retinaface_resnet50',
- help='Face detector. Optional: retinaface_resnet50, retinaface_mobile0.25, YOLOv5l, YOLOv5n, dlib. \
- Default: retinaface_resnet50')
- parser.add_argument('--bg_upsampler', type=str, default='None', help='Background upsampler. Optional: realesrgan')
- parser.add_argument('--face_upsample', action='store_true', help='Face upsampler after enhancement. Default: False')
- parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')
- parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces. Default: None')
- parser.add_argument('--save_video_fps', type=float, default=None, help='Frame rate for saving video. Default: None')
- args = parser.parse_args()
- # ------------------------ input & output ------------------------
- w = args.fidelity_weight
- input_video = False
- 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_{w}'
- elif args.input_path.endswith(('mp4', 'mov', 'avi', 'MP4', 'MOV', 'AVI')): # input video path
- from basicsr.utils.video_util import VideoReader, VideoWriter
- input_img_list = []
- vidreader = VideoReader(args.input_path)
- image = vidreader.get_frame()
- while image is not None:
- input_img_list.append(image)
- image = vidreader.get_frame()
- audio = vidreader.get_audio()
- fps = vidreader.get_fps() if args.save_video_fps is None else args.save_video_fps
- video_name = os.path.basename(args.input_path)[:-4]
- result_root = f'results/{video_name}_{w}'
- input_video = True
- vidreader.close()
- 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 = f'results/{os.path.basename(args.input_path)}_{w}'
- 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 background upsampler ------------------
- if args.bg_upsampler == 'realesrgan':
- bg_upsampler = set_realesrgan()
- else:
- bg_upsampler = None
- # ------------------ set up face upsampler ------------------
- if args.face_upsample:
- if bg_upsampler is not None:
- face_upsampler = bg_upsampler
- else:
- face_upsampler = set_realesrgan()
- else:
- face_upsampler = None
- # ------------------ set up CodeFormer restorer -------------------
- net = ARCH_REGISTRY.get('CodeFormer')(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9,
- connect_list=['32', '64', '128', '256']).to(device)
-
- # ckpt_path = 'weights/CodeFormer/codeformer.pth'
- ckpt_path = load_file_from_url(url=pretrain_model_url['restoration'],
- model_dir='weights/CodeFormer', progress=True, file_name=None)
- checkpoint = torch.load(ckpt_path)['params_ema']
- net.load_state_dict(checkpoint)
- net.eval()
- # ------------------ set up FaceRestoreHelper -------------------
- # large det_model: 'YOLOv5l', 'retinaface_resnet50'
- # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
- if not args.has_aligned:
- print(f'Face detection model: {args.detection_model}')
- if bg_upsampler is not None:
- print(f'Background upsampling: True, Face upsampling: {args.face_upsample}')
- else:
- print(f'Background upsampling: False, Face upsampling: {args.face_upsample}')
- face_helper = FaceRestoreHelper(
- args.upscale,
- face_size=512,
- crop_ratio=(1, 1),
- det_model = args.detection_model,
- save_ext='png',
- use_parse=True,
- device=device)
- # -------------------- start to processing ---------------------
- for i, img_path in enumerate(input_img_list):
- # clean all the intermediate results to process the next image
- face_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)
- else: # for video processing
- basename = str(i).zfill(6)
- img_name = f'{video_name}_{basename}' if input_video else basename
- print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
- img = img_path
- if args.has_aligned:
- # the input faces are already cropped and aligned
- img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
- face_helper.is_gray = is_gray(img, threshold=10)
- if face_helper.is_gray:
- print('Grayscale input: True')
- face_helper.cropped_faces = [img]
- else:
- face_helper.read_image(img)
- # get face landmarks for each face
- num_det_faces = face_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
- face_helper.align_warp_face()
- # face restoration for each cropped face
- for idx, cropped_face in enumerate(face_helper.cropped_faces):
- # prepare data
- cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
- normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
- cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
- try:
- with torch.no_grad():
- output = net(cropped_face_t, w=w, adain=True)[0]
- restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
- del output
- torch.cuda.empty_cache()
- except Exception as error:
- print(f'\tFailed inference for CodeFormer: {error}')
- restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
- restored_face = restored_face.astype('uint8')
- face_helper.add_restored_face(restored_face, cropped_face)
- # paste_back
- if not args.has_aligned:
- # upsample the background
- if bg_upsampler is not None:
- # Now only support RealESRGAN for upsampling background
- bg_img = bg_upsampler.enhance(img, outscale=args.upscale)[0]
- else:
- bg_img = None
- face_helper.get_inverse_affine(None)
- # paste each restored face to the input image
- if args.face_upsample and face_upsampler is not None:
- restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box, face_upsampler=face_upsampler)
- else:
- restored_img = face_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)
- # save enhanced video
- if input_video:
- print('Video Saving...')
- # load images
- video_frames = []
- img_list = sorted(glob.glob(os.path.join(result_root, 'final_results', '*.[jp][pn]g')))
- for img_path in img_list:
- img = cv2.imread(img_path)
- video_frames.append(img)
- # write images to video
- height, width = video_frames[0].shape[:2]
- if args.suffix is not None:
- video_name = f'{video_name}_{args.suffix}.png'
- save_restore_path = os.path.join(result_root, f'{video_name}.mp4')
- vidwriter = VideoWriter(save_restore_path, height, width, fps, audio)
-
- for f in video_frames:
- vidwriter.write_frame(f)
- vidwriter.close()
- print(f'\nAll results are saved in {result_root}')
|