tps-dashboard/cca/src/segment_splitter.py

# -*- coding: utf-8 -*-
"""
节目结构切分器 — 将 ASR 结果按节目结构拆分为 5 段

结构: 导视 + 正片(3段) + 下期预告
标志词:
- 导视结束 / 正片开始: "各位观众你们好" 或 "我是主持人蓝皓"
- 正片结束: "好了观众朋友们" 或 "感谢您.*关注.*军事科技"
- 正片之后 = 下期预告

正片拆3段: 按时长大致均分，优先在角色转换处（speaker_id 变化）切分
"""

import re
from typing import List, Tuple, Optional

# 标志词模式
PATTERN_SHOW_START = re.compile(r"各位观众你们好|我是主持人蓝皓")
PATTERN_SHOW_END = re.compile(r"好了观众朋友们|感谢您.*关注.*军事科技|感谢您持续关注")


def find_segment_boundaries(
    sentences: List[Tuple[int, int, str, int]],
) -> Tuple[int, int]:
    """
    找到正片开始和结束的句子索引。
    返回 (show_start_idx, show_end_idx)
    - show_start_idx: "各位观众你们好"所在句子的索引
    - show_end_idx: "好了观众朋友们"所在句子的索引
    """
    show_start_idx = 0
    show_end_idx = len(sentences) - 1

    for i, (_, _, text, _) in enumerate(sentences):
        if PATTERN_SHOW_START.search(text):
            show_start_idx = i
            break

    for i in range(len(sentences) - 1, -1, -1):
        _, _, text, _ = sentences[i]
        if PATTERN_SHOW_END.search(text):
            show_end_idx = i
            break

    return show_start_idx, show_end_idx


def split_show_into_three(
    sentences: List[Tuple[int, int, str, int]],
    start_idx: int,
    end_idx: int,
) -> Tuple[int, int]:
    """
    将正片（start_idx 到 end_idx）拆成 3 段。
    返回两个切分点索引 (split1_idx, split2_idx)

    策略: 按时长三等分，然后在附近找 speaker_id 变化的位置。
    """
    if end_idx - start_idx < 6:
        # 太短了，均分
        third = (end_idx - start_idx) // 3
        return start_idx + third, start_idx + 2 * third

    show_start_ms = sentences[start_idx][0]
    show_end_ms = sentences[end_idx][1]
    total_duration = show_end_ms - show_start_ms

    target1_ms = show_start_ms + total_duration // 3
    target2_ms = show_start_ms + 2 * total_duration // 3

    split1 = _find_best_split(sentences, start_idx, end_idx, target1_ms)
    split2 = _find_best_split(sentences, split1 + 1, end_idx, target2_ms)

    # 确保 split2 > split1
    if split2 <= split1:
        split2 = split1 + (end_idx - split1) // 2

    return split1, split2


def _find_best_split(
    sentences: List[Tuple[int, int, str, int]],
    range_start: int,
    range_end: int,
    target_ms: int,
    search_window: int = 15,
) -> int:
    """
    在 target_ms 附近（±search_window 句）找最佳切分点。
    优先找 speaker_id 变化的位置，其次找 >2秒 空白。
    """
    # 先找到时间上最接近 target_ms 的句子
    closest_idx = range_start
    min_diff = abs(sentences[range_start][0] - target_ms)
    for i in range(range_start, min(range_end + 1, len(sentences))):
        diff = abs(sentences[i][0] - target_ms)
        if diff < min_diff:
            min_diff = diff
            closest_idx = i

    # 在附近找 speaker 变化点
    search_lo = max(range_start + 1, closest_idx - search_window)
    search_hi = min(range_end, closest_idx + search_window)

    best_idx = closest_idx
    best_score = 0

    for i in range(search_lo, search_hi):
        score = 0
        # speaker 变化加分
        if sentences[i][3] != sentences[i - 1][3] and sentences[i][3] != 0:
            score += 10
        # 空白间隔加分
        gap = sentences[i][0] - sentences[i - 1][1]
        if gap > 2000:
            score += 5
        elif gap > 1000:
            score += 2
        # 离目标越近加分
        time_diff = abs(sentences[i][0] - target_ms)
        time_score = max(0, 5 - time_diff / 10000)
        score += time_score

        if score > best_score:
            best_score = score
            best_idx = i

    return best_idx


def split_into_segments(
    sentences: List[Tuple[int, int, str, int]],
) -> List[Tuple[str, List[Tuple[int, int, str, int]]]]:
    """
    将全部 ASR 句子拆分为 5 段。
    返回: [("导视", [...]), ("正片1", [...]), ("正片2", [...]), ("正片3", [...]), ("预告", [...])]

    如果找不到标志词，则只输出单段。
    """
    if not sentences:
        return [("正片1", [])]

    show_start_idx, show_end_idx = find_segment_boundaries(sentences)

    # 导视: 0 到 show_start_idx-1
    intro = sentences[:show_start_idx] if show_start_idx > 0 else []

    # 正片: show_start_idx 到 show_end_idx
    show_sentences = sentences[show_start_idx:show_end_idx + 1]

    # 预告: show_end_idx+1 到末尾
    trailer = sentences[show_end_idx + 1:] if show_end_idx < len(sentences) - 1 else []

    # 正片拆3段
    if len(show_sentences) > 6:
        split1, split2 = split_show_into_three(sentences, show_start_idx, show_end_idx)
        # 转为相对于 show_sentences 的索引
        rel_split1 = split1 - show_start_idx
        rel_split2 = split2 - show_start_idx
        show_part1 = show_sentences[:rel_split1]
        show_part2 = show_sentences[rel_split1:rel_split2]
        show_part3 = show_sentences[rel_split2:]
    else:
        show_part1 = show_sentences
        show_part2 = []
        show_part3 = []

    segments = []
    if intro:
        segments.append(("导视", intro))
    segments.append(("正片1", show_part1))
    if show_part2:
        segments.append(("正片2", show_part2))
    if show_part3:
        segments.append(("正片3", show_part3))
    if trailer:
        segments.append(("预告", trailer))

    return segments