diff --git a/nodes/frame_extract.py b/nodes/frame_extract.py index 1d8aabf..4165ac6 100644 --- a/nodes/frame_extract.py +++ b/nodes/frame_extract.py @@ -111,6 +111,24 @@ def _parse_progress_line(line: str) -> int | None: except ValueError: return None + +def _sorted_frame_files(frames_dir: Path) -> list[Path]: + """按文件名中的帧号数值排序返回帧文件列表(自然排序,非字典序)。 + + 关键点:ffmpeg 的 %04d 编号在超过 9999 帧后会自动扩为 5 位 + (frame_10000.png 等),此时 sorted() 默认的字典序会把 5 位编号排在 + 4 位编号之前(如 "frame_10009" < "frame_1009"),导致帧号回退、 + manifest 时间与图像错位(曾真实发生于 run_339ec7ee437f 的 14236 帧 + 任务,全片后半段时间轴全部错乱)。必须解析出帧号按数值排序, + 才能保证"第 k 个文件 = 第 k 个选中帧 = 时间 index*step/fps"成立。 + """ + def frame_number(path: Path) -> int: + # 文件名形如 frame_0001.png,取下划线后的数字部分。 + return int(path.stem.split("_", 1)[1]) + + return sorted(frames_dir.glob("frame_*.png"), key=frame_number) + + def invoke(request: InvokeRequest) -> InvokeResponse: """按帧间隔抽取并裁切视频帧,输出 frames.json 清单。""" video_uri = request.inputs.get("video_uri") @@ -203,7 +221,7 @@ def invoke(request: InvokeRequest) -> InvokeResponse: return InvokeResponse(status="failed", error=stderr[-500:] or "ffmpeg failed") # 第 k 个输出文件对应原始帧号 k×step,时间 = 帧号 / fps(帧精确,无累计偏差)。 - files = sorted(frames_dir.glob("frame_*.png")) + files = _sorted_frame_files(frames_dir) manifest = [ {"time": round((index * step) / fps, 3), "image_uri": str(path)} for index, path in enumerate(files) diff --git a/nodes/subtitle_ocr.py b/nodes/subtitle_ocr.py index 9b8c46b..4a76793 100644 --- a/nodes/subtitle_ocr.py +++ b/nodes/subtitle_ocr.py @@ -63,6 +63,28 @@ def _assemble_srt( return lines +def _merge_kept( + manifest: list[dict], texts: list[str] +) -> list[tuple[float, float, str]]: + """按 manifest 时间轴把逐帧 OCR 文本合并为字幕条目。 + + kept 元素为 (起始帧时间, 最后可见帧时间, 文本):空文本(无文字帧)跳过; + 连续帧相同字幕合并为一条(字幕停留多帧属正常现象),仅更新最后可见帧 + 时间,起始时间保持首次出现。要求 texts 与 manifest 按帧顺序一一对应 + (调用方保证),重组装旧数据时也复用此逻辑保证行为一致。 + """ + kept: list[tuple[float, float, str]] = [] + for index, text in enumerate(texts): + if not text: + continue + time = float(manifest[index]["time"]) + if kept and kept[-1][2] == text: + kept[-1] = (kept[-1][0], time, text) + continue + kept.append((time, time, text)) + return kept + + def invoke(request: InvokeRequest) -> InvokeResponse: """逐帧 OCR 并汇总字幕,产物为 subtitle.srt。""" manifest_uri = request.inputs.get("frames_manifest") @@ -138,19 +160,8 @@ def invoke(request: InvokeRequest) -> InvokeResponse: ) texts = pool.map(list(enumerate(manifest))) - # kept 元素为 (起始帧时间, 最后可见帧时间, 文本);按帧顺序合并连续相同字幕。 - kept: list[tuple[float, float, str]] = [] - for index, text in enumerate(texts): - if not text: - continue - time = float(manifest[index]["time"]) - # 连续帧相同字幕合并为一条(字幕停留多帧属正常现象): - # 仅更新最后可见帧时间,起始时间保持首次出现。 - if kept and kept[-1][2] == text: - kept[-1] = (kept[-1][0], time, text) - continue - kept.append((time, time, text)) - + # 按帧顺序合并连续相同字幕(与重组装旧数据共用 _merge_kept)。 + kept = _merge_kept(manifest, texts) output_dir = Path(request.output_dir) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / "subtitle.srt" diff --git a/testdata/frames_boundary/frame_0001.png b/testdata/frames_boundary/frame_0001.png new file mode 100644 index 0000000..ec59afe Binary files /dev/null and b/testdata/frames_boundary/frame_0001.png differ diff --git a/testdata/frames_boundary/frame_0999.png b/testdata/frames_boundary/frame_0999.png new file mode 100644 index 0000000..347fee2 Binary files /dev/null and b/testdata/frames_boundary/frame_0999.png differ diff --git a/testdata/frames_boundary/frame_1000.png b/testdata/frames_boundary/frame_1000.png new file mode 100644 index 0000000..5903ebd Binary files /dev/null and b/testdata/frames_boundary/frame_1000.png differ diff --git a/testdata/frames_boundary/frame_10000.png b/testdata/frames_boundary/frame_10000.png new file mode 100644 index 0000000..3b75861 Binary files /dev/null and b/testdata/frames_boundary/frame_10000.png differ diff --git a/testdata/frames_boundary/frame_10001.png b/testdata/frames_boundary/frame_10001.png new file mode 100644 index 0000000..05a9094 Binary files /dev/null and b/testdata/frames_boundary/frame_10001.png differ diff --git a/testdata/frames_boundary/frame_10009.png b/testdata/frames_boundary/frame_10009.png new file mode 100644 index 0000000..b234866 Binary files /dev/null and b/testdata/frames_boundary/frame_10009.png differ diff --git a/testdata/frames_boundary/frame_1009.png b/testdata/frames_boundary/frame_1009.png new file mode 100644 index 0000000..9c7e06a Binary files /dev/null and b/testdata/frames_boundary/frame_1009.png differ diff --git a/testdata/frames_boundary/frame_14235.png b/testdata/frames_boundary/frame_14235.png new file mode 100644 index 0000000..8f47466 Binary files /dev/null and b/testdata/frames_boundary/frame_14235.png differ diff --git a/testdata/frames_boundary/frame_14236.png b/testdata/frames_boundary/frame_14236.png new file mode 100644 index 0000000..77129a2 Binary files /dev/null and b/testdata/frames_boundary/frame_14236.png differ diff --git a/testdata/frames_boundary/frame_1999.png b/testdata/frames_boundary/frame_1999.png new file mode 100644 index 0000000..d5437c8 Binary files /dev/null and b/testdata/frames_boundary/frame_1999.png differ diff --git a/testdata/frames_boundary/frame_2000.png b/testdata/frames_boundary/frame_2000.png new file mode 100644 index 0000000..c60d43e Binary files /dev/null and b/testdata/frames_boundary/frame_2000.png differ diff --git a/testdata/frames_boundary/frame_5762.png b/testdata/frames_boundary/frame_5762.png new file mode 100644 index 0000000..a6bd097 Binary files /dev/null and b/testdata/frames_boundary/frame_5762.png differ diff --git a/testdata/frames_boundary/frame_9999.png b/testdata/frames_boundary/frame_9999.png new file mode 100644 index 0000000..5fb80ba Binary files /dev/null and b/testdata/frames_boundary/frame_9999.png differ diff --git a/tests/test_integration_reassemble_ocr.py b/tests/test_integration_reassemble_ocr.py new file mode 100644 index 0000000..a7d51ae --- /dev/null +++ b/tests/test_integration_reassemble_ocr.py @@ -0,0 +1,97 @@ +"""真实任务 OCR 数据重组装集成测试。 + +run_339ec7ee437f 是 2026-08 真实跑过的 ocr-subtitle 任务(14236 帧 / 2 小时 +视频)。该任务产出时期存在 frame-extract 帧文件排序 bug:ffmpeg 的 %04d 编号 +超过 9999 帧后扩为 5 位,sorted() 字典序把 5 位编号排在 4 位之前,导致 +frames.json 中 time 与图像错位(13236/14236 条位置↔帧号错位),最终 SRT +后半段时间轴全部错乱。 + +本测试读取该任务**已落盘的逐帧 OCR 文本**(ocr_frames/<位置>/ocr.txt,由 +vlm-ocr 节点写入的清洗后文本),用 image_uri 解析真实帧号重建正确时间轴, +复用生产合并/组装逻辑生成 SRT,并断言"每条字幕的起始时刻 = 其文本来源帧的 +真实时间"这一核心正确性。 + +数据位于 gitignored 的 data/storage,缺失时跳过(与 +test_integration_subtitle_ocr 的真实资产约定一致)。 +""" + +import json +import re +from pathlib import Path + +import pytest + +from nodes.subtitle_ocr import _assemble_srt +from nodes.subtitle_ocr import _merge_kept +from nodes.subtitle_ocr import _sampling_interval + +# 单体根目录:tests/ 的上一级。 +WORKSPACE = Path(__file__).resolve().parent.parent +RUN_DIR = WORKSPACE / "data/storage/runs/run_339ec7ee437f" +MANIFEST = RUN_DIR / "steps/extract/frames.json" +OCR_DIR = RUN_DIR / "steps/ocr/ocr_frames" + +# 源视频时长(ffprobe: 02:00:24.80)。 +VIDEO_DURATION = 7224.8 +# 该视频烧录过的字幕行,被 test_real_hav_sub.png 集成测试确认过识别结果。 +KNOWN_LINE = "还有没有什么困扰" + + +def _frame_number(uri: str) -> int: + """从 image_uri 文件名解析真实帧号(frame_0001.png -> 1)。""" + return int(re.match(r".*frame_(\d+)\.png", Path(uri).name).group(1)) + + +@pytest.mark.integration +def test_reassemble_real_run_ocr_data(tmp_path) -> None: + """真实任务逐帧 OCR 数据按真实帧时间轴重组装为正确 SRT。""" + if not (MANIFEST.is_file() and OCR_DIR.is_dir()): + pytest.skip("缺少真实任务 run_339ec7ee437f 数据,跳过集成测试") + + manifest = json.loads(MANIFEST.read_text(encoding="utf-8")) + total = len(manifest) + # 等间隔采样:step/fps 从旧 manifest 相邻 time 差恢复(本任务 0.507s/帧)。 + dt = manifest[1]["time"] - manifest[0]["time"] + assert dt > 0 + + # ① 读取逐帧 OCR 文本:位置 i -> ocr_frames/{i:04d}/ocr.txt(真实数据)。 + texts_by_pos: list[str] = [] + for i in range(total): + p = OCR_DIR / f"{i:04d}" / "ocr.txt" + if not p.is_file(): + pytest.fail(f"缺少逐帧 OCR 数据: {p}") + texts_by_pos.append(p.read_text(encoding="utf-8").strip()) + + # ② 旧 manifest 位置 -> 真实帧号:image_uri 才是实际被 OCR 的图像, + # 旧 time 字段按位置推导已错位,不可用。 + frame_by_pos = [_frame_number(item["image_uri"]) for item in manifest] + assert sorted(frame_by_pos) == list(range(1, total + 1)) # 一一对应,无缺无重 + + # ③ 按真实帧号重建正确时间轴:第 k 帧(1-based)时间 = (k-1)*dt。 + ocr_by_frame = {n: texts_by_pos[i] for i, n in enumerate(frame_by_pos)} + manifest_corrected = [{"time": round((k - 1) * dt, 3)} for k in range(1, total + 1)] + texts_corrected = [ocr_by_frame[k] for k in range(1, total + 1)] + + # ④ 复用生产合并 + 组装逻辑(与 subtitle_ocr.invoke 完全同一路径)。 + kept = _merge_kept(manifest_corrected, texts_corrected) + interval = _sampling_interval(manifest_corrected, dt) + srt = "\n".join(_assemble_srt(kept, interval)) + "\n" + out = tmp_path / "subtitle.srt" + out.write_text(srt, encoding="utf-8") + + # ⑤ 核心正确性:每条字幕起始时刻对应的帧,其 OCR 文本必须就是本条字幕 + # 文本(修复前旧 SRT 此处大面积不一致:文本来自其他时刻的帧)。 + for start, _end, text in kept: + frame_no = int(round(start / dt)) + 1 + assert ocr_by_frame[frame_no] == text, f"帧 {frame_no} 时间错位: {text!r}" + + # ⑥ 时间轴严格递增且不超出片长(含尾部一个采样间隔的消失余量)。 + assert all( + s1 < s2 for (s1, _1, _), (s2, _2, _) in zip(kept, kept[1:]) + ) + assert kept[0][0] >= 0.0 + assert kept[-1][1] + interval <= VIDEO_DURATION + interval + 1e-6 + + # ⑦ 已知烧录字幕行必须出现,且字幕条数合理(空帧被跳过、连续相同字幕合并)。 + assert KNOWN_LINE in srt + assert len(kept) > 100 diff --git a/tests/test_ocr_flow.py b/tests/test_ocr_flow.py index 1209420..c27aa86 100644 --- a/tests/test_ocr_flow.py +++ b/tests/test_ocr_flow.py @@ -305,6 +305,26 @@ def test_frame_extract_one_second_exact_frames(tmp_path) -> None: assert len(manifest) == 10 +def test_frame_files_read_order_matches_frame_number(tmp_path) -> None: + """真实任务留存数据:帧文件按帧号数值排序读取,而非字典序。 + + 回归用例:ffmpeg 的 %04d 编号在超过 9999 帧后自动扩为 5 位 + (frame_10000.png 等),此时 sorted() 默认字典序会把 5 位编号排在 + 4 位编号之前(如 frame_10009 < frame_1009),导致帧号回退、manifest + 时间与图像错位。testdata/frames_boundary/ 是 2026-08 真实任务 + run_339ec7ee437f(14236 帧 / 2 小时视频)中跨越该边界的真实帧文件。 + """ + from nodes.frame_extract import _sorted_frame_files + + boundary = TESTDATA / "frames_boundary" + files = _sorted_frame_files(boundary) + # 从文件名解析帧号:读取顺序必须等于帧号数值递增序(无回退)。 + nums = [int(p.stem.split("_", 1)[1]) for p in files] + assert nums == sorted(nums) + # 边界关键对:5 位编号必须排在 4 位编号之后,禁止字典序错位。 + assert nums.index(10000) > nums.index(9999) + assert nums.index(10009) > nums.index(1009) + # --------------------------------------------------------------------------- # subtitle-ocr:OCR 循环 + 长度上限 + 合并 + SRT 组装 # ---------------------------------------------------------------------------