340 lines
11 KiB
Python
340 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
KB Whisper Diarization - 语音识别 + 说话者分离
|
|
使用 Whisper 进行语音识别，pyannote.audio 进行说话者分离
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
import re
|
|
|
|
|
|
def parse_arguments():
|
|
"""解析命令行参数"""
|
|
parser = argparse.ArgumentParser(
|
|
description='KB Whisper Diarization - 语音识别 + 说话者分离',
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
示例:
|
|
python whisper_diarization.py --input video.mp4
|
|
python whisper_diarization.py --input interview.mp4 --diarize
|
|
python whisper_diarization.py --input meeting.mp4 --diarize --num-speakers 2
|
|
"""
|
|
)
|
|
parser.add_argument('--input', required=True, help='输入音视频文件路径')
|
|
parser.add_argument('--output', help='输出目录 (默认: 文件所在目录)')
|
|
parser.add_argument('--diarize', action='store_true', help='启用说话者分离')
|
|
parser.add_argument('--num-speakers', type=int, help='估计说话者数量')
|
|
parser.add_argument('--language', help='指定语言 (默认: 自动检测)')
|
|
parser.add_argument('--model', default='base', choices=['tiny', 'base', 'small', 'medium', 'large'],
|
|
help='Whisper 模型 (默认: base)')
|
|
parser.add_argument('--background', action='store_true', help='后台异步执行')
|
|
return parser.parse_args()
|
|
|
|
|
|
def get_whisper_path():
|
|
"""获取 whisper.cpp 路径"""
|
|
path = os.environ.get('WHISPER_CPP_PATH')
|
|
if path:
|
|
path = os.path.expanduser(path)
|
|
if os.path.exists(path):
|
|
return path
|
|
return os.path.expanduser('~/whisper.cpp')
|
|
|
|
|
|
def check_whisper_model(whisper_path, model):
|
|
"""检查 Whisper 模型"""
|
|
model_path = os.path.join(whisper_path, 'models', f'ggml-{model}.bin')
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Whisper 模型不存在: {model_path}")
|
|
return model_path
|
|
|
|
|
|
def convert_to_wav(input_path, output_dir):
|
|
"""转换为 16kHz WAV"""
|
|
if input_path.lower().endswith('.wav'):
|
|
try:
|
|
result = subprocess.run(
|
|
['ffprobe', '-v', 'error', '-show_entries', 'stream=sample_rate',
|
|
'-of', 'default=noprint_wrappers=1:nokey=1', input_path],
|
|
capture_output=True, text=True
|
|
)
|
|
if result.returncode == 0 and int(result.stdout.strip()) == 16000:
|
|
return input_path
|
|
except Exception:
|
|
pass
|
|
|
|
output_path = os.path.join(output_dir, f"{Path(input_path).stem}_audio.wav")
|
|
print(f"转换音频: {input_path} -> {output_path}")
|
|
|
|
subprocess.run([
|
|
'ffmpeg', '-y', '-i', input_path,
|
|
'-acodec', 'pcm_s16le', '-ar', '16000', '-ac', '1',
|
|
'-loglevel', 'error', output_path
|
|
], check=True, capture_output=True)
|
|
|
|
print("✓ 音频转换完成")
|
|
return output_path
|
|
|
|
|
|
def run_whisper(whisper_path, audio_path, model, output_dir, language=None):
|
|
"""运行 Whisper 识别"""
|
|
basename = Path(audio_path).stem.replace('_audio', '')
|
|
output_file = os.path.join(output_dir, basename)
|
|
|
|
cmd = [
|
|
os.path.join(whisper_path, 'whisper-cli'),
|
|
'-m', os.path.join(whisper_path, 'models', f'ggml-{model}.bin'),
|
|
'-f', audio_path,
|
|
'-of', output_file,
|
|
'-l', language or 'auto',
|
|
'--output-txt', '--output-srt', '--output-json',
|
|
]
|
|
|
|
print(f"执行 Whisper: {' '.join(cmd[:6])} ...")
|
|
|
|
result = subprocess.run(
|
|
cmd, capture_output=True, text=True, cwd=whisper_path, timeout=3600
|
|
)
|
|
|
|
if result.returncode != 0 and not os.path.exists(f"{output_file}.txt"):
|
|
raise RuntimeError(f"Whisper 失败: {result.stderr}")
|
|
|
|
return basename
|
|
|
|
|
|
def run_diarization(audio_path, num_speakers=None):
|
|
"""运行 pyannote.audio 说话者分离"""
|
|
print("\n🎙️ 步骤2: 说话者分离...")
|
|
|
|
try:
|
|
import torch
|
|
from pyannote.audio import Pipeline
|
|
except ImportError as e:
|
|
print(f"✗ 缺少依赖: {e}")
|
|
print("请安装: pip install pyannote.audio torch")
|
|
raise
|
|
|
|
# 检查 HuggingFace Token
|
|
hf_token = os.environ.get('HF_TOKEN')
|
|
if not hf_token:
|
|
raise RuntimeError("需要设置 HF_TOKEN 环境变量\n"
|
|
"获取方式: https://huggingface.co/settings/tokens")
|
|
|
|
# 新版 pyannote 需要先登录
|
|
from huggingface_hub import login
|
|
print("登录 HuggingFace...")
|
|
login(token=hf_token)
|
|
|
|
print("加载 pyannote 流水线...")
|
|
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")
|
|
|
|
# 设置设备
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"使用设备: {device}")
|
|
pipeline.to(device)
|
|
|
|
# 运行分离
|
|
print("执行说话者分离...")
|
|
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)
|
|
})
|
|
|
|
print(f"✓ 检测到 {len(speakers)} 个说话者: {list(speakers.keys())}")
|
|
return speakers
|
|
|
|
|
|
def merge_results(whisper_txt, whisper_srt, diarization, output_dir, basename):
|
|
"""合并 Whisper 和说话者分离结果"""
|
|
print("\n📝 步骤3: 合并结果...")
|
|
|
|
# 读取 Whisper 文本
|
|
text_segments = []
|
|
if os.path.exists(whisper_txt):
|
|
with open(whisper_txt, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
|
|
# 简单按行分段
|
|
for i, line in enumerate(content.strip().split('\n')):
|
|
if line.strip():
|
|
text_segments.append({
|
|
"id": i,
|
|
"text": line.strip(),
|
|
"start": i * 3.0, # 估算
|
|
"end": (i + 1) * 3.0
|
|
})
|
|
|
|
# 读取 SRT 获取精确时间
|
|
srt_segments = []
|
|
if os.path.exists(whisper_srt):
|
|
with open(whisper_srt, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
|
|
# 解析 SRT
|
|
pattern = r'(\d+)\n(\d{2}:\d{2}:\d{2},\d{3}) --> (\d{2}:\d{2}:\d{2},\d{3})\n(.+?)(?=\n\n|\Z)'
|
|
matches = re.findall(pattern, content, re.DOTALL)
|
|
|
|
def parse_time(t):
|
|
h, m, s_ms = t.split(':')
|
|
s, ms = s_ms.split(',')
|
|
return float(h) * 3600 + float(m) * 60 + float(s) + float(ms) / 1000
|
|
|
|
for idx, num, start, end, text in matches:
|
|
srt_segments.append({
|
|
"id": int(idx),
|
|
"start": parse_time(start),
|
|
"end": parse_time(end),
|
|
"text": text.strip()
|
|
})
|
|
|
|
# 为每个 Whisper 段落分配说话者
|
|
merged_segments = []
|
|
for seg in srt_segments:
|
|
speaker = assign_speaker(seg['start'], seg['end'], diarization)
|
|
merged_segments.append({
|
|
"id": len(merged_segments),
|
|
"speaker": speaker,
|
|
"start": seg['start'],
|
|
"end": seg['end'],
|
|
"text": seg['text']
|
|
})
|
|
|
|
# 生成带说话者的 SRT
|
|
output_srt = os.path.join(output_dir, f"{basename}_diarized.srt")
|
|
with open(output_srt, 'w', encoding='utf-8') as f:
|
|
for seg in merged_segments:
|
|
start = format_time(seg['start'])
|
|
end = format_time(seg['end'])
|
|
f.write(f"{seg['id'] + 1}\n")
|
|
f.write(f"{start} --> {end}\n")
|
|
f.write(f"[{seg['speaker']}] {seg['text']}\n\n")
|
|
|
|
# 生成合并的 TXT
|
|
output_txt = os.path.join(output_dir, f"{basename}_diarized.txt")
|
|
with open(output_txt, 'w', encoding='utf-8') as f:
|
|
for seg in merged_segments:
|
|
f.write(f"[{seg['speaker']}] {seg['text']}\n")
|
|
|
|
# 生成 JSON
|
|
output_json = os.path.join(output_dir, f"{basename}_diarized.json")
|
|
json_data = {
|
|
"success": True,
|
|
"tool": "whisper-diarization",
|
|
"diarization": True,
|
|
"speakers": list(diarization.keys()),
|
|
"input_file": basename,
|
|
"output_files": {
|
|
"txt": f"{basename}_diarized.txt",
|
|
"srt": f"{basename}_diarized.srt",
|
|
"json": f"{basename}_diarized.json"
|
|
},
|
|
"segments": merged_segments
|
|
}
|
|
with open(output_json, 'w', encoding='utf-8') as f:
|
|
json.dump(json_data, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"✓ 生成文件:")
|
|
print(f" - {output_srt}")
|
|
print(f" - {output_txt}")
|
|
print(f" - {output_json}")
|
|
|
|
return json_data
|
|
|
|
|
|
def assign_speaker(start, end, diarization):
|
|
"""将时间段分配给说话者"""
|
|
best_speaker = "UNKNOWN"
|
|
max_overlap = 0
|
|
|
|
for speaker, segments in diarization.items():
|
|
for seg in segments:
|
|
overlap = min(end, seg['end']) - max(start, seg['start'])
|
|
if overlap > max_overlap:
|
|
max_overlap = overlap
|
|
best_speaker = speaker
|
|
|
|
return best_speaker
|
|
|
|
|
|
def format_time(seconds):
|
|
"""秒数转换为 SRT 时间格式"""
|
|
hours = int(seconds // 3600)
|
|
minutes = int((seconds % 3600) // 60)
|
|
secs = int(seconds % 60)
|
|
ms = int((seconds % 1) * 1000)
|
|
return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"
|
|
|
|
|
|
def main():
|
|
args = parse_arguments()
|
|
|
|
input_file = os.path.abspath(args.input)
|
|
if not os.path.exists(input_file):
|
|
print(f"✗ 文件不存在: {input_file}")
|
|
sys.exit(1)
|
|
|
|
output_dir = os.path.abspath(args.output) if args.output else os.path.dirname(input_file)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
basename = Path(input_file).stem
|
|
|
|
print(f"\n📁 输入: {input_file}")
|
|
print(f"📁 输出: {output_dir}")
|
|
print(f"🎙️ 说话者分离: {'是' if args.diarize else '否'}")
|
|
|
|
try:
|
|
# 步骤1: Whisper 识别
|
|
print(f"\n🎙️ 步骤1: 语音识别...")
|
|
whisper_path = get_whisper_path()
|
|
check_whisper_model(whisper_path, args.model)
|
|
|
|
audio_path = convert_to_wav(input_file, output_dir)
|
|
run_whisper(whisper_path, audio_path, args.model, output_dir, args.language)
|
|
|
|
# 步骤2: 说话者分离
|
|
if args.diarize:
|
|
diarization = run_diarization(audio_path, args.num_speakers)
|
|
else:
|
|
diarization = {}
|
|
|
|
# 步骤3: 合并结果
|
|
whisper_txt = os.path.join(output_dir, f"{basename}.txt")
|
|
whisper_srt = os.path.join(output_dir, f"{basename}.srt")
|
|
result = merge_results(whisper_txt, whisper_srt, diarization, output_dir, basename)
|
|
|
|
# 清理临时音频
|
|
if audio_path != input_path and audio_path.endswith('_audio.wav'):
|
|
try:
|
|
os.remove(audio_path)
|
|
except Exception:
|
|
pass
|
|
|
|
print(f"\n✅ 处理完成!")
|
|
print(f" 说话者数量: {len(diarization) if args.diarize else 1}")
|
|
print(f" 段落数量: {len(result['segments'])}")
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ 错误: {e}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|