AI 视频超分与人脸修复利器：RealESRGAN + GFPGAN 高效处理脚本

🚀 前言

视频超分辨率（Video Super-Resolution, VSR）是提升旧视频或低分辨率视频质量的有效手段。本脚本结合了 RealESRGAN (用于通用背景和细节放大) 和 GFPGAN (用于高清人脸修复) 两个强大的 AI 模型，并针对大规模视频处理中的常见问题（如中断、资源占用、I/O 瓶颈）进行了深度优化。

特别是 V16.0 版本，引入了线程安全的断点续传机制和优化的 I/O 线程池，彻底解决了在处理数万帧视频时容易出现的假死或进度丢失问题。

🛠️ 环境搭建与配置要求

1. 软件环境

您需要安装以下基础软件：

Python 3.8+
CUDA (NVIDIA 显卡用户)：强烈推荐使用 GPU 进行推理，否则速度会非常慢。
FFmpeg：用于视频的读取、音频提取和最终的合成。请确保 ffmpeg 命令在您的系统 PATH 中可用。

2. Python 依赖

通过 pip 安装所有必需的库：

pip install torch torchvision torchaudio
# Real-ESRGAN/GFPGAN 依赖
pip install basicsr
pip install facexlib
pip install gfpgan
pip install opencv-python
pip install tqdm
# 如果您需要手动安装 Real-ESRGAN 和 GFPGAN，请参考官方仓库。

3. 模型文件准备

请下载以下两个预训练模型，并将其放置到您指定的路径（例如 /root/autodl-tmp/models/）：

RealESRGAN 模型 (背景/细节超分):
- **文件**: RealESRGAN_x4plus.pth
- **下载地址**: [官方仓库通常会提供下载链接]
GFPGAN 模型 (人脸修复):
- **文件**: GFPGANv1.4.pth
- **下载地址**: [官方仓库通常会提供下载链接]

📋 代码配置（重要）

在使用前，请务必修改代码开头的四个核心配置项：

编号	变量名	描述	示例值
1	INPUT_VIDEO_PATH	待处理的原始视频的绝对路径。	`"/root/autodl-tmp/video.mp4"`
2	OUTPUT_DIR	临时帧和最终视频的输出目录。	`"/root/autodl-tmp/upscale_output/"`
3	REALESRGAN_MODEL_PATH	步骤 3 下载的 RealESRGAN 模型路径。	`"/root/autodl-tmp/models/RealESRGAN_x4plus.pth"`
4	GFPGAN_MODEL_PATH	步骤 3 下载的 GFPGAN 模型路径。	`"/root/autodl-tmp/models/GFPGANv1.4.pth"`

性能调优参数

变量名	描述	建议值/说明
TILE_SIZE	RealESRGAN 的切块大小。显存越大，可以设置越大。减少显存占用时可调小 (如 4000)。	`6000` (大显存友好)
BATCH_SIZE	每次从视频中读取的帧数。不影响 GPU 推理批次，主要影响 I/O 效率。	`8`
IO_THREADS	异步写入图片的线程数。关键优化点。过多会导致磁盘 I/O 队列阻塞，建议 `4-8`。	`6`
VIDEO_CRF	FFmpeg 视频编码质量，值越小质量越高，文件越大。`18` 是一个高质量的平衡点。	`'18'`

💡 注意事项

断点续传（Checkpointing）：脚本会在 OUTPUT_DIR 中生成一个 checkpoint.txt 文件，记录已完成的帧数。如果程序意外中断（如断电或手动停止），下次运行时将自动从上次中断的位置继续处理，无需修改代码。
线程锁 (CHECKPOINT_LOCK)：这是 V16.0 修复的关键。它确保多线程异步写入图片时，对进度文件 (checkpoint.txt) 的操作是原子性的，彻底避免了文件损坏或多线程争抢导致的程序假死。
单文件原则：脚本将视频分解为图片序列 (frame_000001.jpg, frame_000002.jpg 等) 进行处理，最后再合成视频。请确保 OUTPUT_DIR 有足够的空间（通常是原视频体积的数倍到数十倍）。

