king2025/gaoqqqqq

import cv2 import numpy as np import torch from basicsr.archs.rrdbnet_arch import RRDBNet from realesrgan import RealESRGANer from gfpgan import GFPGANer

class ImageRestorer: def init(self, device='cuda'): # 初始化设备 self.device = torch.device(device)

    # 加载超分辨率模型
    self.upsampler = RealESRGANer(
        scale=4,
        model_path='weights/RealESRGAN_x4plus.pth',
        model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32),
        tile=400,
        tile_pad=10,
        pre_pad=0,
        half=True
    )

    # 加载面部增强模型
    self.face_enhancer = GFPGANer(
        model_path='weights/GFPGANv1.4.pth',
        upscale=4,
        arch='clean',
        channel_multiplier=2,
        bg_upsampler=self.upsampler
    )

def restore_image(self, input_path, output_path, face_enhance=True):
    # 读取图像
    img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
    
    try:
        if face_enhance:
            # 执行面部增强
            _, _, output = self.face_enhancer.enhance(
                img,
                has_aligned=False,
                only_center_face=False,
                paste_back=True
            )
        else:
            # 普通超分辨率处理
            output, _ = self.upsampler.enhance(img, outscale=4)
        
        # 保存结果
        cv2.imwrite(output_path, output)
        print(f"图像已保存至 {output_path}")
        
    except Exception as e:
        print(f"处理出错: {e}")