video-subtitle-remover/backend/tools/video_io.py

import os
import queue
import subprocess
import threading

import cv2
import numpy as np

from .ffmpeg_cli import FFmpegCLI


class FramePrefetcher:
    """
    后台线程预解码视频帧，使 I/O 与模型推理重叠。
    接口兼容 cv2.VideoCapture（read/release）。
    """

    def __init__(self, video_cap, buffer_size=10):
        self.cap = video_cap
        self._buffer = queue.Queue(maxsize=buffer_size)
        self._stopped = False
        self._thread = threading.Thread(target=self._read_loop, daemon=True)
        self._thread.start()

    def _read_loop(self):
        while not self._stopped:
            ret, frame = self.cap.read()
            # 用带超时的 put，避免队列满时永久阻塞
            while not self._stopped:
                try:
                    self._buffer.put((ret, frame), timeout=0.5)
                    break
                except queue.Full:
                    pass
            if not ret:
                break

    def read(self):
        """读取下一帧，接口与 cv2.VideoCapture.read() 一致。"""
        return self._buffer.get()

    def get(self, propId):
        return self.cap.get(propId)

    def stop(self):
        """停止预读取，不释放底层 video_cap。"""
        self._stopped = True
        # 排空队列，解除后台线程可能阻塞的 put()
        while True:
            try:
                self._buffer.get_nowait()
            except queue.Empty:
                break
        self._thread.join(timeout=5)

    def release(self):
        self.stop()
        self.cap.release()


class FFmpegVideoWriter:
    """
    独立 I/O 线程写帧到 FFmpeg stdin，直接输出 H.264。
    主线程只往队列里放帧，永不阻塞在 I/O 上。
    """

    def __init__(self, output_path, fps, size):
        w, h = size
        self._queue = queue.Queue(maxsize=60)
        self._error = None
        # 优先用系统 FFmpeg + NVENC（GPU 编码），回退到 libx264
        import shutil as _shutil
        _sys_ffmpeg = '/usr/bin/ffmpeg'
        _nvenc_ok = (
            _shutil.which(_sys_ffmpeg) is not None and
            subprocess.run(
                [_sys_ffmpeg, '-hide_banner', '-encoders'],
                capture_output=True, text=True
            ).stdout.find('h264_nvenc') >= 0
        )
        ffmpeg_bin = _sys_ffmpeg if _nvenc_ok else FFmpegCLI.instance().ffmpeg_path
        codec_args = (
            ['-vf', 'format=yuv420p', '-vcodec', 'h264_nvenc', '-cq', '18', '-preset', 'p4']
            if _nvenc_ok else
            ['-vcodec', 'libx264', '-crf', '18', '-preset', 'ultrafast']
        )
        cmd = [
            ffmpeg_bin, '-y',
            '-f', 'rawvideo', '-vcodec', 'rawvideo',
            '-s', f'{w}x{h}', '-pix_fmt', 'bgr24',
            '-r', str(fps), '-i', 'pipe:0',
            *codec_args,
            '-loglevel', 'error', output_path,
        ]
        self._proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
        self._thread = threading.Thread(target=self._io_loop, daemon=True)
        self._thread.start()

    def _io_loop(self):
        try:
            while True:
                frame = self._queue.get()
                if frame is None:   # sentinel
                    break
                self._proc.stdin.write(frame.tobytes())
        except Exception as e:
            self._error = e
        finally:
            try:
                self._proc.stdin.close()
            except Exception:
                pass

    def write(self, frame):
        if frame.dtype != np.uint8:
            frame = np.clip(frame, 0, 255).astype(np.uint8)
        self._queue.put(frame)

    def release(self):
        self._queue.put(None)       # 通知 I/O 线程退出
        self._thread.join()         # 等待所有帧写完
        self._proc.wait()           # 等待 FFmpeg 完成编码
        if self._error:
            raise RuntimeError(f"FFmpegVideoWriter error: {self._error}")