💻 完整 Python 代码

以下是经过注释优化和 V16.0 修复的完整脚本，您可以直接保存为 upscale_processor.py。

import cv2
import os
import subprocess
import time
from tqdm import tqdm
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gfpgan import GFPGANer
import torch
import numpy as np
import concurrent.futures
import threading

# -------------------------------------------------------------
# 🎯 配置参数 (V16.0 终极修复版: 解决最后阶段假死问题)
# -------------------------------------------------------------
# *** 1. 替换为您的测试视频路径 (MP4 格式) ***
INPUT_VIDEO_PATH = "/root/autodl-tmp/video.mp4"

# *** 2. 输出文件夹路径 ***
OUTPUT_DIR = "/root/autodl-tmp/upscale_output/"
OUTPUT_VIDEO_NAME = "upscaled_video_1080p.mp4"

# *** 3. 模型路径 ***
REALESRGAN_MODEL_PATH = "/root/autodl-tmp/models/RealESRGAN_x4plus.pth"
GFPGAN_MODEL_PATH = "/root/autodl-tmp/models/GFPGANv1.4.pth"

# -------------------------------------------------------------
# AI/处理参数
# -------------------------------------------------------------
CHECKPOINT_FILE = os.path.join(OUTPUT_DIR, "checkpoint.txt")  # 进度记录文件路径
UPSCALING_FACTOR = 4  # 放大倍数，RealESRGAN x4plus 对应 4

# 目标分辨率 (1080P)
TARGET_WIDTH = 1920
TARGET_HEIGHT = 1080

# 显存优化参数
TILE_SIZE = 6000  # RealESRGAN 切块大小，用于节省显存
BATCH_SIZE = 8  # 视频读取批处理大小

# 图片输出参数
OUTPUT_FRAME_EXTENSION = ".jpg"
JPEG_QUALITY = 95 # 输出 JPEG 质量 (0-100)，95 保证高质量

# [V16.0 关键修复] I/O 线程数
# 用于异步写入已处理的帧，避免 GPU 等待 I/O。
# 之前 16 线程会导致磁盘队列阻塞，改为 6 线程更稳定
IO_THREADS = 6

# 视频编码质量
VIDEO_CRF = '18' # Constant Rate Factor, 越小质量越高 (18-23 常用)

# [V16.0 关键修复] 全局线程锁
# 确保同一时间只有一个线程在写 checkpoint.txt，彻底解决多线程写文件导致的冲突。
CHECKPOINT_LOCK = threading.Lock()


# -------------------------------------------------------------
# 🛠️ 初始化 RealESRGAN 和 GFPGAN 引擎
# -------------------------------------------------------------

def initialize_upsampler():
    """初始化 Real-ESRGAN 和 GFPGAN 引擎，并将其部署到 GPU (如果可用)。"""
    print("INFO: Initializing Real-ESRGAN and GFPGAN engines...")

    # 检查模型文件是否存在
    if not os.path.exists(REALESRGAN_MODEL_PATH) or not os.path.exists(GFPGAN_MODEL_PATH):
        print(f"FATAL ERROR: Model files not found. Please check paths.")
        return None

    # 自动选择设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if device == 'cpu':
        print("WARNING: CUDA not found. Running on CPU, which will be extremely slow.")

    # 1. 初始化 Real-ESRGAN (背景超分)
    print("INFO: Initializing Real-ESRGAN core model...")
    realesrgan_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23)

    bg_upsampler = RealESRGANer(
        scale=UPSCALING_FACTOR,
        model_path=REALESRGAN_MODEL_PATH,
        model=realesrgan_model,
        device=device,
        tile=TILE_SIZE,  # 使用切块处理大图
        tile_pad=10,
    )

    # 2. 初始化 GFPGAN (人脸修复)，并将 RealESRGAN 作为其背景上采样器
    print("INFO: Initializing GFPGAN face enhancer...")
    face_enhancer = GFPGANer(
        model_path=GFPGAN_MODEL_PATH,
        upscale=UPSCALING_FACTOR,
        arch='clean',
        bg_upsampler=bg_upsampler,
        device=device,
    )

    print("INFO: Engines initialized. Device:", face_enhancer.device)
    return face_enhancer


