feat(03-01): implement ASRProcessor with official subtitle download
- Create processors module with ASRProcessor class - Implement _try_official_subtitle for B站字幕下载 - Support process() method with subtitle-first strategy - Add helper functions for whisper.cpp integration - All 5 processor tests passing Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,9 @@
|
||||
"""Processors module for KB Pipeline
|
||||
|
||||
This module contains processors for various content transformation tasks:
|
||||
- ASRProcessor: Audio-to-text conversion (subtitle download + Whisper.cpp)
|
||||
"""
|
||||
|
||||
from .asr import ASRProcessor
|
||||
|
||||
__all__ = ["ASRProcessor"]
|
||||
@@ -0,0 +1,575 @@
|
||||
"""ASRProcessor - Audio-to-text conversion
|
||||
|
||||
Unified processor for B站官方字幕下载 and Whisper.cpp 语音识别.
|
||||
Follows the pipeline pattern established in Phase 1/2.
|
||||
|
||||
Strategy:
|
||||
1. Try official subtitle first (highest quality, no compute cost)
|
||||
2. Fall back to Whisper.cpp if no subtitle available
|
||||
3. Auto-enable speaker diarization if HF_TOKEN is configured
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..config import get_settings
|
||||
from ..schemas import ASRInput, ASROutput
|
||||
from ..state import WorkflowStateManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WhisperNotFoundError(RuntimeError):
|
||||
"""Whisper.cpp 路径未找到"""
|
||||
pass
|
||||
|
||||
|
||||
def get_whisper_cpp_path() -> str:
|
||||
"""获取 whisper.cpp 安装路径
|
||||
|
||||
检查环境变量 WHISPER_CPP_PATH 或默认路径 ~/whisper.cpp
|
||||
|
||||
Returns:
|
||||
whisper.cpp 目录路径
|
||||
|
||||
Raises:
|
||||
WhisperNotFoundError: 如果路径不存在
|
||||
"""
|
||||
# 检查环境变量
|
||||
path = os.environ.get('WHISPER_CPP_PATH')
|
||||
if path and os.path.exists(path):
|
||||
return os.path.expanduser(path)
|
||||
|
||||
# 检查默认路径
|
||||
default = os.path.expanduser('~/whisper.cpp')
|
||||
if os.path.exists(default):
|
||||
return default
|
||||
|
||||
raise WhisperNotFoundError(
|
||||
"未找到 whisper.cpp。请设置 WHISPER_CPP_PATH 环境变量或安装到 ~/whisper.cpp"
|
||||
)
|
||||
|
||||
|
||||
def convert_to_wav(input_path: str, output_path: Optional[str] = None) -> str:
|
||||
"""使用 ffmpeg 将音频/视频转换为 16kHz WAV 格式
|
||||
|
||||
Args:
|
||||
input_path: 输入文件路径
|
||||
output_path: 输出文件路径（可选，默认为 input_path_audio.wav）
|
||||
|
||||
Returns:
|
||||
输出 WAV 文件路径
|
||||
"""
|
||||
if output_path is None:
|
||||
base = os.path.splitext(input_path)[0]
|
||||
output_path = f"{base}_audio.wav"
|
||||
|
||||
cmd = [
|
||||
'ffmpeg', '-y', '-i', input_path,
|
||||
'-ar', '16000', # 16kHz 采样率
|
||||
'-ac', '1', # 单声道
|
||||
'-c:a', 'pcm_s16le', # 16-bit PCM
|
||||
output_path
|
||||
]
|
||||
|
||||
subprocess.run(cmd, capture_output=True, check=True)
|
||||
return output_path
|
||||
|
||||
|
||||
class ASRProcessor:
|
||||
"""ASR 处理器 - 统一字幕下载和语音识别
|
||||
|
||||
处理流程：
|
||||
1. 如果是 B站视频，尝试下载官方字幕
|
||||
2. 无官方字幕则调用 Whisper.cpp
|
||||
3. HF_TOKEN 配置时自动启用说话者分离
|
||||
4. 记录处理状态到 workflow_state.json
|
||||
|
||||
Attributes:
|
||||
STEP_NAME: 步骤名称，用于状态记录
|
||||
settings: KB 配置
|
||||
state_manager: 状态管理器（可选）
|
||||
"""
|
||||
|
||||
STEP_NAME = "asr"
|
||||
|
||||
def __init__(self, state_manager: Optional[WorkflowStateManager] = None):
|
||||
"""初始化 ASR 处理器
|
||||
|
||||
Args:
|
||||
state_manager: 状态管理器（可选，用于记录处理状态）
|
||||
"""
|
||||
self.settings = get_settings()
|
||||
self.state_manager = state_manager
|
||||
|
||||
async def process(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行 ASR 处理
|
||||
|
||||
策略：
|
||||
1. 先尝试官方字幕（仅 B站）
|
||||
2. 无官方字幕则调用 Whisper
|
||||
3. 记录状态
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
|
||||
Raises:
|
||||
Exception: 处理失败时抛出异常
|
||||
"""
|
||||
result = None
|
||||
|
||||
try:
|
||||
# Step 1: 尝试官方字幕（仅 B站）
|
||||
if input_data.platform == "bilibili":
|
||||
result = await self._try_official_subtitle(input_data)
|
||||
|
||||
# Step 2: 无官方字幕则调用 Whisper
|
||||
if result is None:
|
||||
result = await self._run_whisper(input_data)
|
||||
|
||||
# Step 3: 记录成功状态
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
summary = f"tool={result.tool}, segments={len(result.segments)}"
|
||||
if result.speakers:
|
||||
summary += f", speakers={len(result.speakers)}"
|
||||
state.mark_step_completed(self.STEP_NAME, summary)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# 记录失败状态（只在所有重试都失败后）
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_failed(self.STEP_NAME, str(e))
|
||||
self.state_manager.save(state)
|
||||
raise
|
||||
|
||||
async def _try_official_subtitle(self, input_data: ASRInput) -> Optional[ASROutput]:
|
||||
"""尝试下载 B站官方字幕
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
如果找到字幕返回 ASROutput，否则返回 None
|
||||
|
||||
Note:
|
||||
无字幕返回 None（不是抛异常），让流程继续到 Whisper
|
||||
"""
|
||||
video_path = Path(input_data.video_path)
|
||||
output_dir = video_path.parent
|
||||
base_name = video_path.stem
|
||||
|
||||
# 输出文件路径
|
||||
txt_path = output_dir / f"{base_name}.txt"
|
||||
srt_path = output_dir / f"{base_name}.srt"
|
||||
json_path = output_dir / f"{base_name}.json"
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
# Step 1: 获取 CID
|
||||
bvid = input_data.content_id
|
||||
pagelist_url = f"https://api.bilibili.com/x/player/pagelist?bvid={bvid}"
|
||||
async with session.get(pagelist_url) as resp:
|
||||
pagelist_data = await resp.json()
|
||||
|
||||
pages = pagelist_data.get('data', [])
|
||||
if not pages:
|
||||
logger.info(f"No pagelist data for {bvid}")
|
||||
return None
|
||||
|
||||
cid = pages[0].get('cid')
|
||||
if not cid:
|
||||
logger.info(f"No cid found for {bvid}")
|
||||
return None
|
||||
|
||||
# Step 2: 获取字幕列表
|
||||
player_url = f"https://api.bilibili.com/x/player/v2?bvid={bvid}&cid={cid}"
|
||||
async with session.get(player_url) as resp:
|
||||
player_data = await resp.json()
|
||||
|
||||
subtitle_info = player_data.get('data', {}).get('subtitle', {})
|
||||
subtitles = subtitle_info.get('list', [])
|
||||
|
||||
if not subtitles:
|
||||
logger.info(f"No official subtitle for {bvid}")
|
||||
return None
|
||||
|
||||
# Step 3: 下载第一个字幕（优先 AI 生成的中文字幕）
|
||||
subtitle_url = None
|
||||
for sub in subtitles:
|
||||
if sub.get('lan', '').startswith('zh'):
|
||||
subtitle_url = sub.get('subtitle_url')
|
||||
break
|
||||
|
||||
if not subtitle_url and subtitles:
|
||||
subtitle_url = subtitles[0].get('subtitle_url')
|
||||
|
||||
if not subtitle_url:
|
||||
return None
|
||||
|
||||
# 补全 URL
|
||||
if subtitle_url.startswith('//'):
|
||||
subtitle_url = f"https:{subtitle_url}"
|
||||
|
||||
async with session.get(subtitle_url) as resp:
|
||||
subtitle_content = await resp.json()
|
||||
|
||||
# Step 4: 转换为输出格式
|
||||
segments = []
|
||||
text_parts = []
|
||||
body = subtitle_content.get('body', [])
|
||||
|
||||
for i, item in enumerate(body):
|
||||
start = item.get('from', 0)
|
||||
end = item.get('to', 0)
|
||||
content = item.get('content', '')
|
||||
|
||||
segments.append({
|
||||
'id': i,
|
||||
'start': start,
|
||||
'end': end,
|
||||
'text': content
|
||||
})
|
||||
text_parts.append(content)
|
||||
|
||||
full_text = ' '.join(text_parts)
|
||||
duration = segments[-1]['end'] if segments else 0
|
||||
|
||||
# Step 5: 写入文件
|
||||
# TXT 格式
|
||||
txt_path.write_text(full_text, encoding='utf-8')
|
||||
|
||||
# SRT 格式
|
||||
srt_content = self._convert_to_srt(segments)
|
||||
srt_path.write_text(srt_content, encoding='utf-8')
|
||||
|
||||
# JSON 格式
|
||||
json_data = {
|
||||
'success': True,
|
||||
'tool': 'subtitle',
|
||||
'segments': segments,
|
||||
'text': full_text,
|
||||
'duration_seconds': duration
|
||||
}
|
||||
json_path.write_text(json.dumps(json_data, ensure_ascii=False, indent=2), encoding='utf-8')
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="subtitle",
|
||||
has_subtitle=True,
|
||||
output_files={
|
||||
'txt': str(txt_path),
|
||||
'srt': str(srt_path),
|
||||
'json': str(json_path)
|
||||
},
|
||||
segments=segments,
|
||||
text=full_text,
|
||||
duration_seconds=duration
|
||||
)
|
||||
|
||||
except aiohttp.ClientError as e:
|
||||
logger.warning(f"Failed to fetch subtitle: {e}")
|
||||
return None
|
||||
|
||||
def _convert_to_srt(self, segments: List[Dict]) -> str:
|
||||
"""将字幕分段转换为 SRT 格式
|
||||
|
||||
Args:
|
||||
segments: 字幕分段列表，每个包含 start, end, text
|
||||
|
||||
Returns:
|
||||
SRT 格式字符串
|
||||
"""
|
||||
lines = []
|
||||
for i, seg in enumerate(segments, 1):
|
||||
start_time = self._seconds_to_srt_time(seg['start'])
|
||||
end_time = self._seconds_to_srt_time(seg['end'])
|
||||
lines.append(f"{i}")
|
||||
lines.append(f"{start_time} --> {end_time}")
|
||||
lines.append(seg['text'])
|
||||
lines.append("") # 空行
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
def _seconds_to_srt_time(self, seconds: float) -> str:
|
||||
"""将秒数转换为 SRT 时间格式 (HH:MM:SS,mmm)"""
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
millis = int((seconds % 1) * 1000)
|
||||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
|
||||
|
||||
async def _run_whisper(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行 Whisper 语音识别
|
||||
|
||||
根据 HF_TOKEN 配置自动选择普通模式或说话者分离模式。
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
"""
|
||||
hf_token = self.settings.hf_token.get_secret_value()
|
||||
if hf_token:
|
||||
return await self._run_whisper_with_diarization(input_data)
|
||||
else:
|
||||
return await self._run_whisper_basic(input_data)
|
||||
|
||||
async def _run_whisper_basic(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行普通 Whisper.cpp 语音识别
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
|
||||
Raises:
|
||||
RuntimeError: Whisper.cpp 路径不存在或执行失败
|
||||
"""
|
||||
video_path = Path(input_data.video_path)
|
||||
output_dir = video_path.parent
|
||||
base_name = video_path.stem
|
||||
|
||||
# 获取 Whisper.cpp 路径
|
||||
try:
|
||||
whisper_path = get_whisper_cpp_path()
|
||||
except WhisperNotFoundError as e:
|
||||
raise RuntimeError(str(e))
|
||||
|
||||
# 模型路径
|
||||
model_path = os.path.join(whisper_path, 'models', 'ggml-base.bin')
|
||||
if not os.path.exists(model_path):
|
||||
raise RuntimeError(f"Whisper model not found: {model_path}")
|
||||
|
||||
# 转换音频格式
|
||||
audio_path = convert_to_wav(str(video_path))
|
||||
temp_audio_created = audio_path != str(video_path)
|
||||
|
||||
try:
|
||||
# 构建 Whisper 命令
|
||||
output_base = os.path.join(output_dir, base_name)
|
||||
cmd = [
|
||||
os.path.join(whisper_path, 'whisper-cli'),
|
||||
'-m', model_path,
|
||||
'-f', audio_path,
|
||||
'-of', output_base,
|
||||
'-l', 'zh',
|
||||
'--output-txt', '--output-srt', '--output-json',
|
||||
]
|
||||
|
||||
# 执行 Whisper
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"Whisper failed: {result.stderr}")
|
||||
|
||||
# 解析输出
|
||||
return self._parse_whisper_output(output_dir, base_name)
|
||||
|
||||
finally:
|
||||
# 清理临时音频文件
|
||||
if temp_audio_created and os.path.exists(audio_path):
|
||||
try:
|
||||
os.remove(audio_path)
|
||||
except Exception:
|
||||
pass # 静默失败，不阻塞主流程
|
||||
|
||||
def _parse_whisper_output(self, output_dir: Path, base_name: str) -> ASROutput:
|
||||
"""解析 Whisper 输出文件
|
||||
|
||||
Args:
|
||||
output_dir: 输出目录
|
||||
base_name: 文件基础名
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
"""
|
||||
txt_path = output_dir / f"{base_name}.txt"
|
||||
srt_path = output_dir / f"{base_name}.srt"
|
||||
json_path = output_dir / f"{base_name}.json"
|
||||
|
||||
# 读取 JSON 输出
|
||||
segments = []
|
||||
text = ""
|
||||
duration = 0.0
|
||||
|
||||
if json_path.exists():
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
# Whisper.cpp JSON 格式
|
||||
if 'transcription' in data:
|
||||
for i, item in enumerate(data['transcription']):
|
||||
segments.append({
|
||||
'id': i,
|
||||
'start': item.get('offsets', [0, 0])[0] / 1000.0,
|
||||
'end': item.get('offsets', [0, 0])[1] / 1000.0,
|
||||
'text': item.get('text', '')
|
||||
})
|
||||
text = ' '.join(s['text'] for s in segments)
|
||||
if segments:
|
||||
duration = segments[-1]['end']
|
||||
|
||||
# 读取文本
|
||||
if txt_path.exists():
|
||||
text = txt_path.read_text(encoding='utf-8').strip()
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={
|
||||
'txt': str(txt_path),
|
||||
'srt': str(srt_path),
|
||||
'json': str(json_path)
|
||||
},
|
||||
segments=segments,
|
||||
text=text,
|
||||
duration_seconds=duration
|
||||
)
|
||||
|
||||
async def _run_whisper_with_diarization(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行带说话者分离的 Whisper
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果，包含说话者信息
|
||||
"""
|
||||
# 先运行普通 Whisper
|
||||
basic_result = await self._run_whisper_basic(input_data)
|
||||
|
||||
try:
|
||||
# 运行说话者分离
|
||||
diarization_result = run_diarization(input_data.video_path)
|
||||
|
||||
# 合并结果
|
||||
merged_segments = assign_speakers_to_segments(
|
||||
basic_result.segments,
|
||||
diarization_result
|
||||
)
|
||||
|
||||
# 提取说话者列表
|
||||
speakers = list(set(
|
||||
seg.get('speaker') for seg in merged_segments if seg.get('speaker')
|
||||
))
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="whisper-diarization",
|
||||
has_subtitle=False,
|
||||
output_files=basic_result.output_files,
|
||||
segments=merged_segments,
|
||||
text=basic_result.text,
|
||||
duration_seconds=basic_result.duration_seconds,
|
||||
speakers=sorted(speakers)
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
logger.warning("pyannote.audio not available, falling back to basic whisper")
|
||||
return basic_result
|
||||
except Exception as e:
|
||||
logger.warning(f"Diarization failed: {e}, falling back to basic whisper")
|
||||
return basic_result
|
||||
|
||||
|
||||
def run_diarization(audio_path: str, num_speakers: Optional[int] = None) -> Dict[str, List[Dict]]:
|
||||
"""运行 pyannote.audio 说话者分离
|
||||
|
||||
Args:
|
||||
audio_path: 音频文件路径
|
||||
num_speakers: 预期说话者数量（可选）
|
||||
|
||||
Returns:
|
||||
说话者时间线字典 {speaker_name: [{start, end}, ...]}
|
||||
|
||||
Raises:
|
||||
ImportError: 如果 pyannote.audio 未安装
|
||||
RuntimeError: 如果 HF_TOKEN 未配置
|
||||
"""
|
||||
hf_token = os.environ.get('HF_TOKEN')
|
||||
if not hf_token:
|
||||
raise RuntimeError("需要设置 HF_TOKEN 环境变量以使用说话者分离功能")
|
||||
|
||||
from huggingface_hub import login
|
||||
from pyannote.audio import Pipeline
|
||||
import torch
|
||||
|
||||
login(token=hf_token)
|
||||
|
||||
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
pipeline.to(device)
|
||||
|
||||
diarization = pipeline(audio_path, num_speakers=num_speakers)
|
||||
|
||||
# 解析结果
|
||||
speakers = {}
|
||||
for turn, _, speaker in diarization.itertracks(yield_label=True):
|
||||
if speaker not in speakers:
|
||||
speakers[speaker] = []
|
||||
speakers[speaker].append({
|
||||
"start": round(turn.start, 2),
|
||||
"end": round(turn.end, 2)
|
||||
})
|
||||
|
||||
return speakers
|
||||
|
||||
|
||||
def assign_speakers_to_segments(
|
||||
segments: List[Dict],
|
||||
diarization: Dict[str, List[Dict]]
|
||||
) -> List[Dict]:
|
||||
"""为每个段落分配说话者
|
||||
|
||||
根据时间戳匹配说话者分离结果到 Whisper 段落。
|
||||
|
||||
Args:
|
||||
segments: Whisper 段落列表
|
||||
diarization: 说话者分离结果
|
||||
|
||||
Returns:
|
||||
带有 speaker 字段的段落列表
|
||||
"""
|
||||
result = []
|
||||
|
||||
for seg in segments:
|
||||
seg_start = seg.get('start', 0)
|
||||
seg_end = seg.get('end', 0)
|
||||
seg_mid = (seg_start + seg_end) / 2
|
||||
|
||||
# 找到包含中点时间的说话者
|
||||
assigned_speaker = None
|
||||
for speaker, turns in diarization.items():
|
||||
for turn in turns:
|
||||
if turn['start'] <= seg_mid <= turn['end']:
|
||||
assigned_speaker = speaker
|
||||
break
|
||||
if assigned_speaker:
|
||||
break
|
||||
|
||||
new_seg = dict(seg)
|
||||
new_seg['speaker'] = assigned_speaker
|
||||
result.append(new_seg)
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user