extract_audio and transcribe_audio scripts

synthmind/transcribe_audio.py

@@ -0,0 +1,553 @@
+#!/usr/bin/env python3
+"""
+transcribe_audio.py - MP3 音频批量转录文字脚本
+功能:
+  - 扫描单个 MP3 文件或遍历指定目录中的所有 MP3 文件
+  - 维护 manifest.json 进度文件，支持断点续传（key 含模型维度，多模型互不覆盖）
+  - 使用系统已安装的 whisper 命令行工具转录音频
+  - 支持 MP3 切分（处理可能生成超大 txt 的音频文件）
+  - 支持含空格和中文的文件名
+  - 支持预下载模型 (--preload-models)
+  - 支持多模型对比 (--all-models)
+  - 支持自定义输出后缀 (--output-suffix)
+
+支持的模型: tiny / base / small / medium
+推荐语言:  zh (中文) / en (英文)，也接受其他 whisper 支持的语言代码
+
+用法:
+  python transcribe_audio.py <MP3文件或目录> [选项]
+
+示例:
+  # 预下载所有支持的模型
+  python transcribe_audio.py --preload-models
+
+  # 使用 medium 模型转录中文
+  python transcribe_audio.py /path/to/audio/ --model medium --language zh
+
+  # 一次跑 4 个模型对比精度（生成 file.tiny.txt / file.base.txt / ...）
+  python transcribe_audio.py /path/to/audio/ --all-models --language zh
+
+  # 自定义输出后缀
+  python transcribe_audio.py /path/to/audio.mp3 --model base --output-suffix _v1
+
+  # 超1800秒切分
+  python transcribe_audio.py /path/to/audio/ --split-size 1800
+
+  # 重新处理失败项
+  python transcribe_audio.py /path/to/audio/ --retry-failed
+
+  # 查看进度状态
+  python transcribe_audio.py /path/to/audio/ --status
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+# ─── 常量 ──────────────────────────────────────────────────────────────────────
+
+MANIFEST_NAME = "manifest.json"
+AUDIO_EXTS = {".mp3", ".m4a", ".wav", ".flac", ".ogg", ".aac"}
+
+SUPPORTED_MODELS = ["tiny", "base", "small", "medium"]
+DEFAULT_WHISPER_MODEL = "base"
+DEFAULT_SPLIT_SEGMENT_SECS = 1800   # 每片30分钟
+
+
+# ─── 日志 ──────────────────────────────────────────────────────────────────────
+
+def log(msg: str, level: str = "INFO"):
+    ts = time.strftime("%H:%M:%S")
+    prefix = {"INFO": "  ", "OK": "✅", "ERR": "❌", "WARN": "⚠️ ", "STEP": "▶ "}.get(level, "  ")
+    print(f"[{ts}] {prefix} {msg}", flush=True)
+
+
+# ─── Manifest 操作 (key 包含 model，支持同一音频多模型共存) ────────────────────
+
+def load_manifest(manifest_path: Path) -> dict:
+    if manifest_path.exists():
+        with open(manifest_path, encoding="utf-8") as f:
+            return json.load(f)
+    return {"version": 1, "files": {}}
+
+
+def save_manifest(manifest_path: Path, manifest: dict):
+    manifest_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(manifest_path, "w", encoding="utf-8") as f:
+        json.dump(manifest, f, ensure_ascii=False, indent=2)
+
+
+def transcribe_key(model: str, audio_path: str, output_suffix: str = "") -> str:
+    """
+    manifest key 由 model + output_suffix + 音频路径共同决定，
+    因此改变 --output-suffix 会形成新的 key（不会误跳过）。
+    """
+    if output_suffix:
+        return f"transcribe:{model}:{output_suffix}:{audio_path}"
+    return f"transcribe:{model}:{audio_path}"
+
+
+def get_transcribe_entry(manifest: dict, audio_path: str, model: str,
+                         output_suffix: str = "") -> dict:
+    key = transcribe_key(model, audio_path, output_suffix)
+    if key not in manifest["files"]:
+        manifest["files"][key] = {
+            "transcribed": False,
+            "txt_path": None,
+            "split_segments": [],
+            "error": None,
+            "model": model,
+            "language": None,
+            "output_suffix": output_suffix or None,
+            "updated_at": None,
+        }
+    return manifest["files"][key]
+
+
+def mark_transcribed(manifest: dict, manifest_path: Path, audio_path: str,
+                     txt_path: str, model: str, language: str = None,
+                     output_suffix: str = "", segments: list = None):
+    entry = get_transcribe_entry(manifest, audio_path, model, output_suffix)
+    entry["transcribed"] = True
+    entry["txt_path"] = txt_path
+    entry["split_segments"] = segments or []
+    entry["model"] = model
+    entry["language"] = language
+    entry["output_suffix"] = output_suffix or None
+    entry["error"] = None
+    entry["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%S")
+    save_manifest(manifest_path, manifest)
+
+
+def mark_transcribe_failed(manifest: dict, manifest_path: Path,
+                           audio_path: str, model: str, error: str,
+                           output_suffix: str = ""):
+    entry = get_transcribe_entry(manifest, audio_path, model, output_suffix)
+    entry["transcribed"] = False
+    entry["error"] = error
+    entry["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%S")
+    save_manifest(manifest_path, manifest)
+
+
+# ─── 模型预下载 ────────────────────────────────────────────────────────────────
+
+def preload_models(models: list) -> bool:
+    """
+    调用 whisper Python API 预下载模型到 ~/.cache/whisper/。
+    比调用 whisper CLI 更快（无需额外音频输入）。
+    """
+    try:
+        import whisper as whisper_lib
+    except ImportError:
+        log("openai-whisper Python 库未安装，无法预下载", "ERR")
+        log("请安装：pip install openai-whisper", "ERR")
+        return False
+
+    log(f"开始预下载 {len(models)} 个模型：{', '.join(models)}", "STEP")
+    for m in models:
+        log(f"  下载 {m}...")
+        t0 = time.time()
+        try:
+            whisper_lib.load_model(m)
+            log(f"  {m} 就绪 ({time.time() - t0:.1f}s)", "OK")
+        except Exception as e:
+            log(f"  {m} 下载失败: {e}", "ERR")
+            return False
+    log(f"所有模型下载完成 ✅", "OK")
+    return True
+
+
+# ─── Whisper / FFmpeg 工具函数 ─────────────────────────────────────────────────
+
+def check_whisper() -> bool:
+    try:
+        subprocess.run(["whisper", "--help"], capture_output=True, check=False)
+        return True
+    except FileNotFoundError:
+        log("未找到 whisper 命令，请确认已安装：pip install openai-whisper", "ERR")
+        return False
+
+
+def check_ffmpeg() -> bool:
+    try:
+        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return False
+
+
+def get_audio_duration(audio_path: str) -> float:
+    try:
+        r = subprocess.run(
+            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
+             "-of", "default=noprint_wrappers=1:nokey=1", audio_path],
+            capture_output=True, text=True
+        )
+        return float(r.stdout.strip())
+    except Exception:
+        return 0.0
+
+
+def split_audio(audio_path: str, output_dir: str, segment_secs: int) -> list:
+    """
+    将 MP3 切分为多个片段，返回片段路径列表。
+    使用 ffmpeg segment 模式，避免重编码（codec copy）。
+    """
+    stem = Path(audio_path).stem
+    ext = Path(audio_path).suffix
+    out_pattern = os.path.join(output_dir, f"{stem}_seg%03d{ext}")
+
+    cmd = [
+        "ffmpeg", "-y", "-i", audio_path,
+        "-f", "segment",
+        "-segment_time", str(segment_secs),
+        "-c", "copy",
+        "-reset_timestamps", "1",
+        out_pattern
+    ]
+    r = subprocess.run(cmd, capture_output=True)
+    if r.returncode != 0:
+        return []
+
+    return sorted(str(p) for p in Path(output_dir).glob(f"{stem}_seg*{ext}"))
+
+
+def run_whisper(audio_path: str, output_dir: str, model: str,
+                language: str = None) -> tuple:
+    """
+    调用 whisper 命令转录音频，返回 (success: bool, txt_path_or_err: str)。
+    whisper 会自动在 output_dir 生成 <stem>.txt 等文件。
+    注意: 多次调用同 output_dir 且相同 stem 时会覆盖，需由调用方即时重命名。
+    """
+    cmd = [
+        "whisper", audio_path,
+        "--model", model,
+        "--output_dir", output_dir,
+        "--output_format", "txt",
+        "--verbose", "False",
+    ]
+    if language:
+        cmd += ["--language", language]
+
+    r = subprocess.run(cmd, capture_output=True, text=True)
+
+    stem = Path(audio_path).stem
+    txt_path = os.path.join(output_dir, f"{stem}.txt")
+
+    if r.returncode == 0 and os.path.exists(txt_path):
+        return True, txt_path
+
+    err = r.stderr[-500:] if r.stderr else "unknown error"
+    return False, err
+
+
+def merge_txt_files(txt_files: list, merged_path: str) -> bool:
+    """将多个转录片段 txt 按顺序合并为一个完整文件"""
+    try:
+        with open(merged_path, "w", encoding="utf-8") as fout:
+            for i, tf in enumerate(txt_files):
+                if i > 0:
+                    fout.write("\n\n")
+                with open(tf, encoding="utf-8") as fin:
+                    fout.write(fin.read().strip())
+        return True
+    except Exception:
+        return False
+
+
+# ─── 文件扫描 ──────────────────────────────────────────────────────────────────
+
+def scan_audios(source: str) -> list:
+    p = Path(source)
+    if p.is_file():
+        if p.suffix.lower() in AUDIO_EXTS:
+            return [str(p.resolve())]
+        else:
+            log(f"不支持的音频格式: {p.suffix}", "ERR")
+            return []
+    elif p.is_dir():
+        return sorted(str(f.resolve()) for f in p.rglob("*")
+                      if f.suffix.lower() in AUDIO_EXTS)
+    else:
+        log(f"路径不存在: {source}", "ERR")
+        return []
+
+
+def txt_exists(txt_path: str) -> bool:
+    return bool(txt_path) and os.path.exists(txt_path) and os.path.getsize(txt_path) > 0
+
+
+# ─── 输出文件命名 ──────────────────────────────────────────────────────────────
+
+def compute_target_txt(out_dir: Path, stem: str, model: str,
+                       output_suffix: str, tag_model: bool) -> str:
+    """
+    生成最终 txt 路径。命名规则：
+    - tag_model=True (多模型模式):     <stem>.<model>.txt
+    - output_suffix 非空:              <stem><output_suffix>.txt
+    - 都无:                            <stem>.txt (默认)
+    """
+    if tag_model:
+        return str(out_dir / f"{stem}.{model}.txt")
+    if output_suffix:
+        return str(out_dir / f"{stem}{output_suffix}.txt")
+    return str(out_dir / f"{stem}.txt")
+
+
+# ─── 状态展示 ──────────────────────────────────────────────────────────────────
+
+def show_status(manifest: dict, all_audios: list):
+    print("\n" + "=" * 70)
+    print("📊 音频转录进度状态（按 模型 x 音频 展示）")
+    print("=" * 70)
+
+    entries = {k: v for k, v in manifest["files"].items() if k.startswith("transcribe:")}
+    done = sum(1 for v in entries.values() if v.get("transcribed"))
+    failed = sum(1 for v in entries.values() if v.get("error"))
+    print(f"  Manifest 记录: {len(entries)}  已完成: {done}  失败: {failed}")
+    print("=" * 70)
+
+    for a in all_audios:
+        print(f"\n  📄 {Path(a).name}")
+        matching = [(k, v) for k, v in entries.items() if k.endswith(f":{a}")]
+        if not matching:
+            print(f"     ⏳ 无记录（未转录）")
+            continue
+        for key, entry in matching:
+            m = entry.get("model", "?")
+            suffix = entry.get("output_suffix") or ""
+            suffix_tag = f" suffix={suffix}" if suffix else ""
+            if entry.get("transcribed"):
+                size = ""
+                if entry.get("txt_path") and os.path.exists(entry["txt_path"]):
+                    size = f"  {os.path.getsize(entry['txt_path']) // 1024}KB"
+                lang = entry.get("language") or "auto"
+                txt_name = Path(entry['txt_path']).name if entry.get('txt_path') else '?'
+                print(f"     ✅ [{m:6s}] lang={lang}{suffix_tag}{size}  → {txt_name}")
+            elif entry.get("error"):
+                print(f"     ❌ [{m:6s}]{suffix_tag} {entry['error'][:60]}")
+    print()
+
+
+# ─── 核心：单个音频 + 单个模型 转录 ─────────────────────────────────────────────
+
+def transcribe_one(audio: str, model: str, out_dir: Path, args,
+                   manifest: dict, manifest_path: Path,
+                   has_ffmpeg: bool, tag_model: bool) -> str:
+    """
+    返回状态字符串：'ok' / 'skip' / 'skip-failed' / 'error:<msg>'
+    """
+    stem = Path(audio).stem
+    final_txt = compute_target_txt(out_dir, stem, model, args.output_suffix, tag_model)
+
+    # 多模型模式用 .<model>.txt 命名，output_suffix 在这种情况下强制为空
+    effective_suffix = "" if tag_model else args.output_suffix
+    entry = get_transcribe_entry(manifest, audio, model, effective_suffix)
+    already_done = entry.get("transcribed") and txt_exists(entry.get("txt_path"))
+    is_failed = bool(entry.get("error"))
+
+    if already_done:
+        return "skip"
+    if is_failed and not args.retry_failed:
+        return "skip-failed"
+
+    try:
+        duration = get_audio_duration(audio) if has_ffmpeg else 0
+        need_split = args.split_size and duration > 0 and duration > args.split_size
+
+        if need_split:
+            log(f"    音频时长 {duration:.0f}s > {args.split_size}s，启动切分转录")
+            seg_dir = out_dir / f"{stem}_segments"
+            seg_dir.mkdir(exist_ok=True)
+
+            segments = split_audio(audio, str(seg_dir), args.split_segment)
+            if not segments:
+                raise RuntimeError("音频切分失败")
+
+            # 每个模型独立子目录，避免片段 txt 互相覆盖
+            seg_txt_dir = seg_dir / model
+            seg_txt_dir.mkdir(exist_ok=True)
+            log(f"    切分为 {len(segments)} 个片段，逐段转录中...")
+
+            seg_txts = []
+            for j, seg in enumerate(segments, 1):
+                ok, result = run_whisper(seg, str(seg_txt_dir), model, args.language)
+                if ok:
+                    seg_txts.append(result)
+                    log(f"      [{j}/{len(segments)}] {Path(seg).name} → done", "OK")
+                else:
+                    raise RuntimeError(f"片段 {Path(seg).name} 转录失败: {result}")
+
+            if not merge_txt_files(seg_txts, final_txt):
+                raise RuntimeError("合并转录片段失败")
+
+            size = os.path.getsize(final_txt)
+            log(f"    合并 → {Path(final_txt).name} ({size // 1024} KB)", "OK")
+            mark_transcribed(manifest, manifest_path, audio, final_txt,
+                             model, language=args.language,
+                             output_suffix=effective_suffix, segments=seg_txts)
+        else:
+            log(f"    🎙 转录中（模型: {model}, 语言: {args.language or 'auto'}）...")
+            ok, whisper_out = run_whisper(audio, str(out_dir), model, args.language)
+            if not ok:
+                raise RuntimeError(f"转录失败: {whisper_out}")
+
+            if whisper_out != final_txt:
+                os.replace(whisper_out, final_txt)
+
+            size = os.path.getsize(final_txt)
+            log(f"    输出: {Path(final_txt).name} ({size // 1024} KB)", "OK")
+            mark_transcribed(manifest, manifest_path, audio, final_txt,
+                             model, language=args.language,
+                             output_suffix=effective_suffix)
+
+        return "ok"
+
+    except Exception as e:
+        err_msg = str(e)
+        mark_transcribe_failed(manifest, manifest_path, audio, model, err_msg,
+                               output_suffix=effective_suffix)
+        # 清理不完整的输出
+        if os.path.exists(final_txt) and os.path.getsize(final_txt) == 0:
+            os.remove(final_txt)
+        return f"error: {err_msg}"
+
+
+# ─── 主流程 ────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="MP3 音频批量转录工具（基于 openai-whisper）",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__
+    )
+    parser.add_argument("source", nargs="?", default=None,
+                        help="MP3 文件路径或包含音频的目录（--preload-models 时可省略）")
+    parser.add_argument("--output", "-o", default=None,
+                        help="TXT 输出目录（默认与音频文件同目录）")
+    parser.add_argument("--manifest", "-m", default=None,
+                        help="manifest.json 路径（默认在 source 目录下）")
+    parser.add_argument("--model", default=DEFAULT_WHISPER_MODEL,
+                        choices=SUPPORTED_MODELS,
+                        help=f"Whisper 模型（默认 {DEFAULT_WHISPER_MODEL}，可选：{'/'.join(SUPPORTED_MODELS)}）")
+    parser.add_argument("--language", default=None,
+                        help="音频语言代码，推荐 zh(中文) / en(英文)（默认自动检测）")
+    parser.add_argument("--all-models", action="store_true",
+                        help=f"对每个音频依次使用所有支持的模型转录 ({'/'.join(SUPPORTED_MODELS)})，"
+                             f"输出为 <stem>.<model>.txt 便于对比精度")
+    parser.add_argument("--output-suffix", default="",
+                        help="输出文件名后缀，如 _v1 → <stem>_v1.txt（与 --all-models 互斥时被忽略）")
+    parser.add_argument("--split-size", type=int, default=None,
+                        help="切分阈值（秒）：超过此时长的音频会被切分后分别转录")
+    parser.add_argument("--split-segment", type=int, default=DEFAULT_SPLIT_SEGMENT_SECS,
+                        help=f"每个切分片段时长（秒，默认 {DEFAULT_SPLIT_SEGMENT_SECS}）")
+    parser.add_argument("--retry-failed", action="store_true",
+                        help="重新处理上次失败的文件")
+    parser.add_argument("--status", action="store_true",
+                        help="仅查看进度状态，不执行转录")
+    parser.add_argument("--preload-models", action="store_true",
+                        help=f"预下载所有支持的模型（{'/'.join(SUPPORTED_MODELS)}）到本地缓存后退出")
+    args = parser.parse_args()
+
+    # ── 模式1: 仅预下载模型 ──
+    if args.preload_models:
+        if not preload_models(SUPPORTED_MODELS):
+            sys.exit(1)
+        return
+
+    # 从这里开始 source 是必需的
+    if not args.source:
+        parser.error("需要指定 source 参数（音频文件或目录），除非使用 --preload-models")
+
+    if not check_whisper():
+        sys.exit(1)
+
+    has_ffmpeg = check_ffmpeg()
+    if args.split_size and not has_ffmpeg:
+        log("--split-size 需要 ffmpeg，但未找到 ffmpeg 命令", "ERR")
+        sys.exit(1)
+
+    source_path = Path(args.source).resolve()
+    base_dir = source_path if source_path.is_dir() else source_path.parent
+    manifest_path = Path(args.manifest) if args.manifest else base_dir / MANIFEST_NAME
+
+    manifest = load_manifest(manifest_path)
+
+    all_audios = scan_audios(args.source)
+    if not all_audios:
+        log("未找到任何音频文件", "WARN")
+        sys.exit(0)
+
+    # 决定使用哪些模型
+    if args.all_models:
+        models_to_run = SUPPORTED_MODELS
+        if args.output_suffix:
+            log(f"--all-models 已启用，--output-suffix 将被忽略（自动使用 .<model>.txt 命名）", "WARN")
+    else:
+        models_to_run = [args.model]
+
+    tag_model = args.all_models
+    log(f"扫描到 {len(all_audios)} 个音频文件，将使用 {len(models_to_run)} 个模型：{', '.join(models_to_run)}")
+
+    # ── 模式2: 仅查看状态 ──
+    if args.status:
+        show_status(manifest, all_audios)
+        return
+
+    # ── 模式3: 执行转录 ──
+    total_jobs = len(all_audios) * len(models_to_run)
+    stats = {"ok": 0, "skip": 0, "skip-failed": 0, "error": 0}
+    failed_details = []
+
+    log(f"共 {total_jobs} 个 (音频 × 模型) 任务待处理")
+
+    job_idx = 0
+    for audio in all_audios:
+        audio_name = Path(audio).name
+        audio_dir = Path(audio).parent
+        out_dir = Path(args.output).resolve() if args.output else audio_dir
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        log(f"\n📄 {audio_name}", "STEP")
+
+        for model in models_to_run:
+            job_idx += 1
+            log(f"  [{job_idx}/{total_jobs}] 模型: {model}")
+
+            t0 = time.time()
+            result = transcribe_one(audio, model, out_dir, args, manifest,
+                                    manifest_path, has_ffmpeg, tag_model)
+            elapsed = time.time() - t0
+
+            if result == "ok":
+                stats["ok"] += 1
+                log(f"    ⏱  耗时 {elapsed:.1f}s", "OK")
+            elif result == "skip":
+                stats["skip"] += 1
+                log(f"    ⏭  已完成，跳过", "INFO")
+            elif result == "skip-failed":
+                stats["skip-failed"] += 1
+                log(f"    ⚠  上次失败，加 --retry-failed 可重试", "WARN")
+            else:
+                stats["error"] += 1
+                failed_details.append((audio_name, model, result))
+                log(f"    ❌ {result}", "ERR")
+
+    # ── 汇总 ──
+    print(f"\n{'='*70}")
+    log(f"🏁 完成: 成功 {stats['ok']}  已跳过 {stats['skip']}  "
+        f"失败跳过 {stats['skip-failed']}  失败 {stats['error']}", "INFO")
+
+    if failed_details:
+        log("失败详情:", "ERR")
+        for audio_name, model, err in failed_details:
+            print(f"    - {audio_name} [模型:{model}] : {err[:80]}")
+        print("  运行时加 --retry-failed 可重新处理失败项")
+
+
+if __name__ == "__main__":
+    main()