# -------------------------------------------------------------
# [V16.0 修复] 线程安全的异步 I/O 写入
# -------------------------------------------------------------
def _write_frame_and_checkpoint(frame_idx, output_frame, output_path, quality):
    """
    执行文件写入和进度更新。
    此函数在 I/O 线程池中异步运行。
    """
    try:
        # 1. 写入图片 (耗时操作，并发执行)
        cv2.imwrite(
            output_path,
            output_frame,
            [int(cv2.IMWRITE_JPEG_QUALITY), quality]
        )

        # 2. 写入进度 (加锁执行)
        # 只有获得锁的线程才能写文件，防止多线程同时占用文件句柄导致程序挂起
        with CHECKPOINT_LOCK:
            with open(CHECKPOINT_FILE, 'w') as f:
                # 写入下一个待处理的帧索引 (即已完成的帧数)
                f.write(str(frame_idx + 1))

    except Exception as e:
        print(f"\nERROR in I/O worker for frame {frame_idx}: Write failed: {e}")


# -------------------------------------------------------------
# 核心函数：视频处理 (包含断点续传)
# -------------------------------------------------------------

def process_video(upsampler):
    """从视频中读取帧，批量进行 GPU 推理，并异步写入到磁盘。"""

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    cap = cv2.VideoCapture(INPUT_VIDEO_PATH)
    if not cap.isOpened():
        print(f"ERROR: Cannot open video file: {INPUT_VIDEO_PATH}")
        return False, None

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    if total_frames <= 0 or fps <= 0:
        print("ERROR: Cannot retrieve frame count. Check video integrity.")
        return False, None

    # --- Checkpoint/Resume Logic ---
    start_frame = 0
    if os.path.exists(CHECKPOINT_FILE):
        try:
            with open(CHECKPOINT_FILE, 'r') as f:
                content = f.read().strip()
                if content:
                    start_frame_from_file = int(content)
                    
                    # 检查是否已全部完成
                    if start_frame_from_file >= total_frames:
                        print("INFO: All frames already processed based on checkpoint.")
                        cap.release()
                        # 清理 checkpoint 文件
                        if os.path.exists(CHECKPOINT_FILE):
                            os.remove(CHECKPOINT_FILE)
                        return True, fps
                        
                    start_frame = start_frame_from_file
                    print(f"INFO: Checkpoint found. Resuming from frame {start_frame}...")
        except ValueError:
            print("WARNING: Checkpoint corrupted. Starting from 0.")

    # 设置视频读取的起始位置
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    current_frame_idx = start_frame

    # 使用 ThreadPoolExecutor 管理 I/O 线程
    with concurrent.futures.ThreadPoolExecutor(max_workers=IO_THREADS) as executor:
        futures = []  # 记录所有提交的 I/O 任务
        pbar = tqdm(total=total_frames, initial=start_frame, desc="Processing", ascii=True)

        while current_frame_idx < total_frames:
            frame_batch = []
            frame_indices = []
            
            # --- 1. 批量读取与跳过逻辑 ---
            frames_to_read = 0
            temp_idx = current_frame_idx
            first_read_idx = -1

            # 预扫描：确定哪些帧需要真正读取，哪些可以跳过(已存在)
            while frames_to_read < BATCH_SIZE and temp_idx < total_frames:
                output_frame_path = os.path.join(OUTPUT_DIR, f"frame_{temp_idx:06d}{OUTPUT_FRAME_EXTENSION}")

                if os.path.exists(output_frame_path) and os.path.getsize(output_frame_path) > 0:
                    # 文件存在且不为空，跳过，并更新进度条
                    temp_idx += 1
                    pbar.update(1)
                    continue
                else:
                    # 找到第一个需要处理的帧
                    if first_read_idx == -1:
                        first_read_idx = temp_idx
                    frames_to_read += 1
                    temp_idx += 1

            # 实际执行读取操作
            if first_read_idx != -1:
                # 如果跳过了一些帧，需要重新定位视频流
                if first_read_idx > current_frame_idx:
                    cap.set(cv2.CAP_PROP_POS_FRAMES, first_read_idx)
                
                current_frame_idx = first_read_idx # 更新指针到实际读取位置

                for _ in range(frames_to_read):
                    ret, frame = cap.read()
                    if not ret:
                        # 视频提前结束
                        break
                    frame_batch.append(frame)
                    frame_indices.append(current_frame_idx)
                    current_frame_idx += 1
            else:
                # 本批次所有帧都已存在(跳过了)
                if temp_idx >= total_frames:
                    break
                current_frame_idx = temp_idx
                if not frame_batch and current_frame_idx < total_frames:
                    # 即使没有读取新的帧，也要确保继续循环直到结束
                    continue
                elif current_frame_idx >= total_frames:
                    break

            # --- 2. GPU 顺序推理 ---
            for i, frame in enumerate(frame_batch):
                f_idx = frame_indices[i]
                output_frame_path_i = os.path.join(OUTPUT_DIR, f"frame_{f_idx:06d}{OUTPUT_FRAME_EXTENSION}")

                try:
                    # RealESRGAN + GFPGAN 联合处理
                    _, _, output_frame = upsampler.enhance(
                        frame, has_aligned=False, only_center_face=False, paste_back=True
                    )

                except Exception as e:
                    print(f"\nERROR: Processing failed for frame {f_idx}: {e}")
                    cap.release()
                    pbar.close()
                    return False, fps

                # --- 3. 提交异步 I/O 任务 ---
                # 将处理结果和进度更新提交给线程池
                future = executor.submit(
                    _write_frame_and_checkpoint,
                    f_idx,
                    output_frame,
                    output_frame_path_i,
                    JPEG_QUALITY
                )
                futures.append(future)
                pbar.update(1)

                # [内存优化] 定期清理已完成的任务引用，防止 futures 列表无限膨胀
                futures = [f for f in futures if not f.done()]

        # 循环结束
        pbar.close()
        print("\nINFO: GPU processing finished. Waiting for pending I/O writes...")
        
        # [V16.0 关键修复] 优雅关闭线程池
        # wait=True 会阻塞主线程直到所有 I/O 任务完成。
        executor.shutdown(wait=True)

    cap.release()
    
    # 再次检查并清理 checkpoint
    if os.path.exists(CHECKPOINT_FILE):
        try:
            os.remove(CHECKPOINT_FILE)
        except:
            pass
            
    print("\nINFO: All frames processed successfully.")
    return True, fps


