基于ESRGAN和GFPGAN实现人脸图像修复

最新推荐文章于 2025-11-29 16:24:14 发布

原创

最新推荐文章于 2025-11-29 16:24:14 发布 · 343 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python

1.GUI功能如下所示。

2.实现代码。

import sys
import os
import cv2
import glob
import numpy as np
import torch
import shutil
import threading
from basicsr.archs.rrdbnet_arch import RRDBNet
from ultralytics.utils import LOGGER, colorstr
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QFormLayout, QLineEdit, QPushButton, QFileDialog, QRadioButton, QCheckBox, QDoubleSpinBox, QSpinBox, QComboBox)
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="torchvision.models._utils")

def run_super_resolution(config):
    device = torch.device('cuda' if (config['device'] == 'GPU' and torch.cuda.is_available()) else 'cpu')
    if config['magnification'] == 4:
        model_path = config['model_path_x4']
    else:
        model_path = config['model_path_x2']
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32)
    model.load_state_dict(torch.load(model_path, map_location=device)['params'], strict=True)
    model.eval()
    model = model.to(device)
    if config['open_face_enhance']:
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth', upscale=config['gfpgan_magnification'], arch='clean', channel_multiplier=2)
        gfpgan_suffix = f'_gfpgan_x{config["gfpgan_magnification"]}'
    else:
        gfpgan_suffix = ''
    if os.path.exists(config["output_folder"]):
        shutil.rmtree(config["output_folder"])
    os.makedirs(config['output_folder'], exist_ok=True)
    image_paths = sorted(glob.glob(os.path.join(config['input_folder'], '*'))

最低0.47元/天解锁文章