OpenClaw Agent
8bce4ac23b
- 新增 subtitle_generator.py 核心模块 - AudioDownloader: 使用 yt-dlp 下载视频音频 - WhisperTranscriber: Whisper 语音识别器 - SubtitleFormatter: SRT 格式化器 - SubtitleCache: 缓存管理器 - SubtitleGenerator: 统一接口 - 集成到 pipeline.py - 添加 _ensure_subtitle() 方法 - 自动检测无字幕视频并生成字幕 - 支持缓存机制避免重复处理 - 更新配置文件 - config.yaml.example 添加字幕配置项 - requirements.txt 添加依赖包 - 文档更新 - SKILL.md 添加字幕生成说明 - SUBTITLE_FEATURE.md 详细功能文档 - 测试脚本 - test_subtitle_generator.py 完整测试套件 特性： - 本地处理，数据不上传，隐私安全 - 支持中英文识别及自动检测 - 多种模型可选（tiny/base/small/medium/large） - 智能缓存，提升处理速度 - 完善的错误处理和日志输出
494 lines
14 KiB
Python
494 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
字幕生成器 - 为无字幕视频自动生成 SRT 格式字幕
|
||
|
||
使用 OpenAI Whisper 进行语音识别，支持本地离线处理
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import hashlib
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Optional, Callable, Dict, Any
|
||
|
||
# 配置日志
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class AudioDownloader:
|
||
"""使用 yt-dlp 下载视频音频"""
|
||
|
||
def __init__(self):
|
||
"""初始化音频下载器"""
|
||
self.ydl_opts = {
|
||
'format': 'bestaudio/best',
|
||
'postprocessors': [{
|
||
'key': 'FFmpegExtractAudio',
|
||
'preferredcodec': 'mp3',
|
||
'preferredquality': '192',
|
||
}],
|
||
'quiet': True,
|
||
'no_warnings': True,
|
||
}
|
||
|
||
def download_audio(self, url: str, output_path: str) -> str:
|
||
"""
|
||
下载视频音频
|
||
|
||
Args:
|
||
url: 视频URL
|
||
output_path: 输出音频路径（可以是模板，如 'cache/audio/%(id)s.mp3'）
|
||
|
||
Returns:
|
||
str: 实际保存的音频文件路径
|
||
|
||
Raises:
|
||
Exception: 下载失败时抛出异常
|
||
"""
|
||
try:
|
||
import yt_dlp
|
||
|
||
ydl_opts = self.ydl_opts.copy()
|
||
ydl_opts['outtmpl'] = output_path
|
||
|
||
logger.info(f"开始下载音频: {url}")
|
||
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
||
ydl.download([url])
|
||
|
||
# 替换扩展名
|
||
actual_path = output_path.replace('.%(ext)s', '.mp3')
|
||
if not os.path.exists(actual_path):
|
||
# 尝试其他可能的文件名
|
||
actual_path = output_path.rsplit('.', 1)[0] + '.mp3'
|
||
|
||
if os.path.exists(actual_path):
|
||
logger.info(f"音频下载成功: {actual_path}")
|
||
return actual_path
|
||
else:
|
||
raise FileNotFoundError(f"下载的音频文件未找到: {actual_path}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"音频下载失败: {e}")
|
||
raise
|
||
|
||
|
||
class WhisperTranscriber:
|
||
"""Whisper 语音识别器"""
|
||
|
||
# 模型选择
|
||
MODELS = {
|
||
'tiny': 39000000, # 最快，准确度最低 (~39M)
|
||
'base': 74000000, # 推荐，速度/准确度平衡 (~74M)
|
||
'small': 244000000, # 准确度更高 (~244M)
|
||
'medium': 769000000, # 接近专业级 (~769M)
|
||
'large': 1550000000, # 最高准确度 (~1.55GB)
|
||
}
|
||
|
||
def __init__(self, model_size: str = 'base'):
|
||
"""
|
||
初始化识别器
|
||
|
||
Args:
|
||
model_size: 模型大小，默认 'base'
|
||
"""
|
||
self.model_size = model_size
|
||
self.model = None
|
||
self.device = None
|
||
|
||
def load_model(self):
|
||
"""加载 Whisper 模型（懒加载）"""
|
||
if self.model is None:
|
||
logger.info(f"正在加载 Whisper 模型 ({self.model_size})...")
|
||
start_time = time.time()
|
||
|
||
try:
|
||
import whisper
|
||
import torch
|
||
|
||
# 自动检测 GPU
|
||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
logger.info(f"使用设备: {self.device}")
|
||
|
||
self.model = whisper.load_model(
|
||
self.model_size,
|
||
device=self.device
|
||
)
|
||
|
||
load_time = time.time() - start_time
|
||
logger.info(f"模型加载完成 - 耗时: {load_time:.1f}秒")
|
||
|
||
except Exception as e:
|
||
logger.error(f"模型加载失败: {e}")
|
||
raise
|
||
|
||
def transcribe(
|
||
self,
|
||
audio_path: str,
|
||
language: Optional[str] = None,
|
||
progress_callback: Optional[Callable[[str, float], None]] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
语音识别
|
||
|
||
Args:
|
||
audio_path: 音频文件路径
|
||
language: 语言代码（'zh'、'en'、None=自动检测）
|
||
progress_callback: 进度回调函数 callback(message, progress_percent)
|
||
|
||
Returns:
|
||
dict: Whisper 识别结果，包含 'text' 和 'segments'
|
||
|
||
Raises:
|
||
Exception: 识别失败时抛出异常
|
||
"""
|
||
self.load_model()
|
||
|
||
logger.info(f"开始语音识别: {audio_path}")
|
||
if language:
|
||
logger.info(f"指定语言: {language}")
|
||
else:
|
||
logger.info("语言: 自动检测")
|
||
|
||
start_time = time.time()
|
||
|
||
try:
|
||
result = self.model.transcribe(
|
||
audio_path,
|
||
language=language,
|
||
task='transcribe',
|
||
verbose=False,
|
||
fp16=False # CPU 上不使用 FP16
|
||
)
|
||
|
||
duration = time.time() - start_time
|
||
logger.info(f"识别完成 - 耗时: {duration:.1f}秒")
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"语音识别失败: {e}")
|
||
raise
|
||
|
||
def get_segments(self, result: Dict[str, Any]) -> list:
|
||
"""
|
||
获取分段信息
|
||
|
||
Args:
|
||
result: Whisper 识别结果
|
||
|
||
Returns:
|
||
list: segments 列表
|
||
"""
|
||
return result.get('segments', [])
|
||
|
||
|
||
class SubtitleFormatter:
|
||
"""字幕格式化器（SRT 格式）"""
|
||
|
||
@staticmethod
|
||
def to_srt(segments: list) -> str:
|
||
"""
|
||
将 Whisper segments 转换为 SRT 格式
|
||
|
||
Args:
|
||
segments: Whisper segments 列表
|
||
|
||
Returns:
|
||
str: SRT 格式文本
|
||
"""
|
||
srt_lines = []
|
||
|
||
for i, segment in enumerate(segments, 1):
|
||
# SRT 格式：
|
||
# 序号
|
||
# 开始时间 --> 结束时间
|
||
# 文本
|
||
# 空行
|
||
|
||
start_time = SubtitleFormatter._seconds_to_srt_time(segment['start'])
|
||
end_time = SubtitleFormatter._seconds_to_srt_time(segment['end'])
|
||
text = segment['text'].strip()
|
||
|
||
srt_lines.append(str(i))
|
||
srt_lines.append(f"{start_time} --> {end_time}")
|
||
srt_lines.append(text)
|
||
srt_lines.append("")
|
||
|
||
return "\n".join(srt_lines)
|
||
|
||
@staticmethod
|
||
def _seconds_to_srt_time(seconds: float) -> str:
|
||
"""
|
||
将秒数转换为 SRT 时间格式 (HH:MM:SS,mmm)
|
||
|
||
Args:
|
||
seconds: 秒数
|
||
|
||
Returns:
|
||
str: SRT 时间格式
|
||
"""
|
||
hours = int(seconds // 3600)
|
||
minutes = int((seconds % 3600) // 60)
|
||
secs = int(seconds % 60)
|
||
millis = int((seconds - int(seconds)) * 1000)
|
||
|
||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
|
||
|
||
|
||
class SubtitleCache:
|
||
"""字幕缓存管理器"""
|
||
|
||
def __init__(self, cache_dir: str = 'cache/subtitles'):
|
||
"""
|
||
初始化缓存
|
||
|
||
Args:
|
||
cache_dir: 缓存目录
|
||
"""
|
||
self.cache_dir = Path(cache_dir)
|
||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
def get_cache_key(self, url: str) -> str:
|
||
"""
|
||
生成缓存键
|
||
|
||
Args:
|
||
url: 视频URL
|
||
|
||
Returns:
|
||
str: 缓存键（SHA256 哈希的前8位）
|
||
"""
|
||
return hashlib.sha256(url.encode()).hexdigest()[:8]
|
||
|
||
def get_cache_path(self, url: str) -> str:
|
||
"""
|
||
获取缓存文件路径
|
||
|
||
Args:
|
||
url: 视频URL
|
||
|
||
Returns:
|
||
str: 缓存文件路径
|
||
"""
|
||
cache_key = self.get_cache_key(url)
|
||
return str(self.cache_dir / f"{cache_key}.srt")
|
||
|
||
def get(self, url: str) -> Optional[str]:
|
||
"""
|
||
获取缓存字幕
|
||
|
||
Args:
|
||
url: 视频URL
|
||
|
||
Returns:
|
||
str | None: 缓存的字幕内容，如果不存在返回 None
|
||
"""
|
||
cache_path = self.get_cache_path(url)
|
||
|
||
if os.path.exists(cache_path):
|
||
logger.info(f"使用缓存字幕: {cache_path}")
|
||
try:
|
||
with open(cache_path, 'r', encoding='utf-8') as f:
|
||
return f.read()
|
||
except Exception as e:
|
||
logger.warning(f"读取缓存失败: {e}")
|
||
return None
|
||
|
||
return None
|
||
|
||
def set(self, url: str, subtitle: str) -> str:
|
||
"""
|
||
保存缓存字幕
|
||
|
||
Args:
|
||
url: 视频URL
|
||
subtitle: 字幕内容
|
||
|
||
Returns:
|
||
str: 缓存文件路径
|
||
"""
|
||
cache_path = self.get_cache_path(url)
|
||
|
||
try:
|
||
with open(cache_path, 'w', encoding='utf-8') as f:
|
||
f.write(subtitle)
|
||
logger.info(f"字幕已缓存: {cache_path}")
|
||
return cache_path
|
||
except Exception as e:
|
||
logger.error(f"保存缓存失败: {e}")
|
||
raise
|
||
|
||
def exists(self, url: str) -> bool:
|
||
"""检查缓存是否存在"""
|
||
return os.path.exists(self.get_cache_path(url))
|
||
|
||
def clear_cache(self):
|
||
"""清空所有缓存"""
|
||
try:
|
||
for file in self.cache_dir.glob('*.srt'):
|
||
file.unlink()
|
||
logger.info("缓存已清空")
|
||
except Exception as e:
|
||
logger.error(f"清空缓存失败: {e}")
|
||
|
||
|
||
class SubtitleGenerator:
|
||
"""字幕生成器统一接口"""
|
||
|
||
def __init__(
|
||
self,
|
||
model_size: str = 'base',
|
||
cache_dir: str = 'cache/subtitles',
|
||
audio_cache_dir: str = 'cache/audio'
|
||
):
|
||
"""
|
||
初始化生成器
|
||
|
||
Args:
|
||
model_size: Whisper 模型大小
|
||
cache_dir: 字幕缓存目录
|
||
audio_cache_dir: 音频缓存目录
|
||
"""
|
||
self.audio_downloader = AudioDownloader()
|
||
self.transcriber = WhisperTranscriber(model_size)
|
||
self.formatter = SubtitleFormatter()
|
||
self.cache = SubtitleCache(cache_dir)
|
||
self.audio_cache_dir = Path(audio_cache_dir)
|
||
self.audio_cache_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
def generate(
|
||
self,
|
||
url: str,
|
||
content_id: str = None,
|
||
language: Optional[str] = None,
|
||
use_cache: bool = True,
|
||
keep_audio: bool = False
|
||
) -> str:
|
||
"""
|
||
生成字幕
|
||
|
||
Args:
|
||
url: 视频URL
|
||
content_id: 内容ID（用于音频文件命名）
|
||
language: 语言代码（None=自动检测）
|
||
use_cache: 是否使用缓存
|
||
keep_audio: 是否保留音频文件
|
||
|
||
Returns:
|
||
str: SRT 格式字幕
|
||
|
||
Raises:
|
||
Exception: 生成失败时抛出异常
|
||
"""
|
||
# 检查缓存
|
||
if use_cache and self.cache.exists(url):
|
||
logger.info(f"使用缓存字幕: {url}")
|
||
cached_subtitle = self.cache.get(url)
|
||
if cached_subtitle:
|
||
return cached_subtitle
|
||
|
||
logger.info(f"开始生成字幕: {url}")
|
||
audio_path = None
|
||
|
||
try:
|
||
# 生成 content_id
|
||
if content_id is None:
|
||
content_id = self.cache.get_cache_key(url)
|
||
|
||
# 下载音频
|
||
audio_path = str(self.audio_cache_dir / f"{content_id}.mp3")
|
||
logger.info(f"下载音频到: {audio_path}")
|
||
audio_path = self.audio_downloader.download_audio(url, audio_path)
|
||
|
||
# 语音识别
|
||
result = self.transcriber.transcribe(audio_path, language=language)
|
||
segments = self.transcriber.get_segments(result)
|
||
|
||
if not segments:
|
||
logger.warning("未识别到任何语音内容")
|
||
return ""
|
||
|
||
# 转换为 SRT
|
||
srt_text = self.formatter.to_srt(segments)
|
||
logger.info(f"SRT 字幕生成成功，共 {len(segments)} 段")
|
||
|
||
# 保存缓存
|
||
self.cache.set(url, srt_text)
|
||
|
||
# 清理音频文件
|
||
if not keep_audio and os.path.exists(audio_path):
|
||
os.remove(audio_path)
|
||
logger.info(f"已删除音频文件: {audio_path}")
|
||
|
||
return srt_text
|
||
|
||
except Exception as e:
|
||
logger.error(f"字幕生成失败: {e}")
|
||
# 清理失败的音频文件
|
||
if audio_path and os.path.exists(audio_path) and not keep_audio:
|
||
try:
|
||
os.remove(audio_path)
|
||
except:
|
||
pass
|
||
raise
|
||
|
||
|
||
def generate_subtitle(
|
||
url: str,
|
||
model_size: str = 'base',
|
||
language: Optional[str] = None,
|
||
cache_dir: str = 'cache/subtitles',
|
||
audio_cache_dir: str = 'cache/audio'
|
||
) -> str:
|
||
"""
|
||
便捷函数：生成字幕
|
||
|
||
Args:
|
||
url: 视频URL
|
||
model_size: Whisper 模型大小
|
||
language: 语言代码（None=自动检测）
|
||
cache_dir: 字幕缓存目录
|
||
audio_cache_dir: 音频缓存目录
|
||
|
||
Returns:
|
||
str: SRT 格式字幕
|
||
"""
|
||
generator = SubtitleGenerator(
|
||
model_size=model_size,
|
||
cache_dir=cache_dir,
|
||
audio_cache_dir=audio_cache_dir
|
||
)
|
||
return generator.generate(url, language=language)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 测试示例
|
||
import sys
|
||
|
||
if len(sys.argv) < 2:
|
||
print("用法: python subtitle_generator.py <视频URL> [语言]")
|
||
print("示例: python subtitle_generator.py https://www.bilibili.com/video/BV1xx411c7mD zh")
|
||
sys.exit(1)
|
||
|
||
url = sys.argv[1]
|
||
language = sys.argv[2] if len(sys.argv) > 2 else None
|
||
|
||
print("=" * 60)
|
||
print("字幕生成器")
|
||
print("=" * 60)
|
||
|
||
try:
|
||
subtitle = generate_subtitle(url, language=language)
|
||
print("\n" + "=" * 60)
|
||
print("生成成功！字幕预览（前500字符）：")
|
||
print("=" * 60)
|
||
print(subtitle[:500])
|
||
if len(subtitle) > 500:
|
||
print("\n...")
|
||
print(f"\n总字符数: {len(subtitle)}")
|
||
except Exception as e:
|
||
print(f"\n错误: {e}")
|
||
sys.exit(1)
|