diff --git a/nodes/frame_extract.py b/nodes/frame_extract.py index c2fdfb6..1d8aabf 100644 --- a/nodes/frame_extract.py +++ b/nodes/frame_extract.py @@ -12,6 +12,7 @@ from __future__ import annotations import json import re import subprocess +import time from pathlib import Path from nodes.ffmpeg import _ffmpeg_bin @@ -96,6 +97,20 @@ def _frame_step(fps: float, interval: float) -> int: return max(1, int(round(interval * fps))) +def _parse_progress_line(line: str) -> int | None: + """解析 ffmpeg -progress 输出行，返回 frame=N 的数值；其他行返回 None。 + + -progress 每 ~0.5s 输出一组 key=value（frame/fps/out_time_ms/progress 等）， + 这里只关心 frame=（已写出的选中帧数），用于实时进度与速度统计。 + """ + stripped = line.strip() + if not stripped.startswith("frame="): + return None + try: + return int(stripped.split("=", 1)[1]) + except ValueError: + return None + def invoke(request: InvokeRequest) -> InvokeResponse: """按帧间隔抽取并裁切视频帧，输出 frames.json 清单。""" video_uri = request.inputs.get("video_uri") @@ -137,13 +152,18 @@ def invoke(request: InvokeRequest) -> InvokeResponse: # 单次解码全片：select 按帧号（n mod step == 0）精确取帧，随后对选中帧 # 裁切字幕区域并压缩到 720p 内（仅缩小）——过大的输入图会触发 glm-ocr # 的重复循环 bug（M-RoPE delta），源头规避。 - result = subprocess.run( + # -progress 把 frame=N 进度写入管道，实时打印抽帧进度与速度。 + expected_frames = (int(duration * fps) + step - 1) // step + process = subprocess.Popen( [ ffmpeg_bin, "-y", "-hide_banner", "-loglevel", "error", + "-nostats", + "-progress", + "pipe:1", "-i", str(video), "-vf", @@ -157,11 +177,30 @@ def invoke(request: InvokeRequest) -> InvokeResponse: "vfr", str(frames_dir / "frame_%04d.png"), ], - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, ) - if result.returncode != 0: - return InvokeResponse(status="failed", error=result.stderr[-500:] or "ffmpeg failed") + extract_started = time.monotonic() + last_logged = 0 + progress_step = max(1, expected_frames // 10) + for line in process.stdout: + done = _parse_progress_line(line) + if done is None: + continue + # 每写出约 10 个检查点打一次日志，避免刷屏；同时给出处理速度。 + if done - last_logged >= progress_step: + elapsed = time.monotonic() - extract_started + logger.info( + "抽帧进度 %d/%d 帧 (%.1f 帧/s)", + min(done, expected_frames), expected_frames, + done / elapsed if elapsed > 0 else 0.0, + ) + last_logged = done + process.wait() + stderr = process.stderr.read() + if process.returncode != 0: + return InvokeResponse(status="failed", error=stderr[-500:] or "ffmpeg failed") # 第 k 个输出文件对应原始帧号 k×step，时间 = 帧号 / fps（帧精确，无累计偏差）。 files = sorted(frames_dir.glob("frame_*.png")) diff --git a/tests/test_ocr_flow.py b/tests/test_ocr_flow.py index b8b0bbb..1209420 100644 --- a/tests/test_ocr_flow.py +++ b/tests/test_ocr_flow.py @@ -159,6 +159,84 @@ def test_video_duration_unreadable(monkeypatch) -> None: 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))