Files
vrsub/tests/test_ocr_flow.py
T
cat-shark 885cf07921 feat(extract): 抽帧打印进度日志
ffmpeg 增加 -progress pipe:1 并解析 frame=N:每约 10 个检查点打印
'抽帧进度 X/Y 帧 (Z 帧/s)',长视频抽帧期间可见实时进度与处理速度。
测试覆盖进度行解析、进度日志输出与失败路径(Popen mock)。
2026-08-17 00:03:51 +08:00

480 lines
19 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""抽帧与字幕 OCR 节点单元测试。
frame-extract 用 testdata 真实视频抽帧+裁切+720p 压缩;subtitle-ocr 的 OCR
网络调用(vlm-ocr)按 I/O 边界 mock,但喂给它的帧图片是真实的 testdata 资产。
应用层不再加工模型输出文本,只做长度上限校验(超长报错跳过)。
"""
import json
from pathlib import Path
from wov_sdk.models import InvokeRequest, InvokeResponse
from nodes.frame_extract import invoke as frame_invoke
from nodes.subtitle_ocr import _assemble_srt
from nodes.subtitle_ocr import invoke as ocr_invoke
# 单体根目录:tests/ 的上一级。
WORKSPACE = Path(__file__).resolve().parent.parent
TESTDATA = WORKSPACE / "testdata"
# 10s 测试视频:SUB 001 在 1-4s、SUB 002 在 6-9s。
VIDEO = TESTDATA / "subtitle_10s.mp4"
TEXT_IMG = TESTDATA / "ocr_text.png"
def _png_size(path: Path) -> tuple[int, int]:
"""从 PNG 头读取宽高(真实图片尺寸断言)。"""
data = path.read_bytes()
assert data[:8] == b"\x89PNG\r\n\x1a\n", "not a real png"
width = int.from_bytes(data[16:20], "big")
height = int.from_bytes(data[20:24], "big")
return width, height
def _frame_request(tmp_path, video=VIDEO, **params) -> InvokeRequest:
"""构造 frame-extract 调用请求。"""
return InvokeRequest(
run_id="run_fx",
node_instance_id="",
inputs={"video_uri": str(video)},
params=params,
output_dir=str(tmp_path / "out"),
)
# ---------------------------------------------------------------------------
# frame-extract:抽帧 + 裁切 + 720p 压缩
# ---------------------------------------------------------------------------
def test_frame_extract_crop_and_manifest(tmp_path) -> None:
"""真实视频抽帧:裁切下半 50% 后帧尺寸为 1280x360,清单时间轴正确。"""
response = frame_invoke(
_frame_request(tmp_path, interval_seconds=1, crop=[0, 0.5, 1, 0.5])
)
assert response.status == "completed", response.error
manifest = json.loads(Path(response.outputs["frames_manifest"]).read_text(encoding="utf-8"))
assert len(manifest) >= 9
assert [round(item["time"], 3) for item in manifest] == [
round(i * 1.0, 3) for i in range(len(manifest))
]
first = Path(manifest[0]["image_uri"])
assert first.is_file()
# 1280x360 已在 720p 内,压缩不改变尺寸。
assert _png_size(first) == (1280, 360)
def test_frame_extract_default_params(tmp_path) -> None:
"""未指定参数时使用默认值:抽帧间隔 0.5 秒 + 默认底部裁切区域。"""
response = frame_invoke(_frame_request(tmp_path))
assert response.status == "completed", response.error
manifest = json.loads(Path(response.outputs["frames_manifest"]).read_text(encoding="utf-8"))
assert manifest
# 默认间隔 0.5s25fps 下 step=round(12.5)=12(银行家舍入),
# 帧时间按 step/fps=12/25=0.48s 步进(帧号精确,采样周期由帧量化决定)。
assert [round(item["time"], 3) for item in manifest] == [
round(i * 12 / 25, 3) for i in range(len(manifest))
]
def test_frame_extract_missing_video(tmp_path) -> None:
"""缺少 video_uri 时返回失败。"""
response = frame_invoke(
InvokeRequest(
run_id="r", node_instance_id="", inputs={}, output_dir=str(tmp_path)
)
)
assert response.status == "failed"
def test_frame_extract_bad_crop(tmp_path) -> None:
"""crop 比例越界(超出画面)时返回失败。"""
assert frame_invoke(_frame_request(tmp_path, crop=[0, 0.5, 1, 1.5])).status == "failed"
assert frame_invoke(_frame_request(tmp_path, crop=[-0.1, 0, 1, 0.5])).status == "failed"
assert frame_invoke(_frame_request(tmp_path, crop="abc")).status == "failed"
assert frame_invoke(_frame_request(tmp_path, crop=[0, 0.5, 1])).status == "failed"
# 各值域合法但 x+w 越出画面。
assert frame_invoke(_frame_request(tmp_path, crop=[0.6, 0, 0.5, 0.3])).status == "failed"
def test_frame_extract_video_missing_file(tmp_path) -> None:
"""video_uri 指向不存在的文件时返回失败。"""
response = frame_invoke(_frame_request(tmp_path, video=tmp_path / "none.mp4"))
assert response.status == "failed"
assert "not found" in response.error
def test_frame_extract_bad_interval(tmp_path) -> None:
"""间隔 <= 0 时返回失败。"""
response = frame_invoke(_frame_request(tmp_path, interval_seconds=0))
assert response.status == "failed"
def test_frame_extract_ffmpeg_fails(monkeypatch, tmp_path) -> None:
"""ffmpeg 抽帧失败时透传错误。"""
import subprocess as sp
monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720))
monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: 25.0)
monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0)
monkeypatch.setattr(
"nodes.frame_extract.subprocess.run",
lambda *a, **k: sp.CompletedProcess([], 1, stderr="boom"),
)
response = frame_invoke(_frame_request(tmp_path))
assert response.status == "failed"
assert "boom" in response.error
def test_video_size_unreadable(monkeypatch) -> None:
"""ffmpeg -i 输出不含视频流信息时返回 None。"""
import subprocess as sp
from nodes.frame_extract import _video_size
monkeypatch.setattr(
"nodes.frame_extract.subprocess.run",
lambda *a, **k: sp.CompletedProcess([], 0, stderr="no video stream"),
)
assert _video_size(Path("/tmp/x.mp4"), "ffmpeg") is None
def test_frame_extract_video_size_unknown(monkeypatch, tmp_path) -> None:
"""无法读取视频分辨率时返回失败。"""
monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: None)
response = frame_invoke(_frame_request(tmp_path))
assert response.status == "failed"
assert "video size" in response.error
def test_video_duration_unreadable(monkeypatch) -> None:
"""ffmpeg -i 输出缺少 Duration 时返回 None。"""
import subprocess as sp
from nodes.frame_extract import _video_duration
monkeypatch.setattr(
"nodes.frame_extract.subprocess.run",
lambda *a, **k: sp.CompletedProcess([], 0, stderr="no duration info"),
)
assert _video_duration(Path("/tmp/x.mp4"), "ffmpeg") is None
class _FakeReader:
"""模拟 stderr 读取对象。"""
def __init__(self, content: str = "") -> None:
self._content = content
def read(self) -> str:
return self._content
class FakePopen:
"""模拟 ffmpeg 进程:stdout 可迭代 -progress 行,可配置返回码与 stderr。"""
def __init__(self, lines=(), returncode: int = 0, stderr: str = "") -> None:
self.stdout = list(lines)
self.stderr = _FakeReader(stderr)
self.returncode = returncode
def wait(self) -> int:
return self.returncode
def _fake_popen(lines=(), returncode: int = 0, stderr: str = ""):
"""构造替换 subprocess.Popen 的工厂函数。"""
return lambda *a, **k: FakePopen(lines, returncode, stderr)
def test_parse_progress_line() -> None:
"""-progress 行解析:frame=N 返回数值,其他行与非法值返回 None。"""
from nodes.frame_extract import _parse_progress_line
assert _parse_progress_line("frame=25\n") == 25
assert _parse_progress_line("progress=continue") is None
assert _parse_progress_line("fps=25.0") is None
assert _parse_progress_line("frame=abc") is None
def test_frame_extract_progress_logging(monkeypatch, tmp_path) -> None:
"""ffmpeg -progress 的 frame=N 被解析并打印抽帧进度与速度。"""
import logging
monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720))
monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: 25.0)
monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0)
monkeypatch.setattr(
"nodes.frame_extract.subprocess.Popen",
_fake_popen(lines=["fps=25.0", "frame=20", "progress=continue", "frame=50", "progress=continue"]),
)
captured: list[str] = []
class CaptureHandler(logging.Handler):
def emit(self, record):
captured.append(record.getMessage())
logger = logging.getLogger("vrsub.frame-extract")
logger.addHandler(CaptureHandler())
try:
response = frame_invoke(_frame_request(tmp_path))
finally:
logger.removeHandler(logger.handlers[-1])
assert response.status == "completed", response.error
assert any("抽帧进度" in message and "帧/s" in message for message in captured)
def test_frame_extract_ffmpeg_fails(monkeypatch, tmp_path) -> None:
"""ffmpeg 抽帧失败时透传错误。"""
monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720))
monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: 25.0)
monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0)
monkeypatch.setattr(
"nodes.frame_extract.subprocess.Popen",
_fake_popen(lines=["frame=1", "progress=end"], returncode=1, stderr="boom"),
)
response = frame_invoke(_frame_request(tmp_path))
assert response.status == "failed"
assert "boom" in response.error
def test_frame_extract_duration_unknown(monkeypatch, tmp_path) -> None:
"""无法读取视频时长时返回失败。"""
monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720))
monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: None)
response = frame_invoke(_frame_request(tmp_path))
assert response.status == "failed"
assert "duration" in response.error
def test_frame_extract_fps_unknown(monkeypatch, tmp_path) -> None:
"""无法读取视频帧率时返回失败。"""
monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720))
monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0)
monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: None)
response = frame_invoke(_frame_request(tmp_path))
assert response.status == "failed"
assert "fps" in response.error
def test_frame_step_conversion() -> None:
"""帧间隔换算:step=round(间隔秒×fps),至少为 1。"""
from nodes.frame_extract import _frame_step
assert _frame_step(fps=25.0, interval=0.2) == 5
assert _frame_step(fps=25.0, interval=1.0) == 25
assert _frame_step(fps=29.97, interval=1.0) == 30
# fps 很低时 step 也不会小于 1(每帧都取)。
assert _frame_step(fps=1.0, interval=0.2) == 1
def test_video_fps_parse(monkeypatch) -> None:
"""帧率解析:支持小数(29.97)与有理数(30000/1001)。"""
import subprocess as sp
from nodes.frame_extract import _video_fps
monkeypatch.setattr(
"nodes.frame_extract.subprocess.run",
lambda *a, **k: sp.CompletedProcess(
[], 0, stderr="Stream #0:0: Video: h264, 1280x720, 30000/1001 fps, 30000/1001 tbr"
),
)
assert _video_fps(Path("/tmp/x.mp4"), "ffmpeg") == 30000 / 1001
monkeypatch.setattr(
"nodes.frame_extract.subprocess.run",
lambda *a, **k: sp.CompletedProcess(
[], 0, stderr="Stream #0:0: Video: h264, 1280x720, 25 fps, 25 tbr"
),
)
assert _video_fps(Path("/tmp/x.mp4"), "ffmpeg") == 25.0
monkeypatch.setattr(
"nodes.frame_extract.subprocess.run",
lambda *a, **k: sp.CompletedProcess([], 0, stderr="no video stream"),
)
assert _video_fps(Path("/tmp/x.mp4"), "ffmpeg") is None
def test_frame_extract_one_second_exact_frames(tmp_path) -> None:
"""真实视频 1s 间隔按秒 seek 精确抽帧:10s 视频应得 10 帧,时间 0..9。"""
response = frame_invoke(
_frame_request(tmp_path, interval_seconds=1, crop=[0, 0.7, 1, 0.3])
)
assert response.status == "completed", response.error
manifest = json.loads(Path(response.outputs["frames_manifest"]).read_text(encoding="utf-8"))
assert [round(item["time"], 3) for item in manifest] == [
round(i * 1.0, 3) for i in range(len(manifest))
]
assert len(manifest) == 10
# ---------------------------------------------------------------------------
# subtitle-ocrOCR 循环 + 长度上限 + 合并 + SRT 组装
# ---------------------------------------------------------------------------
def _frames_manifest(tmp_path, frame_specs) -> Path:
"""构造真实 frames.jsonframe_specs=[(time, image_path), ...]。"""
items = [{"time": time, "image_uri": str(image)} for time, image in frame_specs]
path = tmp_path / "frames.json"
path.write_text(json.dumps(items), encoding="utf-8")
return path
def test_assemble_srt_real_timeline() -> None:
"""SRT 组装:起始=帧时间,结束=最后可见帧时间+采样间隔。"""
lines = _assemble_srt([(0.0, 6.0, "A"), (6.0, 8.0, "B")], interval_seconds=2.0)
text = "\n".join(lines)
assert text.startswith("1\n")
# A 最后可见帧 6.0 + 间隔 2.0 = 8.0(而非下一条字幕的出现时间)。
assert "00:00:00,000 --> 00:00:08,000" in text
assert "00:00:06,000 --> 00:00:10,000" in text
def test_sampling_interval_from_manifest() -> None:
"""采样间隔从帧清单时间轴推导:均匀间隔取相邻差,退化清单回退默认值。"""
from nodes.subtitle_ocr import _sampling_interval
manifest = [{"time": i * 0.2, "image_uri": f"f{i}.png"} for i in range(10)]
assert _sampling_interval(manifest, 2.0) == 0.2
# 单帧(无法算差)与异常时间序:回退默认值。
assert _sampling_interval([{"time": 0.0, "image_uri": "f0.png"}], 2.0) == 2.0
assert _sampling_interval(
[{"time": 0.0}, {"time": 0.0}, {"time": 0.2}], 2.0
) == 0.2
def test_ocr_merges_consecutive_same_text(monkeypatch, tmp_path) -> None:
"""连续帧相同字幕合并为一条;消失时间=最后可见帧+间隔,空白段保留。"""
# SUB 001 在 0/2s4s 为空帧,SUB 002 在 6/8s。
frame_texts = [
(0.0, "SUB 001"), (2.0, "SUB 001"), (4.0, ""),
(6.0, "SUB 002"), (8.0, "SUB 002"),
]
frames = []
for index, (time, _text) in enumerate(frame_texts):
image = tmp_path / f"f{index}.png"
image.write_bytes(TEXT_IMG.read_bytes())
frames.append((time, image))
mapping = {str(image): text for (time, image), (_, text) in zip(frames, frame_texts)}
def fake_vlm(node_id, request):
return InvokeResponse(status="completed", outputs={"text": mapping[request.inputs["image_uri"]]})
monkeypatch.setattr("wov_app.registry.invoke", fake_vlm)
manifest = _frames_manifest(tmp_path, frames)
response = ocr_invoke(
InvokeRequest(
run_id="run_ocr",
node_instance_id="",
inputs={"frames_manifest": str(manifest)},
params={},
output_dir=str(tmp_path / "out"),
)
)
assert response.status == "completed", response.error
srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8")
assert srt.count("SUB 001") == 1
assert srt.count("SUB 002") == 1
# SUB 001 最后可见帧 2.0 + 间隔 2.0 = 4.0 消失(而非拖到 SUB 002 出现)。
assert "00:00:00,000 --> 00:00:04,000" in srt
assert "00:00:06,000 --> 00:00:10,000" in srt
def test_ocr_skips_failed_frames(monkeypatch, tmp_path) -> None:
"""个别帧 OCR 失败时跳过,不影响其余帧汇总。"""
frames = []
for index in range(3):
image = tmp_path / f"f{index}.png"
image.write_bytes(TEXT_IMG.read_bytes())
frames.append((index * 2.0, image))
def fake_vlm(node_id, request):
if "f1" in request.inputs["image_uri"]:
return InvokeResponse(status="failed", error="boom")
return InvokeResponse(status="completed", outputs={"text": "SUB 001"})
monkeypatch.setattr("wov_app.registry.invoke", fake_vlm)
manifest = _frames_manifest(tmp_path, frames)
response = ocr_invoke(
InvokeRequest(
run_id="run_ocr",
node_instance_id="",
inputs={"frames_manifest": str(manifest)},
params={},
output_dir=str(tmp_path / "out"),
)
)
assert response.status == "completed", response.error
srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8")
assert "SUB 001" in srt
def test_ocr_missing_manifest(tmp_path) -> None:
"""缺少 frames_manifest 或清单文件不存在时返回失败。"""
response = ocr_invoke(
InvokeRequest(run_id="r", node_instance_id="", inputs={}, output_dir=str(tmp_path))
)
assert response.status == "failed"
response = ocr_invoke(
InvokeRequest(
run_id="r", node_instance_id="",
inputs={"frames_manifest": str(tmp_path / "none.json")},
output_dir=str(tmp_path),
)
)
assert response.status == "failed"
def test_ocr_skips_oversized_output(monkeypatch, tmp_path) -> None:
"""超长输出(模型重复循环等)直接报错跳过该帧,不进入 SRT。"""
frames = []
for index in range(3):
image = tmp_path / f"o{index}.png"
image.write_bytes(TEXT_IMG.read_bytes())
frames.append((index * 2.0, image))
def fake_vlm(node_id, request):
if "o0" in request.inputs["image_uri"]:
return InvokeResponse(status="completed", outputs={"text": "重复字幕\n" * 50})
if "o1" in request.inputs["image_uri"]:
# 空输出帧跳过。
return InvokeResponse(status="completed", outputs={"text": ""})
return InvokeResponse(status="completed", outputs={"text": "SUB 001"})
monkeypatch.setattr("wov_app.registry.invoke", fake_vlm)
manifest = _frames_manifest(tmp_path, frames)
response = ocr_invoke(
InvokeRequest(
run_id="run_ocr",
node_instance_id="",
inputs={"frames_manifest": str(manifest)},
params={"max_result_chars": 200},
output_dir=str(tmp_path / "out"),
)
)
assert response.status == "completed", response.error
srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8")
assert "SUB 001" in srt
assert "重复字幕" not in srt
def test_ocr_passes_short_text_through(monkeypatch, tmp_path) -> None:
"""不超过上限的模型输出原样进入 SRT(不再做应用层过滤)。"""
image = tmp_path / "s0.png"
image.write_bytes(TEXT_IMG.read_bytes())
frames = [(0.0, image)]
def fake_vlm(node_id, request):
return InvokeResponse(status="completed", outputs={"text": " SUB 001 "})
monkeypatch.setattr("wov_app.registry.invoke", fake_vlm)
manifest = _frames_manifest(tmp_path, frames)
response = ocr_invoke(
InvokeRequest(
run_id="run_ocr", node_instance_id="",
inputs={"frames_manifest": str(manifest)},
params={},
output_dir=str(tmp_path / "out"),
)
)
assert response.status == "completed", response.error
srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8")
assert "SUB 001" in srt