Python 语音转文本(wav2vec2)

作者：追风剑情发布于：2026-6-9 17:17 分类：AI

一、下载 FFmpeg

将 ffmpeg 的 bin 目录配置到环境变量中。

示例：

import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import argparse
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
from pydub import AudioSegment
import soundfile as sf
from tqdm import tqdm

def convert_and_split_audio(input_file, output_dir, chunk_length_ms=60*1000, target_sample_rate=16000):
    """
    将音频文件转换为目标采样率，并按指定时长分割成多个小块。
    返回分割后的临时文件路径列表。
    """
    # 加载音频（支持 mp3, wav 等）
    audio = AudioSegment.from_file(input_file)
    # 转换采样率
    audio = audio.set_frame_rate(target_sample_rate)
    # 按块切割
    chunks = [audio[i:i+chunk_length_ms] for i in range(0, len(audio), chunk_length_ms)]
    
    output_files = []
    base_name = os.path.splitext(os.path.basename(input_file))[0]
    os.makedirs(output_dir, exist_ok=True)
    
    for idx, chunk in enumerate(chunks):
        chunk_name = os.path.join(output_dir, f"{base_name}_chunk{idx}.wav")
        chunk.export(chunk_name, format="wav")
        output_files.append(chunk_name)
    
    return output_files

def main():
    parser = argparse.ArgumentParser(description="使用 Wav2Vec2 进行长音频转录")
    parser.add_argument("--input_file", type=str, default="888.mp3", help="输入音频文件路径")
    parser.add_argument("--model_name", type=str, default="jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn", help="预训练模型名称")
    parser.add_argument("--chunk_minutes", type=int, default=1, help="每块时长（分钟）")
    parser.add_argument("--temp_dir", type=str, default="./temp_chunks", help="临时分块存储目录")
    args = parser.parse_args()

    # 1. 分割音频
    print("正在分割音频...")
    chunk_files = convert_and_split_audio(
        args.input_file,
        args.temp_dir,
        chunk_length_ms=args.chunk_minutes * 60 * 1000,
        target_sample_rate=16000
    )
    print(f"共分割为 {len(chunk_files)} 个片段")

    # 2. 加载模型和处理器（使用 GPU 如果可用）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    processor = Wav2Vec2Processor.from_pretrained(args.model_name)
    model = Wav2Vec2ForCTC.from_pretrained(args.model_name).to(device)
    model.eval()

    # 3. 逐块转录
    transcriptions = []
    for chunk_file in tqdm(chunk_files, desc="转录进度"):
        try:
            # 读取音频块
            audio_input, sample_rate = sf.read(chunk_file)
            # 确保是单声道（如果立体声则取均值）
            if len(audio_input.shape) > 1:
                audio_input = audio_input.mean(axis=1)
            # 处理输入
            inputs = processor(audio_input, sampling_rate=16000, return_tensors="pt", padding=True)
            input_values = inputs.input_values.to(device)
            # 推理
            with torch.no_grad():
                logits = model(input_values).logits
            # 解码
            predicted_ids = torch.argmax(logits, dim=-1)
            transcription = processor.batch_decode(predicted_ids)[0]
            transcriptions.append(transcription)
        except Exception as e:
            print(f"处理文件 {chunk_file} 时出错: {e}")
            transcriptions.append("")

    # 4. 保存结果
    output_txt = "transcriptions.txt"
    with open(output_txt, "w", encoding="utf-8") as f:
        for idx, text in enumerate(transcriptions):
            f.write(f"块 {idx}:\n{text}\n\n")
    print(f"转录结果已保存到 {output_txt}")

    # 可选：删除临时分块文件
    # import shutil
    # shutil.rmtree(args.temp_dir)

if __name__ == "__main__":
    main()

运行测试