278 lines
13 KiB
Python
278 lines
13 KiB
Python
"""字幕 OCR 汇总节点。
|
||
|
||
读取 frame-extract 产出的 frames.json，逐帧调用 vlm-ocr 节点识别字幕文字；
|
||
过滤无文字帧的垃圾输出（glm-ocr 在空帧上会输出无用文字），折叠模型重复
|
||
循环输出，合并连续相同的字幕（记录最后可见帧时间），最终组装为带时间轴的
|
||
SRT 基准数据：每条字幕消失时间 = 最后可见帧时间 + 采样间隔（间隔从帧清单
|
||
时间轴推导），与视频烧录时间对齐。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import threading
|
||
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")
|
||
|
||
# 节点级断点存档的写锁：多线程 OCR 并发完成时串行化追加写，避免行交错。
|
||
_partial_lock = threading.Lock()
|
||
|
||
# 断点存档文件名：每行 {"frame": 帧序号(0-based), "text": 该帧OCR文本}。
|
||
# OCR 每帧完成后立即追加一行；进程重启后从存档恢复已处理帧，只对未处理
|
||
# 帧重新调用 vlm-ocr——2 小时视频级任务中断后不浪费已完成的帧。
|
||
_PARTIAL_NAME = "ocr_partial.jsonl"
|
||
|
||
# 暂停信号文件名：位于 run 根目录（<storage>/runs/<run_id>/paused.flag），
|
||
# 由暂停接口写入、继续/重试/删除时清除；节点逐帧检查，存在即中止。
|
||
_PAUSE_FLAG = "paused.flag"
|
||
|
||
|
||
class PauseRequested(Exception):
|
||
"""节点内暂停信号：OCR 检测到任务被暂停后抛出，由调度器保持 PAUSED。
|
||
|
||
不把暂停误报为 FAILED：调度器捕获异常时若任务状态已是 PAUSED，
|
||
则保持暂停等待用户 resume，从断点存档继续未处理帧。
|
||
"""
|
||
|
||
|
||
def _load_partial(output_dir: Path) -> dict[int, str]:
|
||
"""读取节点级断点存档，返回 {帧序号: OCR 文本}；无存档时返回空字典。"""
|
||
path = output_dir / _PARTIAL_NAME
|
||
if not path.is_file():
|
||
return {}
|
||
result: dict[int, str] = {}
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
if not line.strip():
|
||
continue
|
||
try:
|
||
item = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
# 进程被杀时可能残留半行写入：跳过该行，对应帧视为未处理。
|
||
continue
|
||
result[int(item["frame"])] = str(item["text"])
|
||
return result
|
||
|
||
|
||
def _format_eta(seconds: float) -> str:
|
||
"""把剩余秒数格式化为可读的预计完成时间（如 34分13秒 / 2小时05分）。"""
|
||
total = max(0, int(seconds))
|
||
hours, remainder = divmod(total, 3600)
|
||
minutes, secs = divmod(remainder, 60)
|
||
if hours:
|
||
return f"{hours}小时{minutes:02d}分"
|
||
if minutes:
|
||
return f"{minutes}分{secs:02d}秒"
|
||
return f"{secs}秒"
|
||
|
||
|
||
def _eta_suffix(done: int, total: int, rate: float) -> str:
|
||
"""根据当前处理速度计算剩余时间后缀（供进度日志追加）。
|
||
|
||
剩余时间 = 剩余帧数 / 当前速度；速度为 0（刚开始或耗时不可测）时
|
||
返回空串，进度提示不显示 ETA。
|
||
"""
|
||
if rate <= 0:
|
||
return ""
|
||
return f", 预计剩余 {_format_eta((total - done) / rate)}"
|
||
# 默认垃圾词：无文字帧的模型输出可能反复出现这些词。
|
||
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
|
||
}
|
||
|
||
# 节点级断点：读取已处理帧存档，只对未处理帧调用 vlm-ocr。
|
||
output_dir = Path(request.output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
partial_path = output_dir / _PARTIAL_NAME
|
||
partial_texts = _load_partial(output_dir)
|
||
pending = [(i, item) for i, item in enumerate(manifest) if i not in partial_texts]
|
||
if partial_texts:
|
||
logger.info(
|
||
"检测到节点级断点：%d/%d 帧已处理，本次只处理剩余 %d 帧",
|
||
len(partial_texts), len(manifest), len(pending),
|
||
)
|
||
|
||
# 单帧 OCR：并行池的工作函数，返回该帧识别文本（失败/空/超长均返回空串）。
|
||
# 每帧无论结果如何都把 {frame, text} 追加到断点存档，重启后不再重跑该帧。
|
||
def ocr_frame(payload) -> str:
|
||
index, item = payload
|
||
# 暂停检查：调度器置 PAUSED 并向 run 根写入 paused.flag 后，工作线程
|
||
# 立即中止（不 OCR、不写存档，该帧恢复时重跑），让 map 快速结束。
|
||
if (Path(request.output_dir).parent.parent / _PAUSE_FLAG).exists():
|
||
# 抑制后续进度回调：队列中剩余大量帧会快速失败退出，避免逐项
|
||
# 打印"OCR 进度"导致日志井喷。
|
||
pool.cancel()
|
||
raise PauseRequested(f"OCR 被暂停（run {request.run_id}）")
|
||
text = ""
|
||
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)
|
||
else:
|
||
logger.info("帧 %d/%d OCR 完成: %r", index + 1, len(manifest), response.outputs.get("text"))
|
||
text = str(response.outputs.get("text", "")).strip()
|
||
if len(text) > max_result_chars:
|
||
# 超长输出视为模型异常（重复循环等），直接报错并跳过该帧。
|
||
logger.warning(
|
||
"帧 %d OCR 输出超长(%d > %d)，跳过: %r",
|
||
index, len(text), max_result_chars, text[:60],
|
||
)
|
||
text = ""
|
||
# 断点存档：成功/失败/空串都记录"已处理"，恢复时保持一致行为。
|
||
with _partial_lock:
|
||
with partial_path.open("a", encoding="utf-8") as fh:
|
||
fh.write(json.dumps({"frame": index, "text": text}, ensure_ascii=False) + "\n")
|
||
return text
|
||
|
||
if pending:
|
||
# 进度日志：打印已识别帧数、总数、平均处理速度（帧/s）、最近窗口平均
|
||
# 单帧耗时、当前线程数与**预计剩余完成时间**（剩余帧/当前速度）——
|
||
# 便于判断多线程是否因单帧处理过慢而未启用（窗口平均响应 ≥
|
||
# fast_threshold 时自适应池不会扩容）以及整体还需要多久。
|
||
def log_progress(done: int, total: int, rate: float, avg_time: float, workers: int) -> None:
|
||
logger.info(
|
||
"OCR 进度 %d/%d 帧 (%.1f 帧/s, 平均 %.2fs/帧, 线程 %d/%d%s)",
|
||
done, total, rate, avg_time, workers, pool.max_workers,
|
||
_eta_suffix(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)),
|
||
)
|
||
pending_texts = pool.map(pending)
|
||
# 任一工作线程检测到暂停信号即整体中止：已写盘的断点存档保留，
|
||
# resume 后从剩余帧续跑；以 failed 返回让调度器保持 PAUSED（不误报失败）。
|
||
if any(isinstance(text, PauseRequested) for text in pending_texts):
|
||
return InvokeResponse(status="failed", error=f"OCR 被暂停（run {request.run_id}）")
|
||
# 新增结果按帧号归位，与断点存档合并成完整帧序文本列表。
|
||
new_by_index = {i: t for (i, _item), t in zip(pending, pending_texts)}
|
||
else:
|
||
new_by_index = {}
|
||
|
||
# 完整帧序文本：优先取存档，其次取本次新增（未处理帧必有其一条目）。
|
||
texts = [
|
||
partial_texts.get(i, new_by_index.get(i, ""))
|
||
for i in range(len(manifest))
|
||
]
|
||
# 按帧顺序合并连续相同字幕（与重组装旧数据共用 _merge_kept）。
|
||
kept = _merge_kept(manifest, texts)
|
||
# output_dir 已在断点初始化时创建（mkdir 幂等），此处直接使用。
|
||
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)},
|
||
)
|