devbox
小蓝xlanll
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

一键快速还原修复人脸，CodeFormer 助力人脸图像修复_codeformer解压密码

作者：小蓝xlanll | 2024-02-22 22:56:58

赞

踩

codeformer解压密码

今天在查资料的时候无意间看到了一个很有意思的工具，就是CodeFormer ，作者给出来的说明是用于人脸修复任务的，觉得很有意思就拿来实践了一下，这里记录分享一下。

首先对人脸修复任务进行简单的回顾总结：

人脸修复是指对损坏、有缺陷或者遮挡的人脸图像进行恢复、修复或重建的过程。这涉及到图像处理、计算机视觉和机器学习等领域的技术。

人脸修复可以应用于多个场景，例如：

  1. 人脸去除遮挡：当人脸图像被遮挡或者被其他对象遮挡时，人脸修复可以通过推断遮挡区域的内容进行修复。这可以包括使用图像修复算法（如基于纹理填充、基于复制粘贴的方法）或者利用人脸识别和人脸重建技术进行修复。

  2. 人脸恢复和重建：当人脸图像受到严重损坏或者缺陷时，通过人脸修复技术可以进行图像的恢复和重建。这可以包括利用图像补全算法、人脸超分辨率重建、人脸合成等技术来重建完整的人脸图像。

  3. 人脸修复和美化：人脸修复还可以包括对人脸图像的美化和修饰。这可以包括去除皱纹、瑕疵、眼袋等改善面部外貌的操作，以及改变皮肤色调、增强对比度等优化图像质量的操作。

在人脸修复领域，常用的算法模型包括以下几种：

  1. 自编码器（Autoencoder）：自编码器是一种无监督学习方法，将输入数据压缩为低维编码，然后再将编码恢复为原始数据。在人脸修复中，自编码器可以用于重建缺失或损坏的人脸部分。

  2. 卷积神经网络（Convolutional Neural Network, CNN）：CNN是一种深度学习模型，广泛用于图像处理任务。在人脸修复中，CNN可以用于将损坏的人脸图像与完整的人脸图像进行对齐，从而进行修复。

  3. 生成对抗网络（Generative Adversarial Network, GAN）：GAN由生成器和判别器组成，通过对抗学习的方式，生成逼真的以假乱真的图像。在人脸修复中，GAN可以用于生成缺失或损坏的人脸部分。

  4. 基于纹理填充的方法：基于纹理填充的方法使用纹理合成和纹理迁移技术，将周围的纹理信息应用于损坏区域，来修复人脸图像。

  5. 结构生成模型：结构生成模型基于人脸的结构信息，如关键点或网格，通过形状优化和插值方法，实现对人脸的修复和重建。

除了上述算法模型，还有一些其他的变种和组合方法，如图像修补算法、超分辨率重建算法、形状恢复算法等。这些算法模型都在不同场景下有其适用性和局限性，具体选择哪种模型需要根据具体任务和需求来进行权衡和选择。

接下来回归正题，来看本文要介绍学习的CodeFormer。

作者项目在这里，如下所示：

 盲脸恢复是一个高度不适定的问题，通常需要辅助指导来

1）改进从退化输入到期望输出的映射，或

2）补充输入中丢失的高质量细节。

在本文中，我们证明了在小代理空间中学习的离散码本先验通过将人脸恢复作为代码预测任务，大大降低了恢复映射的不确定性和模糊性，同时为生成高质量人脸提供了丰富的视觉原子。在这种范式下，我们提出了一种基于Transformer的预测网络，名为CodeFormer，用于对低质量人脸的全局组成和上下文进行建模，以进行代码预测，即使在输入严重退化的情况下，也能够发现与目标人脸非常接近的自然人脸。为了增强对不同退化的适应性，我们还提出了一个可控的特征转换模块，该模块允许在保真度和质量之间进行灵活的权衡。得益于富有表现力的码本先验和全局建模，CodeFormer在质量和保真度方面都优于现有技术，对退化表现出卓越的鲁棒性。在合成和真实世界数据集上的大量实验结果验证了我们方法的有效性。

文章的核心方法：

 （a） 我们首先通过自重构学习学习离散码本和解码器来存储人脸图像的高质量视觉部分。

（b） 在固定码本和解码器的情况下，我们引入了一个用于代码序列预测的Transformer模块，对低质量输入的全局人脸组成进行建模。此外，使用可控特征变换模块来控制从LQ编码器到解码器的信息流。请注意，此连接是可选的，当输入严重降级时，可以禁用此连接以避免不利影响，并且可以调整标量权重w以在质量和保真度之间进行权衡。

