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

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'], '*')))
    for idx, path in enumerate(image_paths):
        imgname, imgsuffix = os.path.splitext(os.path.basename(path))
        LOGGER.info(colorstr('red', f"Processing image {idx + 1}/{len(image_paths)}: {os.path.basename(path)}"))
        img = cv2.imread(path, cv2.IMREAD_COLOR)
        img_float = img.astype(np.float32) / 255.0
        img_tensor = torch.from_numpy(np.transpose(img_float[:, :, [2, 1, 0]], (2, 0, 1))).float().unsqueeze(0).to(device)
        try:
            with torch.no_grad():
                output = model(img_tensor)
                output = output.data.squeeze().float().cpu().clamp(0, 1).numpy()
                output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
                output = (output * 255.0).round().astype(np.uint8)
                if config['open_face_enhance']:
                    _, _, output = face_enhancer.enhance(output, has_aligned=False, only_center_face=False, paste_back=True)
        except:
            LOGGER.info(colorstr('red', 'Error processing image!'))
            continue
        cv2.imwrite(os.path.join(config['output_folder'], f"{imgname}_eesrgan_x{config['magnification']}{gfpgan_suffix}{imgsuffix}"), output)
        if config['display_image']:
            cv2.imshow("Super Resolution Result", cv2.resize(output, None, fx=config['zoom_ratio'], fy=config['zoom_ratio'], interpolation=cv2.INTER_LINEAR))
            if cv2.waitKey(1) & 0xFF == 27:
                break
    if config['display_image']:
        cv2.destroyAllWindows()
    LOGGER.info(colorstr('green', '超分完成.'))


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("人脸图像超分辨率重建操作界面(Powered By Basicsr And Anperlanch)")
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout()
        form_layout = QFormLayout()

        self.input_folder_lineedit = QLineEdit()
        self.input_folder_lineedit.setText("C:/Users/15231/Desktop/detection_results/noHelmet")
        self.input_folder_button = QPushButton("浏览")
        self.input_folder_button.clicked.connect(self.browse_input_folder)
        h_input = QHBoxLayout()
        h_input.addWidget(self.input_folder_lineedit)
        h_input.addWidget(self.input_folder_button)
        form_layout.addRow("低分辨率图片文件夹路径:", h_input)

        self.output_folder_lineedit = QLineEdit()
        self.output_folder_lineedit.setText("C:/Users/15231/Desktop/super_resolution_results")
        self.output_folder_button = QPushButton("浏览")
        self.output_folder_button.clicked.connect(self.browse_output_folder)
        h_output = QHBoxLayout()
        h_output.addWidget(self.output_folder_lineedit)
        h_output.addWidget(self.output_folder_button)
        form_layout.addRow("超分后图片文件夹路径:", h_output)

        self.device_combobox = QComboBox()
        self.device_combobox.addItems(["GPU", "CPU"])
        form_layout.addRow("超分运行设备:", self.device_combobox)

        self.mag2_radio = QRadioButton("x2")
        self.mag4_radio = QRadioButton("x4")
        self.mag4_radio.setChecked(True)
        h_mag = QHBoxLayout()
        h_mag.addWidget(self.mag2_radio)
        h_mag.addWidget(self.mag4_radio)
        form_layout.addRow("超分倍率:", h_mag)

        self.face_enhance_checkbox = QCheckBox()
        self.face_enhance_checkbox.setChecked(True)
        form_layout.addRow("是否启用GPFGAN进行面部增强:", self.face_enhance_checkbox)

        self.gfpgan_magnification_spinbox = QSpinBox()
        self.gfpgan_magnification_spinbox.setRange(1, 4)
        self.gfpgan_magnification_spinbox.setValue(1)
        form_layout.addRow("GFPGAN放大倍数:", self.gfpgan_magnification_spinbox)

        self.model_path_x4_lineedit = QLineEdit()
        self.model_path_x4_lineedit.setText("D:/personalFiles/fsr/Improve_esrgan/experiments/train_RRDBNet_PSNR_x4_miniCeleba/models/net_g_32000.pth")
        self.model_path_x4_button = QPushButton("浏览")
        self.model_path_x4_button.clicked.connect(self.browse_model_path_x4)
        h_model4 = QHBoxLayout()
        h_model4.addWidget(self.model_path_x4_lineedit)
        h_model4.addWidget(self.model_path_x4_button)
        form_layout.addRow("ENHANCED-ESRGAN模型文件路径 (x4):", h_model4)

        self.model_path_x2_lineedit = QLineEdit()
        self.model_path_x2_lineedit.setText("")
        self.model_path_x2_button = QPushButton("浏览")
        self.model_path_x2_button.clicked.connect(self.browse_model_path_x2)
        h_model2 = QHBoxLayout()
        h_model2.addWidget(self.model_path_x2_lineedit)
        h_model2.addWidget(self.model_path_x2_button)
        form_layout.addRow("ENHANCED-ESRGAN模型文件路径 (x2):", h_model2)

        self.display_checkbox = QCheckBox()
        self.display_checkbox.setChecked(True)
        form_layout.addRow("显示超分后图片:", self.display_checkbox)

        self.zoom_ratio_spinbox = QDoubleSpinBox()
        self.zoom_ratio_spinbox.setRange(1.0, 3.0)
        self.zoom_ratio_spinbox.setSingleStep(0.1)
        self.zoom_ratio_spinbox.setValue(2.0)
        form_layout.addRow("显示图像缩放比:", self.zoom_ratio_spinbox)

        layout.addLayout(form_layout)
        self.start_button = QPushButton("超分辨率 启动")
        self.start_button.clicked.connect(self.start_super_resolution)
        layout.addWidget(self.start_button)
        central_widget.setLayout(layout)
        self.resize(870, 400)

    def browse_input_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择低分辨率图片文件夹")
        if folder:
            self.input_folder_lineedit.setText(folder)

    def browse_output_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择超分后输出图片文件夹")
        if folder:
            self.output_folder_lineedit.setText(folder)

    def browse_model_path_x4(self):
        filename, _ = QFileDialog.getOpenFileName(self, "选择x4模型文件", "", "Model Files (*.pth)")
        if filename:
            self.model_path_x4_lineedit.setText(filename)

    def browse_model_path_x2(self):
        filename, _ = QFileDialog.getOpenFileName(self, "选择x2模型文件", "", "Model Files (*.pth)")
        if filename:
            self.model_path_x2_lineedit.setText(filename)

    def start_super_resolution(self):
        config = {
            'input_folder': self.input_folder_lineedit.text(),
            'output_folder': self.output_folder_lineedit.text(),
            'device': self.device_combobox.currentText(),
            'magnification': 4 if self.mag4_radio.isChecked() else 2,
            'open_face_enhance': self.face_enhance_checkbox.isChecked(),
            'gfpgan_magnification': self.gfpgan_magnification_spinbox.value(),
            'model_path_x4': self.model_path_x4_lineedit.text(),
            'model_path_x2': self.model_path_x2_lineedit.text(),
            'display_image': self.display_checkbox.isChecked(),
            'zoom_ratio': self.zoom_ratio_spinbox.value()
        }
        LOGGER.info(colorstr('blue', '开始人脸图像超分辨率重建.') + colorstr('red', '（已启用GFPGAN人脸图像增强算法）' if self.face_enhance_checkbox.isChecked() else ''))
        threading.Thread(target=run_super_resolution, args=(config,)).start()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())