Python使用Whisper + Transformers自动生成中英文双语字幕的完整流程

pip install openai-whisper transformers pydub librosa tqdm torch ffmpeg-python modelscope

from modelscope import snapshot_download

# 下载模型到当前目录
model_dir = snapshot_download('Helsinki-NLP/opus-mt-en-zh', cache_dir='./')

print(f"✅ 模型已下载到当前目录: {model_dir}")

model_dir = snapshot_download('Helsinki-NLP/opus-mt-zh-en', cache_dir='./')

print(f"✅ 模型已下载到当前目录: {model_dir}")

import whisper
import warnings
from datetime import timedelta
from tqdm import tqdm
import librosa
import time
from transformers import pipeline
import torch
from pydub import AudioSegment
import os
import ffmpeg

# 忽略警告信息
warnings.filterwarnings("ignore", category=UserWarning)

# 输入视频路径（可修改）
input_video = r"./Gesture Drawing Practice _ 20 and 40 sec. poses.mp4"
# 输出音频路径
output_audio = "Gesture Drawing Practice _ 20 and 40 sec. poses.wav"

# =========================
# Step 1: 提取视频中的音频
# =========================
ffmpeg.input(input_video).output(output_audio, ac=1, ar=16000).run()

# =========================
# Step 2: 时间戳格式化函数
# =========================
def format_timestamp(seconds):
    """将秒数转换为 SRT 时间格式（HH:MM:SS,mmm）"""
    td = timedelta(seconds=seconds)
    hours, remainder = divmod(td.seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    milliseconds = int(td.microseconds / 1000)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"

# =========================
# Step 3: 验证音频文件是否有效
# =========================
def validate_audio(audio_file):
    """检查音频文件是否有效，并返回其基本信息"""
    try:
        audio = AudioSegment.from_file(audio_file)
        duration = len(audio) / 1000.0  # 秒
        sample_rate = audio.frame_rate
        channels = audio.channels
        print(f"✅ 音频验证成功：时长={duration:.2f}s, 采样率={sample_rate}Hz, 声道={channels}")
        return True, duration
    except Exception as e:
        print(f"❌ 音频验证失败: {e}")
        return False, 0

# =========================
# Step 4: 初始化翻译模型
# =========================
print("🔤 正在加载翻译模型，请稍候...")
translator_en_to_zh = pipeline("translation", model="./Helsinki-NLP/opus-mt-en-zh")  # 英译中
translator_zh_to_en = pipeline("translation", model="./Helsinki-NLP/opus-mt-zh-en")  # 中译英
print("✅ 翻译模型加载完成。")

# =========================
# Step 5: 验证音频文件
# =========================
audio_file = output_audio
is_valid, duration = validate_audio(audio_file)
if not is_valid:
    raise ValueError(f"无效的音频文件: {audio_file}")

# 备用方案：使用 librosa 检查时长
try:
    duration = librosa.get_duration(path=audio_file)
except Exception as e:
    print(f"Librosa 检查失败: {e}")

# =========================
# Step 6: 初始化 Whisper 模型
# =========================
print("🎙️ 正在加载 Whisper 模型...")
try:
    model = whisper.load_model("base")  # 可选 tiny, base, small, medium, large
except Exception as e:
    print(f"加载模型失败: {e}，回退到 'medium' 模型。")
    model = whisper.load_model("medium")

# 进度条初始化
progress_bar = tqdm(total=100, desc="语音识别中", unit="%")

# =========================
# Step 7: 执行语音识别（带进度）
# =========================
def transcribe_with_progress(model, audio_file, language=None):
    start_time = time.time()
    
    # 自动检测语言
    if language is None:
        try:
            audio = whisper.load_audio(audio_file)
            mel = whisper.log_mel_spectrogram(audio, n_mels=80).to(model.device)
            _, probs = model.detect_language(mel)
            detected_language = max(probs, key=probs.get)
            print(f"🌍 自动检测语言: {detected_language}")
        except Exception as e:
            print(f"语言检测失败: {e}，默认使用英语(en)。")
            detected_language = "en"
    else:
        detected_language = language
        print(f"🌐 使用指定语言: {detected_language}")
    
    # 执行转录
    try:
        result = model.transcribe(audio_file, language=detected_language)
    except Exception as e:
        print(f"转录失败: {e}")
        raise
    
    # 更新进度条
    progress_bar.update(100 - progress_bar.n)
    progress_bar.close()
    
    elapsed_time = time.time() - start_time
    print(f"✅ 转录完成，用时 {elapsed_time:.2f} 秒。")
    
    return result, detected_language

# 执行转录（可指定语言）
result, detected_language = transcribe_with_progress(model, audio_file, language="en")

# =========================
# Step 8: 生成中英双语字幕文件
# =========================
srt_path = f"{output_audio}.srt"
with open(srt_path, "w", encoding="utf-8") as f:
    for i, segment in enumerate(result["segments"]):
        start_time = format_timestamp(segment["start"])
        end_time = format_timestamp(segment["end"])
        text = segment["text"].strip()
        
        # 判断语言并翻译
        if detected_language == "zh":
            # 中文音频：原文中文 + 翻译英文
            zh_text = text
            en_text = translator_zh_to_en(text)[0]["translation_text"]
        else:
            # 英语音频：原文英文 + 翻译中文
            zh_text = translator_en_to_zh(text)[0]["translation_text"]
            en_text = text
        
        # 写入 SRT 文件（中文在上，英文在下）
        f.write(f"{i+1}\n")
        f.write(f"{start_time} --> {end_time}\n")
        f.write(f"{zh_text}\n{en_text}\n\n")

print(f"🎬 字幕文件生成成功：{srt_path}")