实例结果：

 官方项目地址在这里，如下所示：

 目前已经有将近1w的star量了，看来还是蛮受欢迎的。

官方的论文地址在这里，如下所示：

 感兴趣的话可以自行下载研读一下，这里我主要是想实际试用一下。

下载官方项目，本地解压缩如下所示：

 这里提供了可以直接使用的人脸修复推理模块，如下所示：

  1. import os
  2. import cv2
  3. import argparse
  4. import glob
  5. import torch
  6. from torchvision.transforms.functional import normalize
  7. from basicsr.utils import imwrite, img2tensor, tensor2img
  8. from basicsr.utils.download_util import load_file_from_url
  9. from basicsr.utils.misc import gpu_is_available, get_device
  10. from facelib.utils.face_restoration_helper import FaceRestoreHelper
  11. from facelib.utils.misc import is_gray
  12.  
  13. from basicsr.utils.registry import ARCH_REGISTRY
  14.  
  15. pretrain_model_url = {
  16.     'restoration': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
  17. }
  18.  
  19. def set_realesrgan():
  20.     from basicsr.archs.rrdbnet_arch import RRDBNet
  21.     from basicsr.utils.realesrgan_utils import RealESRGANer
  22.  
  23.     use_half = False
  24.     if torch.cuda.is_available(): # set False in CPU/MPS mode
  25.         no_half_gpu_list = ['1650', '1660'] # set False for GPUs that don't support f16
  26.         if not True in [gpu in torch.cuda.get_device_name(0) for gpu in no_half_gpu_list]:
  27.             use_half = True
  28.  
  29.     model = RRDBNet(
  30.         num_in_ch=3,
  31.         num_out_ch=3,
  32.         num_feat=64,
  33.         num_block=23,
  34.         num_grow_ch=32,
  35.         scale=2,
  36.     )
  37.     upsampler = RealESRGANer(
  38.         scale=2,
  39.         model_path="https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth",
  40.         model=model,
  41.         tile=args.bg_tile,
  42.         tile_pad=40,
  43.         pre_pad=0,
  44.         half=use_half
  45.     )
  46.  
  47.     if not gpu_is_available():  # CPU
  48.         import warnings
  49.         warnings.warn('Running on CPU now! Make sure your PyTorch version matches your CUDA.'
  50.                         'The unoptimized RealESRGAN is slow on CPU. '
  51.                         'If you want to disable it, please remove `--bg_upsampler` and `--face_upsample` in command.',
  52.                         category=RuntimeWarning)
  53.     return upsampler
  54.  
  55. if __name__ == '__main__':
  56.     # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  57.     device = get_device()
  58.     parser = argparse.ArgumentParser()
  59.  
  60.     parser.add_argument('-i', '--input_path', type=str, default='./inputs/whole_imgs', 
  61.             help='Input image, video or folder. Default: inputs/whole_imgs')
  62.     parser.add_argument('-o', '--output_path', type=str, default=None, 
  63.             help='Output folder. Default: results/<input_name>_<w>')
  64.     parser.add_argument('-w', '--fidelity_weight', type=float, default=0.5, 
  65.             help='Balance the quality and fidelity. Default: 0.5')
  66.     parser.add_argument('-s', '--upscale', type=int, default=2, 
  67.             help='The final upsampling scale of the image. Default: 2')
  68.     parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces. Default: False')
  69.     parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face. Default: False')
  70.     parser.add_argument('--draw_box', action='store_true', help='Draw the bounding box for the detected faces. Default: False')
  71.     # large det_model: 'YOLOv5l', 'retinaface_resnet50'
  72.     # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
  73.     parser.add_argument('--detection_model', type=str, default='retinaface_resnet50', 
  74.             help='Face detector. Optional: retinaface_resnet50, retinaface_mobile0.25, YOLOv5l, YOLOv5n, dlib. \
  75.                 Default: retinaface_resnet50')
  76.     parser.add_argument('--bg_upsampler', type=str, default='None', help='Background upsampler. Optional: realesrgan')
  77.     parser.add_argument('--face_upsample', action='store_true', help='Face upsampler after enhancement. Default: False')
  78.     parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')
  79.     parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces. Default: None')
  80.     parser.add_argument('--save_video_fps', type=float, default=None, help='Frame rate for saving video. Default: None')
  81.  
  82.     args = parser.parse_args()
  83.  
  84.     # ------------------------ input & output ------------------------
  85.     w = args.fidelity_weight
  86.     input_video = False
  87.     if args.input_path.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')): # input single img path
  88.         input_img_list = [args.input_path]
  89.         result_root = f'results/test_img_{w}'
  90.     elif args.input_path.endswith(('mp4', 'mov', 'avi', 'MP4', 'MOV', 'AVI')): # input video path
  91.         from basicsr.utils.video_util import VideoReader, VideoWriter
  92.         input_img_list = []
  93.         vidreader = VideoReader(args.input_path)
  94.         image = vidreader.get_frame()
  95.         while image is not None:
  96.             input_img_list.append(image)
  97.             image = vidreader.get_frame()
  98.         audio = vidreader.get_audio()
  99.         fps = vidreader.get_fps() if args.save_video_fps is None else args.save_video_fps   
  100.         video_name = os.path.basename(args.input_path)[:-4]
  101.         result_root = f'results/{video_name}_{w}'
  102.         input_video = True
  103.         vidreader.close()
  104.     else: # input img folder
  105.         if args.input_path.endswith('/'):  # solve when path ends with /
  106.             args.input_path = args.input_path[:-1]
  107.         # scan all the jpg and png images
  108.         input_img_list = sorted(glob.glob(os.path.join(args.input_path, '*.[jpJP][pnPN]*[gG]')))
  109.         result_root = f'results/{os.path.basename(args.input_path)}_{w}'
  110.  
  111.     if not args.output_path is None: # set output path
  112.         result_root = args.output_path
  113.  
  114.     test_img_num = len(input_img_list)
  115.     if test_img_num == 0:
  116.         raise FileNotFoundError('No input image/video is found...\n' 
  117.             '\tNote that --input_path for video should end with .mp4|.mov|.avi')
  118.  
  119.     # ------------------ set up background upsampler ------------------
  120.     if args.bg_upsampler == 'realesrgan':
  121.         bg_upsampler = set_realesrgan()
  122.     else:
  123.         bg_upsampler = None
  124.  
  125.     # ------------------ set up face upsampler ------------------
  126.     if args.face_upsample:
  127.         if bg_upsampler is not None:
  128.             face_upsampler = bg_upsampler
  129.         else:
  130.             face_upsampler = set_realesrgan()
  131.     else:
  132.         face_upsampler = None
  133.  
  134.     # ------------------ set up CodeFormer restorer -------------------
  135.     net = ARCH_REGISTRY.get('CodeFormer')(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, 
  136.                                             connect_list=['32', '64', '128', '256']).to(device)
  137.     
  138.     # ckpt_path = 'weights/CodeFormer/codeformer.pth'
  139.     ckpt_path = load_file_from_url(url=pretrain_model_url['restoration'], 
  140.                                     model_dir='weights/CodeFormer', progress=True, file_name=None)
  141.     checkpoint = torch.load(ckpt_path)['params_ema']
  142.     net.load_state_dict(checkpoint)
  143.     net.eval()
  144.  
  145.     # ------------------ set up FaceRestoreHelper -------------------
  146.     # large det_model: 'YOLOv5l', 'retinaface_resnet50'
  147.     # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
  148.     if not args.has_aligned: 
  149.         print(f'Face detection model: {args.detection_model}')
  150.     if bg_upsampler is not None: 
  151.         print(f'Background upsampling: True, Face upsampling: {args.face_upsample}')
  152.     else:
  153.         print(f'Background upsampling: False, Face upsampling: {args.face_upsample}')
  154.  
  155.     face_helper = FaceRestoreHelper(
  156.         args.upscale,
  157.         face_size=512,
  158.         crop_ratio=(1, 1),
  159.         det_model = args.detection_model,
  160.         save_ext='png',
  161.         use_parse=True,
  162.         device=device)
  163.  
  164.     # -------------------- start to processing ---------------------
  165.     for i, img_path in enumerate(input_img_list):
  166.         # clean all the intermediate results to process the next image
  167.         face_helper.clean_all()
  168.         
  169.         if isinstance(img_path, str):
  170.             img_name = os.path.basename(img_path)
  171.             basename, ext = os.path.splitext(img_name)
  172.             print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
  173.             img = cv2.imread(img_path, cv2.IMREAD_COLOR)
  174.         else: # for video processing
  175.             basename = str(i).zfill(6)
  176.             img_name = f'{video_name}_{basename}' if input_video else basename
  177.             print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
  178.             img = img_path
  179.  
  180.         if args.has_aligned: 
  181.             # the input faces are already cropped and aligned
  182.             img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
  183.             face_helper.is_gray = is_gray(img, threshold=10)
  184.             if face_helper.is_gray:
  185.                 print('Grayscale input: True')
  186.             face_helper.cropped_faces = [img]
  187.         else:
  188.             face_helper.read_image(img)
  189.             # get face landmarks for each face
  190.             num_det_faces = face_helper.get_face_landmarks_5(
  191.                 only_center_face=args.only_center_face, resize=640, eye_dist_threshold=5)
  192.             print(f'\tdetect {num_det_faces} faces')
  193.             # align and warp each face
  194.             face_helper.align_warp_face()
  195.  
  196.         # face restoration for each cropped face
  197.         for idx, cropped_face in enumerate(face_helper.cropped_faces):
  198.             # prepare data
  199.             cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
  200.             normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
  201.             cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
  202.  
  203.             try:
  204.                 with torch.no_grad():
  205.                     output = net(cropped_face_t, w=w, adain=True)[0]
  206.                     restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
  207.                 del output
  208.                 torch.cuda.empty_cache()
  209.             except Exception as error:
  210.                 print(f'\tFailed inference for CodeFormer: {error}')
  211.                 restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
  212.  
  213.             restored_face = restored_face.astype('uint8')
  214.             face_helper.add_restored_face(restored_face, cropped_face)
  215.  
  216.         # paste_back
  217.         if not args.has_aligned:
  218.             # upsample the background
  219.             if bg_upsampler is not None:
  220.                 # Now only support RealESRGAN for upsampling background
  221.                 bg_img = bg_upsampler.enhance(img, outscale=args.upscale)[0]
  222.             else:
  223.                 bg_img = None
  224.             face_helper.get_inverse_affine(None)
  225.             # paste each restored face to the input image
  226.             if args.face_upsample and face_upsampler is not None: 
  227.                 restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box, face_upsampler=face_upsampler)
  228.             else:
  229.                 restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box)
  230.  
  231.         # save faces
  232.         for idx, (cropped_face, restored_face) in enumerate(zip(face_helper.cropped_faces, face_helper.restored_faces)):
  233.             # save cropped face
  234.             if not args.has_aligned: 
  235.                 save_crop_path = os.path.join(result_root, 'cropped_faces', f'{basename}_{idx:02d}.png')
  236.                 imwrite(cropped_face, save_crop_path)
  237.             # save restored face
  238.             if args.has_aligned:
  239.                 save_face_name = f'{basename}.png'
  240.             else:
  241.                 save_face_name = f'{basename}_{idx:02d}.png'
  242.             if args.suffix is not None:
  243.                 save_face_name = f'{save_face_name[:-4]}_{args.suffix}.png'
  244.             save_restore_path = os.path.join(result_root, 'restored_faces', save_face_name)
  245.             imwrite(restored_face, save_restore_path)
  246.  
  247.         # save restored img
  248.         if not args.has_aligned and restored_img is not None:
  249.             if args.suffix is not None:
  250.                 basename = f'{basename}_{args.suffix}'
  251.             save_restore_path = os.path.join(result_root, 'final_results', f'{basename}.png')
  252.             imwrite(restored_img, save_restore_path)
  253.  
  254.     # save enhanced video
  255.     if input_video:
  256.         print('Video Saving...')
  257.         # load images
  258.         video_frames = []
  259.         img_list = sorted(glob.glob(os.path.join(result_root, 'final_results', '*.[jp][pn]g')))
  260.         for img_path in img_list:
  261.             img = cv2.imread(img_path)
  262.             video_frames.append(img)
  263.         # write images to video
  264.         height, width = video_frames[0].shape[:2]
  265.         if args.suffix is not None:
  266.             video_name = f'{video_name}_{args.suffix}.png'
  267.         save_restore_path = os.path.join(result_root, f'{video_name}.mp4')
  268.         vidwriter = VideoWriter(save_restore_path, height, width, fps, audio)
  269.          
  270.         for f in video_frames:
  271.             vidwriter.write_frame(f)
  272.         vidwriter.close()
  273.  
  274.     print(f'\nAll results are saved in {result_root}')

