tps-dashboard/cca/cca_pipeline.py

# -*- coding: utf-8 -*-
"""
CCA 唱词助手 — 脚本版流水线入口

用法:
    # 完整流程: A稿热词 + ASR + AI校对 + AI折行 → 5个SRT
    python cca_pipeline.py --audio data/xxx.mp3 --script data/xxx.docx

    # 手动指定热词
    python cca_pipeline.py --audio data/xxx.mp3 --hotwords "热词1|热词2"

    # 跳过ASR，从缓存处理（调试折行/校对用）
    python cca_pipeline.py --asr-cache output/asr_raw.json --script data/xxx.docx

    # 省token模式（不用AI折行和校对）
    python cca_pipeline.py --asr-cache output/asr_raw.json --no-ai

流水线:
    A稿 → 热词提取 → 音频+热词 → 讯飞ASR → AI校对 → 节目结构切分 → AI折行 → 5个SRT
"""

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent / "src"))

from asr_client import transcribe, parse_result
from line_breaker import process_sentences
from ai_line_breaker import process_sentences_with_ai
from srt_writer import write_srt, ms_to_srt_time
from segment_splitter import split_into_segments
from hotword_extractor import extract_hotwords
from ai_proofreader import proofread_batch


def main():
    parser = argparse.ArgumentParser(description="CCA 唱词助手 - 自动生成拍词 SRT")
    parser.add_argument("--audio", type=str, help="音频文件路径 (mp3/wav)")
    parser.add_argument("--script", type=str, help="A稿路径 (.docx/.txt)，用于热词提取和AI校对")
    parser.add_argument("--hotwords", type=str, default="", help="手动指定热词，用|分隔（与--script可叠加）")
    parser.add_argument("--asr-cache", type=str, help="ASR 缓存 JSON 路径（跳过ASR调用）")
    parser.add_argument("--output-dir", type=str, default="output", help="输出目录 (默认: output/)")
    parser.add_argument("--no-ai", action="store_true", help="不使用AI折行和校对（省token）")
    parser.add_argument("--no-proofread", action="store_true", help="跳过AI校对（只省校对的token）")
    args = parser.parse_args()

    if not args.audio and not args.asr_cache:
        parser.error("必须提供 --audio 或 --asr-cache")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    use_ai = not args.no_ai

    # ====== Step 0: 热词提取（从A稿）======
    hot_words = []
    script_text = ""

    if args.hotwords:
        hot_words = [w.strip() for w in args.hotwords.split("|") if w.strip()]

    if args.script:
        print(f"[流水线] 从A稿提取热词: {args.script}")
        script_hot = extract_hotwords(args.script, use_ai=use_ai)
        # 合并手动热词和A稿热词
        seen = set(hot_words)
        for w in script_hot:
            if w not in seen:
                hot_words.append(w)
                seen.add(w)

        # 读取A稿全文（校对用）
        ext = os.path.splitext(args.script)[1].lower()
        if ext == ".docx":
            from hotword_extractor import read_docx_text
            script_text = read_docx_text(args.script)
        else:
            from hotword_extractor import read_text_file
            script_text = read_text_file(args.script)

    if hot_words:
        print(f"[流水线] 热词共 {len(hot_words)} 个: {', '.join(hot_words[:10])}...")
        # 保存热词列表
        hotwords_path = output_dir / "hotwords.txt"
        with open(hotwords_path, "w", encoding="utf-8") as f:
            f.write("|".join(hot_words))

    # ====== Step 1: ASR 转写 ======
    if args.asr_cache:
        print(f"[流水线] 从缓存加载 ASR 结果: {args.asr_cache}")
        with open(args.asr_cache, "r", encoding="utf-8") as f:
            raw_json = f.read()
        sentences = parse_result(raw_json)
    else:
        print(f"[流水线] 开始 ASR 转写: {args.audio}")
        sentences, raw_json = transcribe(args.audio, hot_words=hot_words if hot_words else None)

        cache_path = output_dir / "asr_raw.json"
        with open(cache_path, "w", encoding="utf-8") as f:
            f.write(raw_json)
        print(f"[流水线] ASR 原始结果已缓存: {cache_path}")

    print(f"[流水线] ASR 共 {len(sentences)} 句")

    # ====== Step 2: AI 校对 ======
    if use_ai and not args.no_proofread and script_text:
        print("[流水线] AI 校对中 (DeepSeek)...")
        sentences = proofread_batch(sentences, script_text)
    elif not script_text and not args.no_proofread and use_ai:
        print("[流水线] 未提供A稿(--script)，跳过AI校对")

    # ====== Step 3: 节目结构切分 ======
    print("[流水线] 切分节目结构...")
    segments = split_into_segments(sentences)
    print(f"[流水线] 切分结果: {[name for name, _ in segments]}")

    # ====== Step 4: 折行 + 生成 SRT ======
    if use_ai:
        print("[流水线] 使用 AI 折行 (DeepSeek)...")
    else:
        print("[流水线] 使用机械折行规则...")

    for seg_name, seg_sentences in segments:
        if not seg_sentences:
            print(f"  [{seg_name}] 空段，跳过")
            continue

        if use_ai:
            subtitle_lines = process_sentences_with_ai(seg_sentences)
        else:
            subtitle_lines = process_sentences(seg_sentences)

        seg_offset = subtitle_lines[0][0] if subtitle_lines else 0
        seg_end = subtitle_lines[-1][1] if subtitle_lines else 0
        seg_duration = seg_end - seg_offset

        srt_filename = f"{seg_name}.srt"
        srt_path = output_dir / srt_filename
        write_srt(subtitle_lines, str(srt_path))  # 绝对时间戳，方便在时间线上对位
        print(f"  [{seg_name}] 时长 {ms_to_srt_time(seg_duration)}, 在音频中的位置: {ms_to_srt_time(seg_offset)} ~ {ms_to_srt_time(seg_end)}")

    # ====== 完成 ======
    print(f"\n[流水线] 完成! 输出目录: {output_dir}/")
    print("[流水线] 生成的文件:")
    for f in sorted(output_dir.glob("*.srt")):
        print(f"  {f.name}")


if __name__ == "__main__":
    main()