diff --git a/app/tree/subtitle.py b/app/tree/subtitle.py new file mode 100644 index 0000000..081e85c --- /dev/null +++ b/app/tree/subtitle.py @@ -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 标签（如 、） + - 跳过纯音乐符号行（仅含空白和 ♪♫） + - 多行字幕合并为单行（空格连接） + - 跳过格式异常的块（容错处理） + + 参数: + 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), + ) diff --git a/tests/unit/test_subtitle.py b/tests/unit/test_subtitle.py new file mode 100644 index 0000000..7c01b57 --- /dev/null +++ b/tests/unit/test_subtitle.py @@ -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 +This is italic 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