diff --git a/kb_pipeline/processors/__init__.py b/kb_pipeline/processors/__init__.py new file mode 100644 index 0000000..be2f254 --- /dev/null +++ b/kb_pipeline/processors/__init__.py @@ -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"] diff --git a/kb_pipeline/processors/asr.py b/kb_pipeline/processors/asr.py new file mode 100644 index 0000000..824a94b --- /dev/null +++ b/kb_pipeline/processors/asr.py @@ -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