feat(tree): 字幕模块 — SRT 解析 + 完整性检查 + Voronoi 分配
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,308 @@
|
||||
"""字幕模块：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),
|
||||
)
|
||||
@@ -0,0 +1,118 @@
|
||||
"""字幕模块单元测试。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.tree.subtitle import (
|
||||
SRTEntry,
|
||||
assign_subtitles_voronoi,
|
||||
check_subtitle_completeness,
|
||||
extract_subtitle_for_range,
|
||||
parse_srt,
|
||||
)
|
||||
|
||||
_SAMPLE_SRT = """\
|
||||
1
|
||||
00:00:01,000 --> 00:00:03,500
|
||||
Hello world.
|
||||
|
||||
2
|
||||
00:00:05,000 --> 00:00:08,000
|
||||
<i>This is italic</i> text.
|
||||
|
||||
3
|
||||
00:00:10,000 --> 00:00:12,000
|
||||
Final line.
|
||||
"""
|
||||
|
||||
|
||||
class TestParseSrt:
|
||||
def test_basic_parse(self, tmp_path):
|
||||
srt_file = tmp_path / "test.srt"
|
||||
srt_file.write_text(_SAMPLE_SRT, encoding="utf-8")
|
||||
entries = parse_srt(str(srt_file))
|
||||
assert len(entries) == 3
|
||||
assert entries[0] == SRTEntry(start=1.0, end=3.5, text="Hello world.")
|
||||
assert entries[1].text == "This is italic text."
|
||||
|
||||
def test_empty_srt(self, tmp_path):
|
||||
srt_file = tmp_path / "empty.srt"
|
||||
srt_file.write_text("", encoding="utf-8")
|
||||
entries = parse_srt(str(srt_file))
|
||||
assert entries == []
|
||||
|
||||
def test_malformed_srt_skips_bad_blocks(self, tmp_path):
|
||||
bad_srt = "garbage\n\n1\n00:00:01,000 --> 00:00:02,000\nGood line.\n"
|
||||
srt_file = tmp_path / "bad.srt"
|
||||
srt_file.write_text(bad_srt, encoding="utf-8")
|
||||
entries = parse_srt(str(srt_file))
|
||||
assert len(entries) == 1
|
||||
assert entries[0].text == "Good line."
|
||||
|
||||
|
||||
class TestCompletenessCheck:
|
||||
def test_good_coverage(self):
|
||||
entries = [SRTEntry(0.0, 5.0, "a"), SRTEntry(5.0, 10.0, "b")]
|
||||
report = check_subtitle_completeness(entries, duration=10.0, min_coverage=0.5)
|
||||
assert report.usable is True
|
||||
assert report.coverage_ratio >= 0.5
|
||||
|
||||
def test_poor_coverage(self):
|
||||
entries = [SRTEntry(0.0, 1.0, "short")]
|
||||
report = check_subtitle_completeness(entries, duration=100.0, min_coverage=0.3)
|
||||
assert report.usable is False
|
||||
|
||||
def test_max_gap(self):
|
||||
entries = [SRTEntry(0.0, 1.0, "a"), SRTEntry(50.0, 51.0, "b")]
|
||||
report = check_subtitle_completeness(entries, duration=60.0)
|
||||
assert report.max_gap_sec >= 49.0
|
||||
|
||||
|
||||
class TestExtractForRange:
|
||||
def test_overlap(self):
|
||||
entries = [
|
||||
SRTEntry(0.0, 5.0, "first"),
|
||||
SRTEntry(4.0, 8.0, "second"),
|
||||
SRTEntry(10.0, 12.0, "third"),
|
||||
]
|
||||
text = extract_subtitle_for_range(entries, (3.0, 9.0))
|
||||
assert "first" in text
|
||||
assert "second" in text
|
||||
assert "third" not in text
|
||||
|
||||
|
||||
class TestVoronoiAssign:
|
||||
def test_assigns_to_l3_nodes(self):
|
||||
from app.tree.index import (
|
||||
IndexMeta,
|
||||
L1Card,
|
||||
L1Node,
|
||||
L2Card,
|
||||
L2Node,
|
||||
L3Card,
|
||||
L3Node,
|
||||
TreeIndex,
|
||||
)
|
||||
|
||||
l3_0 = L3Node(id="l1_0_l2_0_l3_0", card=L3Card("desc0", [], [], [], "", {}), timestamp=2.0)
|
||||
l3_1 = L3Node(id="l1_0_l2_0_l3_1", card=L3Card("desc1", [], [], [], "", {}), timestamp=6.0)
|
||||
l2 = L2Node(
|
||||
id="l1_0_l2_0",
|
||||
card=L2Card("evt", [], [], [], [], "", None),
|
||||
time_range=(0.0, 10.0),
|
||||
children=[l3_0, l3_1],
|
||||
)
|
||||
l1 = L1Node(
|
||||
id="l1_0",
|
||||
card=L1Card("scene", "", [], [], [], [], ""),
|
||||
time_range=(0.0, 10.0),
|
||||
children=[l2],
|
||||
)
|
||||
index = TreeIndex(metadata=IndexMeta("/test.mp4", "video"), roots=[l1])
|
||||
|
||||
entries = [SRTEntry(1.0, 3.0, "hello"), SRTEntry(5.0, 7.0, "world")]
|
||||
assign_subtitles_voronoi(index, entries)
|
||||
|
||||
assert l3_0.subtitle is not None
|
||||
assert "hello" in l3_0.subtitle
|
||||
assert l3_1.subtitle is not None
|
||||
assert "world" in l3_1.subtitle
|
||||
Reference in New Issue
Block a user