# -------------------------------------------------------------
# 核心函数：视频合成（使用 FFmpeg）
# -------------------------------------------------------------

def combine_video(original_fps):
    """使用 FFmpeg 将高分辨率帧与原始音频合成为最终视频。"""

    image_sequence = os.path.join(OUTPUT_DIR, f"frame_%06d{OUTPUT_FRAME_EXTENSION}")
    final_video_path = os.path.join(OUTPUT_DIR, OUTPUT_VIDEO_NAME)
    temp_audio_path = os.path.join(OUTPUT_DIR, "temp_audio.aac")

    # 1. 提取音频
    print("INFO: Extracting audio...")
    # 尝试使用 FFmpeg 提取音频流，-vn 跳过视频流，-acodec copy 复制音频编码
    audio_command = [
        'ffmpeg', '-i', INPUT_VIDEO_PATH,
        '-vn', '-acodec', 'copy', temp_audio_path,
        '-y'
    ]

    has_audio = False
    try:
        # 使用 subprocess.run 执行命令并检查返回码
        subprocess.run(audio_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        has_audio = True
    except subprocess.CalledProcessError:
        print("WARNING: No audio stream found or extraction failed. Video will be silent.")
        has_audio = False

    # 2. 合并帧序列和音频
    print(f"INFO: Combining frames into {OUTPUT_VIDEO_NAME}...")
    ffmpeg_command = [
        'ffmpeg',
        '-r', str(original_fps), # 设置输入帧率
        '-i', image_sequence, # 输入图片序列
    ]

    # 如果音频提取成功，则添加音频输入
    if has_audio and os.path.exists(temp_audio_path):
        ffmpeg_command.extend(['-i', temp_audio_path])

    # 视频输出设置
    ffmpeg_command.extend([
        '-c:v', 'libx264',  # 编码器
        '-crf', VIDEO_CRF,  # 质量因子
        '-pix_fmt', 'yuv420p', # 像素格式，保证兼容性
        # 视频滤镜：缩放并居中到目标分辨率 (TARGET_WIDTH x TARGET_HEIGHT)
        '-vf', f'scale={TARGET_WIDTH}:{TARGET_HEIGHT}:force_original_aspect_ratio=decrease, pad={TARGET_WIDTH}:{TARGET_HEIGHT}:(ow-iw)/2:(oh-ih)/2',
        '-y', final_video_path
    ])

    start_time = time.time()
    try:
        subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        print(f"SUCCESS: Video saved to {final_video_path}")
        print(f"Time taken: {time.time() - start_time:.2f}s")
        
        # 清理临时音频文件
        if has_audio and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)

    except subprocess.CalledProcessError as e:
        print(f"FATAL ERROR: FFmpeg failed. Please check if FFmpeg is installed and accessible.")
        print(f"Error details: {e.stderr.decode()}")


