309 lines
9.6 KiB
Python
309 lines
9.6 KiB
Python
"""字幕模块：SRT 解析、完整性检查、时间范围提取、Voronoi 分配。
|
||
|
||
提供四个核心函数：
|
||
- parse_srt: 解析 SRT 文件为结构化条目列表
|
||
- check_subtitle_completeness: 检查字幕覆盖率与完整性
|
||
- extract_subtitle_for_range: 提取指定时间范围内的字幕文本
|
||
- assign_subtitles_voronoi: 使用 Voronoi 中点策略将字幕分配给 L3 节点
|
||
|
||
迁移来源:
|
||
- TRM4 core/tree/enhance/merge.py (parse_srt)
|
||
- TRM3 tools/generate_subtitles.py (Voronoi 逻辑)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING
|
||
|
||
from loguru import logger
|
||
|
||
if TYPE_CHECKING:
|
||
from app.tree.index import TreeIndex
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 正则表达式
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_HTML_TAG_RE = re.compile(r"<[^>]+>")
|
||
_MUSIC_ONLY_RE = re.compile(r"^[\s♪♫]*$")
|
||
_TIMECODE_RE = re.compile(r"(\d+):(\d+):(\d+)[,.](\d+)\s*-->\s*(\d+):(\d+):(\d+)[,.](\d+)")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 数据类型
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SRTEntry:
|
||
"""单条 SRT 字幕条目。
|
||
|
||
属性:
|
||
start: 开始时间（秒）。
|
||
end: 结束时间（秒）。
|
||
text: 字幕文本（已清洗 HTML 标签）。
|
||
"""
|
||
|
||
start: float
|
||
end: float
|
||
text: str
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SubtitleReport:
|
||
"""字幕完整性检查报告。
|
||
|
||
属性:
|
||
total_entries: 字幕条目总数。
|
||
coverage_ratio: SRT 覆盖时长 / 视频总时长。
|
||
max_gap_sec: 最大连续无字幕间隔（秒）。
|
||
usable: 覆盖率是否达到最低要求。
|
||
"""
|
||
|
||
total_entries: int
|
||
coverage_ratio: float
|
||
max_gap_sec: float
|
||
usable: bool
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 内部辅助
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _ts_to_seconds(h: str, m: str, s: str, ms: str) -> float:
|
||
"""SRT 时间戳组件 (HH:MM:SS,mmm) 转秒数。
|
||
|
||
参数:
|
||
h: 小时。
|
||
m: 分钟。
|
||
s: 秒。
|
||
ms: 毫秒。
|
||
|
||
返回:
|
||
浮点秒数。
|
||
"""
|
||
return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 公共 API
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def parse_srt(srt_path: str) -> list[SRTEntry]:
|
||
"""解析 SRT 字幕文件，返回结构化条目列表。
|
||
|
||
- 剥离 HTML 标签（如 <i>、<b>）
|
||
- 跳过纯音乐符号行（仅含空白和 ♪♫）
|
||
- 多行字幕合并为单行（空格连接）
|
||
- 跳过格式异常的块（容错处理）
|
||
|
||
参数:
|
||
srt_path: SRT 文件的绝对路径。
|
||
|
||
返回:
|
||
按时间顺序排列的 SRTEntry 列表；空文件或无有效条目返回空列表。
|
||
|
||
迁移来源:
|
||
TRM4 core/tree/enhance/merge.py parse_srt
|
||
TRM3 tools/generate_subtitles.py parse_srt
|
||
"""
|
||
with open(srt_path, encoding="utf-8") as f:
|
||
content = f.read()
|
||
|
||
if not content.strip():
|
||
return []
|
||
|
||
entries: list[SRTEntry] = []
|
||
blocks = re.split(r"\n\s*\n", content.strip())
|
||
|
||
for block in blocks:
|
||
lines = block.strip().split("\n")
|
||
if len(lines) < 2:
|
||
continue
|
||
|
||
# 在块内搜索时间码行（可能是第 1 行或第 2 行）
|
||
ts_match = None
|
||
ts_line_idx = -1
|
||
for i, line in enumerate(lines):
|
||
ts_match = _TIMECODE_RE.search(line)
|
||
if ts_match:
|
||
ts_line_idx = i
|
||
break
|
||
|
||
if not ts_match:
|
||
continue
|
||
|
||
groups = [int(x) for x in ts_match.groups()]
|
||
start = _ts_to_seconds(str(groups[0]), str(groups[1]), str(groups[2]), str(groups[3]))
|
||
end = _ts_to_seconds(str(groups[4]), str(groups[5]), str(groups[6]), str(groups[7]))
|
||
|
||
# 时间码行之后的所有行为字幕文本
|
||
text_lines = lines[ts_line_idx + 1 :]
|
||
raw_text = " ".join(text_lines)
|
||
clean_text = _HTML_TAG_RE.sub("", raw_text).strip()
|
||
|
||
# 跳过空文本和纯音乐符号行
|
||
if not clean_text or _MUSIC_ONLY_RE.match(clean_text):
|
||
continue
|
||
|
||
entries.append(SRTEntry(start=start, end=end, text=clean_text))
|
||
|
||
logger.debug("SRT 解析完成: {} 条有效条目, 文件={}", len(entries), srt_path)
|
||
return entries
|
||
|
||
|
||
def check_subtitle_completeness(
|
||
entries: list[SRTEntry],
|
||
duration: float,
|
||
min_coverage: float = 0.3,
|
||
) -> SubtitleReport:
|
||
"""检查字幕完整性：覆盖率、最大间隔、可用性判定。
|
||
|
||
参数:
|
||
entries: 已排序的 SRTEntry 列表。
|
||
duration: 视频总时长（秒），必须 > 0。
|
||
min_coverage: 最低可用覆盖率阈值（0~1）。
|
||
|
||
返回:
|
||
SubtitleReport 包含覆盖率、最大间隔和可用性判定。
|
||
"""
|
||
assert duration > 0, f"视频时长必须 > 0，实际={duration}"
|
||
|
||
if not entries:
|
||
return SubtitleReport(
|
||
total_entries=0,
|
||
coverage_ratio=0.0,
|
||
max_gap_sec=duration,
|
||
usable=False,
|
||
)
|
||
|
||
# 按开始时间排序
|
||
sorted_entries = sorted(entries, key=lambda e: e.start)
|
||
|
||
# 计算覆盖时长（合并重叠区间）
|
||
merged_intervals: list[tuple[float, float]] = []
|
||
for entry in sorted_entries:
|
||
if merged_intervals and entry.start <= merged_intervals[-1][1]:
|
||
# 与上一区间重叠，扩展
|
||
merged_intervals[-1] = (
|
||
merged_intervals[-1][0],
|
||
max(merged_intervals[-1][1], entry.end),
|
||
)
|
||
else:
|
||
merged_intervals.append((entry.start, entry.end))
|
||
|
||
covered = sum(end - start for start, end in merged_intervals)
|
||
coverage_ratio = min(covered / duration, 1.0)
|
||
|
||
# 计算最大间隔（包括视频开头到第一条字幕、最后一条到视频结尾）
|
||
max_gap = merged_intervals[0][0] # 视频开头到第一条字幕
|
||
for i in range(1, len(merged_intervals)):
|
||
gap = merged_intervals[i][0] - merged_intervals[i - 1][1]
|
||
max_gap = max(max_gap, gap)
|
||
# 最后一条字幕到视频结尾
|
||
max_gap = max(max_gap, duration - merged_intervals[-1][1])
|
||
|
||
return SubtitleReport(
|
||
total_entries=len(entries),
|
||
coverage_ratio=coverage_ratio,
|
||
max_gap_sec=max_gap,
|
||
usable=coverage_ratio >= min_coverage,
|
||
)
|
||
|
||
|
||
def extract_subtitle_for_range(
|
||
entries: list[SRTEntry],
|
||
time_range: tuple[float, float],
|
||
) -> str:
|
||
"""提取与指定时间范围重叠的字幕文本。
|
||
|
||
重叠判定：entry.start < range_end 且 entry.end > range_start。
|
||
|
||
参数:
|
||
entries: SRTEntry 列表。
|
||
time_range: (start, end) 时间范围（秒）。
|
||
|
||
返回:
|
||
匹配的字幕文本，多条用换行符连接；无匹配返回空字符串。
|
||
"""
|
||
range_start, range_end = time_range
|
||
matched = [
|
||
entry.text for entry in entries if entry.start < range_end and entry.end > range_start
|
||
]
|
||
return "\n".join(matched)
|
||
|
||
|
||
def assign_subtitles_voronoi(
|
||
index: TreeIndex,
|
||
entries: list[SRTEntry],
|
||
) -> None:
|
||
"""使用 Voronoi 中点策略将字幕分配给 L3 节点。
|
||
|
||
对每个 L2 节点内的 L3 子节点，按 timestamp 排序后计算 Voronoi 有效范围：
|
||
- 相邻 L3 节点之间取中点作为边界
|
||
- 首个 L3 的左边界扩展到 L2 的 time_range 起点
|
||
- 末个 L3 的右边界扩展到 L2 的 time_range 终点
|
||
|
||
然后用 extract_subtitle_for_range 提取每个 L3 有效范围内的字幕文本。
|
||
|
||
参数:
|
||
index: 树索引，包含 L1→L2→L3 嵌套结构。
|
||
entries: 已解析的 SRTEntry 列表。
|
||
|
||
副作用:
|
||
直接修改每个 L3Node.subtitle 字段。
|
||
|
||
迁移来源:
|
||
TRM3 tools/generate_subtitles.py compute_effective_ranges + assign_subtitles
|
||
"""
|
||
for l1 in index.roots:
|
||
for l2 in l1.children:
|
||
if not l2.children:
|
||
continue
|
||
|
||
# 按 timestamp 排序 L3 子节点（保留原列表引用以便赋值）
|
||
siblings = sorted(
|
||
l2.children,
|
||
key=lambda n: n.timestamp if n.timestamp is not None else 0.0,
|
||
)
|
||
|
||
# L2 的时间范围作为边界
|
||
l2_start = l2.time_range[0] if l2.time_range else 0.0
|
||
l2_end = l2.time_range[1] if l2.time_range else 0.0
|
||
|
||
for idx, l3 in enumerate(siblings):
|
||
ts = l3.timestamp if l3.timestamp is not None else 0.0
|
||
|
||
# 计算 Voronoi 有效范围
|
||
if idx == 0:
|
||
left = l2_start
|
||
else:
|
||
prev_ts = (
|
||
siblings[idx - 1].timestamp
|
||
if siblings[idx - 1].timestamp is not None
|
||
else 0.0
|
||
)
|
||
left = (prev_ts + ts) / 2.0
|
||
|
||
if idx == len(siblings) - 1:
|
||
right = l2_end
|
||
else:
|
||
next_ts = (
|
||
siblings[idx + 1].timestamp
|
||
if siblings[idx + 1].timestamp is not None
|
||
else 0.0
|
||
)
|
||
right = (ts + next_ts) / 2.0
|
||
|
||
subtitle_text = extract_subtitle_for_range(entries, (left, right))
|
||
l3.subtitle = subtitle_text if subtitle_text else None
|
||
|
||
logger.debug(
|
||
"Voronoi 字幕分配完成: {} 个 L1 节点, {} 条字幕条目",
|
||
len(index.roots),
|
||
len(entries),
|
||
)
|