faster whisper从多媒体语音材料中抽取出文本-4

"""
批量转录当前目录下的 .mp3 文件，使用 faster-whisper
转录完成后立即删除原始 .mp3 文件，仅保留 .txt
"""
import os
import sys
import logging
import argparse
from pathlib import Path
from faster_whisper import WhisperModel

# ================== 日志配置 ==================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('Extract Text')

# ================== 配置区 ==================
MODEL_SIZE = "base"        # 可选: tiny, base, small, medium, large
DEVICE = "cpu"              # cpu 或 cuda
COMPUTE_TYPE = "int8"       # int8, float16, float32 (CPU 推荐 int8)
VAD_FILTER = True           # 启用语音活动检测，去除静音
OUTPUT_FORMAT = "txt"       # 输出格式
VERBOSE = True              # 是否显示详细日志

# ===========================================

def transcribe_audio(audio_path: Path, model: WhisperModel) -> str:
    """转录单个音频文件，返回文本内容"""
    print(f"转录: {audio_path.name} → {audio_path.stem}.txt")

    segments, info = model.transcribe(
        str(audio_path),
        language=None,
        beam_size=5,
        vad_filter=VAD_FILTER,
        vad_parameters=dict(min_silence_duration_ms=500),
        word_timestamps=False,
    )

    text_lines = []
    for segment in segments:
        line = segment.text.strip()
        if line:
            text_lines.append(line)
        if VERBOSE:
            logger.info(f"[{segment.start:06.2f}s --> {segment.end:06.2f}s] {line}")

    return "\n".join(text_lines)


def safe_delete(file_path: Path):
    """安全删除文件，捕获异常并记录"""
    try:
        file_path.unlink()
        logger.info(f"已删除: {file_path.name}")
    except Exception as e:
        logger.error(f"删除失败 {file_path.name}: {e}")


def main(keep_audio: bool = False):
    print("=== faster-whisper 批量转录 + 自动清理音频 ===")

    current_dir = Path(".")
    mp3_files = sorted(current_dir.glob("*.mp3"))

    if not mp3_files:
        print("未找到 .mp3 文件，退出。")
        return

    # 加载模型（只加载一次）
    print(f"正在加载模型 {MODEL_SIZE} ({DEVICE}, {COMPUTE_TYPE})...")
    model = WhisperModel(MODEL_SIZE, device=DEVICE, compute_type=COMPUTE_TYPE)

    processed = 0
    for mp3_path in mp3_files:
        txt_path = mp3_path.with_suffix(".txt")

        # 如果 .txt 已存在，跳过（但仍可选择删除 .mp3）
        if txt_path.exists():
            print(f"跳过: {txt_path.name} 已存在")
            if not keep_audio and mp3_path.exists():
                safe_delete(mp3_path)
            continue

        try:
            # 转录
            text = transcribe_audio(mp3_path, model)
            if not text.strip():
                print(f"警告: {mp3_path.name} 转录结果为空，跳过删除")
                continue

            # 写入文本
            txt_path.write_text(text, encoding="utf-8")
            processed += 1
            print(f"已保存: {txt_path.name}")

            # 删除原始音频（仅当不保留时）
            if not keep_audio:
                safe_delete(mp3_path)

        except Exception as e:
            print(f"错误转录 {mp3_path.name}: {e}", file=sys.stderr)
            logger.error(f"处理失败: {mp3_path.name} - {e}")

    mode = "保留音频" if keep_audio else "已删除音频"
    print(f"全部完成！共处理 {processed} 个文件，{mode}。")


if __name__ == "__main__":
    main()