test.jpg是我自己网上下载的测试图片，终端执行下面的命令：

python inference_codeformer.py -w 0.5 --has_aligned --input_path test.jpg

原图如下：

 结果图片如下所示：

 感觉有点变形了，不知道是不是图像的问题。

接下来看下灰度图的效果，原图如下所示：

 结果图片如下所示：

 感觉变高清了很多很多，我以为会变成彩色的，结果并没有。

接下来看下画笔涂抹严重的脸能修复成什么样，终端执行：

python inference_inpainting.py --input_path test.jpg

原图如下所示：

 结果图片如下所示：

 图像输入：

 修复输出:

 看起来效果还不错。

当然了，CodeFormer也可以实现对面部上色，终端执行：

python inference_colorization.py --input_path test.jpg

结果对比显示如下所示：

 当然了，如果不喜欢终端命令行的形式，官方同样提供了基于Gradio模块开发的APP，如下所示：

  1. """
  2. This file is used for deploying hugging face demo:
  3. https://huggingface.co/spaces/sczhou/CodeFormer
  4. """
  5.  
  6. import sys
  7. sys.path.append('CodeFormer')
  8. import os
  9. import cv2
  10. import torch
  11. import torch.nn.functional as F
  12. import gradio as gr
  13.  
  14. from torchvision.transforms.functional import normalize
  15.  
  16. from basicsr.archs.rrdbnet_arch import RRDBNet
  17. from basicsr.utils import imwrite, img2tensor, tensor2img
  18. from basicsr.utils.download_util import load_file_from_url
  19. from basicsr.utils.misc import gpu_is_available, get_device
  20. from basicsr.utils.realesrgan_utils import RealESRGANer
  21. from basicsr.utils.registry import ARCH_REGISTRY
  22.  
  23. from facelib.utils.face_restoration_helper import FaceRestoreHelper
  24. from facelib.utils.misc import is_gray
  25.  
  26.  
  27. os.system("pip freeze")
  28.  
  29. pretrain_model_url = {
  30.     'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
  31.     'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
  32.     'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',
  33.     'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'
  34. }
  35. # download weights
  36. if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'):
  37.     load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None)
  38. if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'):
  39.     load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
  40. if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'):
  41.     load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
  42. if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'):
  43.     load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None)
  44.  
  45. # download images
  46. torch.hub.download_url_to_file(
  47.     'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png',
  48.     '01.png')
  49. torch.hub.download_url_to_file(
  50.     'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg',
  51.     '02.jpg')
  52. torch.hub.download_url_to_file(
  53.     'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg',
  54.     '03.jpg')
  55. torch.hub.download_url_to_file(
  56.     'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg',
  57.     '04.jpg')
  58. torch.hub.download_url_to_file(
  59.     'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',
  60.     '05.jpg')
  61.  
  62. def imread(img_path):
  63.     img = cv2.imread(img_path)
  64.     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  65.     return img
  66.  
  67. # set enhancer with RealESRGAN
  68. def set_realesrgan():
  69.     # half = True if torch.cuda.is_available() else False
  70.     half = True if gpu_is_available() else False
  71.     model = RRDBNet(
  72.         num_in_ch=3,
  73.         num_out_ch=3,
  74.         num_feat=64,
  75.         num_block=23,
  76.         num_grow_ch=32,
  77.         scale=2,
  78.     )
  79.     upsampler = RealESRGANer(
  80.         scale=2,
  81.         model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",
  82.         model=model,
  83.         tile=400,
  84.         tile_pad=40,
  85.         pre_pad=0,
  86.         half=half,
  87.     )
  88.     return upsampler
  89.  
  90. upsampler = set_realesrgan()
  91. # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  92. device = get_device()
  93. codeformer_net = ARCH_REGISTRY.get("CodeFormer")(
  94.     dim_embd=512,
  95.     codebook_size=1024,
  96.     n_head=8,
  97.     n_layers=9,
  98.     connect_list=["32", "64", "128", "256"],
  99. ).to(device)
  100. ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth"
  101. checkpoint = torch.load(ckpt_path)["params_ema"]
  102. codeformer_net.load_state_dict(checkpoint)
  103. codeformer_net.eval()
  104.  
  105. os.makedirs('output', exist_ok=True)
  106.  
  107. def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
  108.     """Run a single prediction on the model"""
  109.     try: # global try
  110.         # take the default setting for the demo
  111.         has_aligned = False
  112.         only_center_face = False
  113.         draw_box = False
  114.         detection_model = "retinaface_resnet50"
  115.         print('Inp:', image, background_enhance, face_upsample, upscale, codeformer_fidelity)
  116.  
  117.         img = cv2.imread(str(image), cv2.IMREAD_COLOR)
  118.         print('\timage size:', img.shape)
  119.  
  120.         upscale = int(upscale) # convert type to int
  121.         if upscale > 4: # avoid memory exceeded due to too large upscale
  122.             upscale = 4 
  123.         if upscale > 2 and max(img.shape[:2])>1000: # avoid memory exceeded due to too large img resolution
  124.             upscale = 2 
  125.         if max(img.shape[:2]) > 1500: # avoid memory exceeded due to too large img resolution
  126.             upscale = 1
  127.             background_enhance = False
  128.             face_upsample = False
  129.  
  130.         face_helper = FaceRestoreHelper(
  131.             upscale,
  132.             face_size=512,
  133.             crop_ratio=(1, 1),
  134.             det_model=detection_model,
  135.             save_ext="png",
  136.             use_parse=True,
  137.             device=device,
  138.         )
  139.         bg_upsampler = upsampler if background_enhance else None
  140.         face_upsampler = upsampler if face_upsample else None
  141.  
  142.         if has_aligned:
  143.             # the input faces are already cropped and aligned
  144.             img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
  145.             face_helper.is_gray = is_gray(img, threshold=5)
  146.             if face_helper.is_gray:
  147.                 print('\tgrayscale input: True')
  148.             face_helper.cropped_faces = [img]
  149.         else:
  150.             face_helper.read_image(img)
  151.             # get face landmarks for each face
  152.             num_det_faces = face_helper.get_face_landmarks_5(
  153.             only_center_face=only_center_face, resize=640, eye_dist_threshold=5
  154.             )
  155.             print(f'\tdetect {num_det_faces} faces')
  156.             # align and warp each face
  157.             face_helper.align_warp_face()
  158.  
  159.         # face restoration for each cropped face
  160.         for idx, cropped_face in enumerate(face_helper.cropped_faces):
  161.             # prepare data
  162.             cropped_face_t = img2tensor(
  163.                 cropped_face / 255.0, bgr2rgb=True, float32=True
  164.             )
  165.             normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
  166.             cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
  167.  
  168.             try:
  169.                 with torch.no_grad():
  170.                     output = codeformer_net(
  171.                         cropped_face_t, w=codeformer_fidelity, adain=True
  172.                     )[0]
  173.                     restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
  174.                 del output
  175.                 torch.cuda.empty_cache()
  176.             except RuntimeError as error:
  177.                 print(f"Failed inference for CodeFormer: {error}")
  178.                 restored_face = tensor2img(
  179.                     cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
  180.                 )
  181.  
  182.             restored_face = restored_face.astype("uint8")
  183.             face_helper.add_restored_face(restored_face)
  184.  
  185.         # paste_back
  186.         if not has_aligned:
  187.             # upsample the background
  188.             if bg_upsampler is not None:
  189.                 # Now only support RealESRGAN for upsampling background
  190.                 bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
  191.             else:
  192.                 bg_img = None
  193.             face_helper.get_inverse_affine(None)
  194.             # paste each restored face to the input image
  195.             if face_upsample and face_upsampler is not None:
  196.                 restored_img = face_helper.paste_faces_to_input_image(
  197.                     upsample_img=bg_img,
  198.                     draw_box=draw_box,
  199.                     face_upsampler=face_upsampler,
  200.                 )
  201.             else:
  202.                 restored_img = face_helper.paste_faces_to_input_image(
  203.                     upsample_img=bg_img, draw_box=draw_box
  204.                 )
  205.  
  206.         # save restored img
  207.         save_path = f'output/out.png'
  208.         imwrite(restored_img, str(save_path))
  209.  
  210.         restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
  211.         return restored_img, save_path
  212.     except Exception as error:
  213.         print('Global exception', error)
  214.         return None, None
  215.  
  216.  
  217. title = "CodeFormer: Robust Face Restoration and Enhancement Network"
  218. description = r"""<center><img src='https://user-images.githubusercontent.com/14334509/189166076-94bb2cac-4f4e-40fb-a69f-66709e3d98f5.png' alt='CodeFormer logo'></center>
  219. <b>Official Gradio demo</b> for <a href='https://github.com/sczhou/CodeFormer' target='_blank'><b>Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)</b></a>.<br>
  220. 声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小蓝xlanll/article/detail/131772
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。