cat-shark
2b3a650612
- frame-extract 新增 _sorted_frame_files：ffmpeg %04d 编号超过 9999 帧后扩为 5 位，sorted() 字典序会把 5 位编号排在 4 位之前，导致 frames.json 时间与 图像错位（真实发生于 run_339ec7ee437f 的 14236 帧任务） - subtitle-ocr 提取 _merge_kept 供重组装复用 - 回归测试用真实任务留存数据（testdata/frames_boundary/），并新增真实 OCR 数据按正确时间轴重组装为 SRT 的集成测试（test_integration_reassemble_ocr）
174 lines
7.5 KiB
Python
174 lines
7.5 KiB
Python
"""字幕 OCR 汇总节点。
|
||
|
||
读取 frame-extract 产出的 frames.json，逐帧调用 vlm-ocr 节点识别字幕文字；
|
||
过滤无文字帧的垃圾输出（glm-ocr 在空帧上会输出无用文字），折叠模型重复
|
||
循环输出，合并连续相同的字幕（记录最后可见帧时间），最终组装为带时间轴的
|
||
SRT 基准数据：每条字幕消失时间 = 最后可见帧时间 + 采样间隔（间隔从帧清单
|
||
时间轴推导），与视频烧录时间对齐。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from pathlib import Path
|
||
|
||
from nodes.adaptive_pool import AdaptiveThreadPool
|
||
from nodes.whisper import format_timestamp
|
||
from wov_app import registry
|
||
from wov_app.logging import get_logger
|
||
from wov_sdk.models import InvokeRequest, InvokeResponse
|
||
|
||
logger = get_logger("subtitle-ocr")
|
||
|
||
# 默认垃圾词：无文字帧的模型输出可能反复出现这些词。
|
||
def _sampling_interval(manifest: list[dict], default: float) -> float:
|
||
"""从帧清单时间轴推导采样间隔（相邻帧时间差的中位数）。
|
||
|
||
帧时间由 frame-extract 按固定间隔生成，取相邻差的中位数即可还原真实
|
||
采样间隔，避免结束时间与抽取参数不一致。清单不足两帧时回退默认值。
|
||
"""
|
||
diffs = [
|
||
float(manifest[i + 1]["time"]) - float(manifest[i]["time"])
|
||
for i in range(len(manifest) - 1)
|
||
if float(manifest[i + 1]["time"]) > float(manifest[i]["time"])
|
||
]
|
||
if not diffs:
|
||
return default
|
||
diffs.sort()
|
||
# 取 3 位小数：与 frame-extract 的 round(秒,3) 时间戳精度一致，避免浮点漂移。
|
||
return round(diffs[len(diffs) // 2], 3)
|
||
|
||
|
||
def _assemble_srt(
|
||
kept: list[tuple[float, float, str]],
|
||
interval_seconds: float,
|
||
) -> list[str]:
|
||
"""把 (起始帧时间, 最后可见帧时间, 文本) 序列组装为 SRT 行列表。
|
||
|
||
每条字幕的结束时间 = 最后可见帧时间 + 采样间隔：字幕在最后一个被识别
|
||
到的帧之后的一个采样间隔内消失，与视频烧录时间对齐；两段字幕之间的
|
||
空白段（无字幕帧）不再被并入前一条字幕。
|
||
"""
|
||
lines: list[str] = []
|
||
for index, (start, last_seen, text) in enumerate(kept):
|
||
end = last_seen + interval_seconds
|
||
lines.extend(
|
||
[
|
||
str(index + 1),
|
||
f"{format_timestamp(start)} --> {format_timestamp(end)}",
|
||
text,
|
||
"",
|
||
]
|
||
)
|
||
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")
|
||
if not manifest_uri:
|
||
return InvokeResponse(status="failed", error="frames_manifest is required")
|
||
manifest_path = Path(manifest_uri)
|
||
if not manifest_path.is_file():
|
||
return InvokeResponse(status="failed", error="frames manifest not found")
|
||
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||
|
||
# 结果长度上限：超过即视为模型异常（重复循环等），该帧直接报错跳过。
|
||
max_result_chars = int(request.params.get("max_result_chars", 200))
|
||
# 采样间隔优先从帧清单时间轴推导（与 frame-extract 实际抽取间隔一致），
|
||
# 参数仅作清单退化时的兜底。
|
||
interval = _sampling_interval(
|
||
manifest, float(request.params.get("interval_seconds", 2.0))
|
||
)
|
||
# 透传给 vlm-ocr 的参数（仅传已提供的，避免覆盖其默认值）。
|
||
vlm_params = {
|
||
key: request.params.get(key)
|
||
for key in (
|
||
"model", "ollama_host", "prompt", "timeout_seconds", "keep_alive",
|
||
"temperature", "repeat_penalty", "num_predict",
|
||
)
|
||
if request.params.get(key) is not None
|
||
}
|
||
|
||
# 单帧 OCR：并行池的工作函数，返回该帧识别文本（失败/空/超长均返回空串）。
|
||
def ocr_frame(payload) -> str:
|
||
index, item = payload
|
||
response = registry.invoke(
|
||
"vlm-ocr",
|
||
InvokeRequest(
|
||
run_id=request.run_id,
|
||
node_instance_id="",
|
||
inputs={"image_uri": str(item["image_uri"])},
|
||
params=vlm_params,
|
||
output_dir=str(Path(request.output_dir) / "ocr_frames" / f"{index:04d}"),
|
||
),
|
||
)
|
||
if response.status != "completed":
|
||
# 单帧失败不中断整体，跳过该帧继续汇总。
|
||
logger.warning("帧 %d OCR 失败，跳过: %s", index, response.error)
|
||
return ""
|
||
logger.info("帧 %d/%d OCR 完成: %r", index + 1, len(manifest), response.outputs.get("text"))
|
||
text = str(response.outputs.get("text", "")).strip()
|
||
if not text:
|
||
return ""
|
||
# 超长输出视为模型异常（重复循环等），直接报错并跳过该帧。
|
||
if len(text) > max_result_chars:
|
||
logger.warning(
|
||
"帧 %d OCR 输出超长(%d > %d)，跳过: %r",
|
||
index, len(text), max_result_chars, text[:60],
|
||
)
|
||
return ""
|
||
return text
|
||
|
||
# 进度日志：打印已识别帧数、总数与平均处理速度（帧/s）。
|
||
def log_progress(done: int, total: int, rate: float) -> None:
|
||
logger.info("OCR 进度 %d/%d 帧 (%.1f 帧/s)", done, total, rate)
|
||
|
||
# 自适应并发调用 vlm-ocr：10s 窗口内平均响应 < 0.3s 则加 1 线程（上限
|
||
# pool_max_workers），> pool_slow_threshold 则减 1 线程（下限 1），
|
||
# 按实测负载弹性伸缩，避免盲目并发压垮本地 Ollama。
|
||
pool = AdaptiveThreadPool(
|
||
worker=ocr_frame,
|
||
on_progress=log_progress,
|
||
min_workers=int(request.params.get("pool_min_workers", 1)),
|
||
max_workers=int(request.params.get("pool_max_workers", 16)),
|
||
window_seconds=float(request.params.get("pool_window_seconds", 10.0)),
|
||
fast_threshold=float(request.params.get("pool_fast_threshold", 0.3)),
|
||
slow_threshold=float(request.params.get("pool_slow_threshold", 1.0)),
|
||
)
|
||
texts = pool.map(list(enumerate(manifest)))
|
||
|
||
# 按帧顺序合并连续相同字幕（与重组装旧数据共用 _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"
|
||
output_path.write_text("\n".join(_assemble_srt(kept, interval)), encoding="utf-8")
|
||
logger.info("字幕汇总完成: %d 条", len(kept))
|
||
return InvokeResponse(
|
||
status="completed",
|
||
outputs={"srt_uri": str(output_path), "count": len(kept)},
|
||
)
|