mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-04-27 11:07:31 +08:00
6c61ddc6b0
- 字幕检测：每3帧采样一次OCR，中间帧插值填充，检测速度提升约3倍 - 视频编码：cv2.VideoWriter(mp4v) 替换为 FFmpeg libx264 管道编码，画质更好、体积更小 - 帧预读取：后台线程预解码视频帧，I/O 与模型推理重叠 - 消除 deepcopy：numpy 数组改用 .copy() 替代 copy.deepcopy，降低内存开销 - 清理冗余颜色空间转换中的 np.array() 包装 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
101 lines
2.6 KiB
Python
101 lines
2.6 KiB
Python
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()
|
||
self._buffer.put((ret, frame))
|
||
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
|
||
try:
|
||
while not self._buffer.empty():
|
||
self._buffer.get_nowait()
|
||
except queue.Empty:
|
||
pass
|
||
self._thread.join(timeout=5)
|
||
|
||
def release(self):
|
||
self.stop()
|
||
self.cap.release()
|
||
|
||
|
||
class FFmpegVideoWriter:
|
||
"""
|
||
通过 FFmpeg 管道写入帧，使用 libx264 编码。
|
||
接口兼容 cv2.VideoWriter（write/release）。
|
||
"""
|
||
|
||
def __init__(self, output_path, fps, size):
|
||
w, h = size
|
||
cmd = [
|
||
FFmpegCLI.instance().ffmpeg_path,
|
||
'-y',
|
||
'-f', 'rawvideo',
|
||
'-vcodec', 'rawvideo',
|
||
'-s', f'{w}x{h}',
|
||
'-pix_fmt', 'bgr24',
|
||
'-r', str(fps),
|
||
'-i', '-',
|
||
'-c:v', 'libx264',
|
||
'-pix_fmt', 'yuv420p',
|
||
'-crf', '18',
|
||
'-preset', 'fast',
|
||
'-loglevel', 'error',
|
||
output_path
|
||
]
|
||
self._process = subprocess.Popen(
|
||
cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.DEVNULL,
|
||
)
|
||
|
||
def write(self, frame):
|
||
"""写入一帧（numpy BGR 数组）。"""
|
||
if frame.dtype != np.uint8:
|
||
frame = np.clip(frame, 0, 255).astype(np.uint8)
|
||
try:
|
||
self._process.stdin.write(frame.tobytes())
|
||
except BrokenPipeError:
|
||
pass
|
||
|
||
def release(self):
|
||
"""关闭管道并等待编码完成。"""
|
||
try:
|
||
self._process.stdin.close()
|
||
except BrokenPipeError:
|
||
pass
|
||
self._process.wait()
|