mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-04-14 16:47:31 +08:00
93d822d067
- 限制字幕区间end不超过frame_count，防止内循环消费哨兵后外层永久阻塞 - LAMA批量推理改为mini-batch(4帧)，避免GPU OOM - 各inpaint模型空inpaint_area时返回原始帧 - FFmpeg子进程添加600s超时保护 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
105 lines
2.8 KiB
Python
105 lines
2.8 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
|
||
try:
|
||
self._process.wait(timeout=600)
|
||
except subprocess.TimeoutExpired:
|
||
self._process.terminate()
|
||
self._process.wait(timeout=5)
|