# -------------------------------------------------------------
# 🏁 主程序入口
# -------------------------------------------------------------

if __name__ == "__main__":
    
    # 环境检查
    if not os.path.exists(INPUT_VIDEO_PATH):
        print(f"FATAL ERROR: Input video not found: {INPUT_VIDEO_PATH}")
    else:
        # 1. 初始化模型
        upsampler = initialize_upsampler()

        if upsampler:
            # 2. 处理帧
            success, fps = process_video(upsampler)

            # 3. 合成视频
            if success and fps:
                combine_video(fps)

    print("\n------------------------------------")
    print("🎬 Job Done.")
    print("------------------------------------")

🏃 使用步骤

保存代码: 将上述代码保存为 upscale_processor.py。
修改配置: 在代码开头修改 INPUT_VIDEO_PATH, OUTPUT_DIR, REALESRGAN_MODEL_PATH, 和 GFPGAN_MODEL_PATH 四个变量，确保路径正确。
运行脚本: 在命令行中执行：
```
python upscale_processor.py
```
查看结果: 程序将在 OUTPUT_DIR 目录下输出中间帧序列，并在处理完成后生成最终视频 upscaled_video_1080p.mp4。
处理中断: 如果程序中断，重新运行即可。它将自动检测 checkpoint.txt 并从上次中断的帧继续处理。

希望这个 V16.0 终极修复版能帮助您高效稳定地完成视频超分任务！如果您想调整人脸修复的力度或背景超分的细节，可以在代码中研究 RealESRGANer 和 GFPGANer 的更多参数（例如 fidelity_weight）。
项目代码和模型数据地址：

通过网盘分享的文件：VividRestore.zip
链接: https://pan.baidu.com/s/1baXHkciH1_8-mQoDE4enUg 提取码: t6gp 复制这段内容后打开百度网盘手机App，操作更方便哦

posted @ 2025-11-27 14:19 lvye001 阅读(93) 评论(0) 收藏举报

刷新页面返回顶部

lvye001