404 lines
12 KiB
Python
404 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
KB Whisper - 本地语音识别工具
|
|
基于 Whisper.cpp 的音视频转文字
|
|
|
|
功能：
|
|
- 自动将视频/音频转换为 16kHz WAV 格式
|
|
- 支持指定语言（默认中文）
|
|
- 输出标准 JSON 格式供 KB Workflow 使用
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
import re
|
|
|
|
|
|
def parse_arguments():
|
|
"""解析命令行参数"""
|
|
parser = argparse.ArgumentParser(
|
|
description='KB Whisper - 本地语音识别工具',
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
示例:
|
|
python whisper.py --input video.mp4
|
|
python whisper.py --input video.mp4 --model base
|
|
python whisper.py --input video.mp4 --output ./output --background
|
|
"""
|
|
)
|
|
parser.add_argument(
|
|
'--input',
|
|
required=True,
|
|
help='输入音视频文件路径'
|
|
)
|
|
parser.add_argument(
|
|
'--model',
|
|
default='base',
|
|
choices=['tiny', 'base', 'small', 'medium', 'large'],
|
|
help='Whisper 模型类型 (默认: base)'
|
|
)
|
|
parser.add_argument(
|
|
'--output',
|
|
help='输出目录 (默认: 输入文件所在目录)'
|
|
)
|
|
parser.add_argument(
|
|
'--background',
|
|
action='store_true',
|
|
help='后台异步执行'
|
|
)
|
|
parser.add_argument(
|
|
'--language',
|
|
default='zh',
|
|
help='指定语言 (默认: zh)，例如: en, zh, ja'
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def get_whisper_cpp_path():
|
|
"""获取 whisper.cpp 路径"""
|
|
# 优先从环境变量获取
|
|
path = os.environ.get('WHISPER_CPP_PATH')
|
|
if path:
|
|
path = os.path.expanduser(path)
|
|
if os.path.exists(path):
|
|
return path
|
|
|
|
# 尝试常见位置
|
|
common_paths = [
|
|
os.path.expanduser('~/whisper.cpp'),
|
|
'/usr/local/bin/whisper',
|
|
'/opt/whisper.cpp',
|
|
]
|
|
|
|
for cp in common_paths:
|
|
if os.path.exists(cp):
|
|
return cp
|
|
|
|
raise FileNotFoundError(
|
|
"未找到 whisper.cpp。请设置环境变量 WHISPER_CPP_PATH "
|
|
"或安装到 ~/whisper.cpp\n"
|
|
"安装方法: git clone https://github.com/ggerganov/whisper.cpp.git ~/whisper.cpp && cd ~/whisper.cpp && make"
|
|
)
|
|
|
|
|
|
def check_model(whisper_path, model):
|
|
"""检查模型是否存在"""
|
|
model_path = os.path.join(whisper_path, 'models', f'ggml-{model}.bin')
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(
|
|
f"模型文件不存在: {model_path}\n"
|
|
f"请运行以下命令下载模型:\n"
|
|
f"cd {whisper_path} && ./models/download-ggml-model.sh {model}"
|
|
)
|
|
return model_path
|
|
|
|
|
|
def convert_to_wav(input_path, output_dir):
|
|
"""将音视频转换为 16kHz WAV 格式
|
|
|
|
Whisper 需要 WAV 格式的输入文件。
|
|
"""
|
|
input_path = os.path.abspath(input_path)
|
|
|
|
# 如果已经是 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:
|
|
print(f"输入文件已是 16kHz WAV 格式，直接使用")
|
|
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}")
|
|
|
|
try:
|
|
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
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"音频转换失败: {e}")
|
|
raise RuntimeError(f"无法转换音频格式: {e}")
|
|
except FileNotFoundError:
|
|
raise RuntimeError("未找到 ffmpeg，请安装: sudo apt-get install ffmpeg")
|
|
|
|
|
|
def run_whisper(whisper_path, audio_path, model, output_dir, language='zh'):
|
|
"""运行 Whisper 语音识别
|
|
|
|
Args:
|
|
whisper_path: whisper.cpp 目录
|
|
audio_path: WAV 音频文件路径
|
|
model: 模型名称
|
|
output_dir: 输出目录
|
|
language: 语言代码
|
|
"""
|
|
# 输出文件名（使用原始输入文件名，不是 audio.wav）
|
|
# 从 output_dir 和 audio_path 推断原始文件名
|
|
original_stem = Path(audio_path).stem.replace('_audio', '')
|
|
output_basename = original_stem
|
|
|
|
# 构建命令
|
|
cmd = [
|
|
os.path.join(whisper_path, 'whisper-cli'),
|
|
'-m', os.path.join(whisper_path, 'models', f'ggml-{model}.bin'),
|
|
'-f', audio_path,
|
|
'-of', os.path.join(output_dir, output_basename),
|
|
'-l', language,
|
|
'--output-txt',
|
|
'--output-srt',
|
|
'--output-json',
|
|
]
|
|
|
|
print(f"执行命令: {' '.join(cmd)}")
|
|
|
|
# 运行
|
|
result = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
cwd=whisper_path,
|
|
timeout=3600 # 1小时超时
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
# 检查是否有警告但仍成功
|
|
if "error" in result.stderr.lower() and not os.path.exists(
|
|
os.path.join(output_dir, f"{output_basename}.txt")
|
|
):
|
|
raise RuntimeError(
|
|
f"Whisper 执行失败:\n"
|
|
f"标准输出: {result.stdout}\n"
|
|
f"错误输出: {result.stderr}"
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def parse_srt_for_duration(srt_file):
|
|
"""从 SRT 文件解析时长"""
|
|
if not os.path.exists(srt_file):
|
|
return 0
|
|
|
|
with open(srt_file, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
|
|
# 找到最后一个时间戳
|
|
pattern = r'(\d{2}:\d{2}:\d{2}),\d{3} --> \d{2}:\d{2}:\d{2},\d{3}'
|
|
matches = re.findall(pattern, content)
|
|
|
|
if not matches:
|
|
return 0
|
|
|
|
# 解析最后一个时间戳
|
|
last_time = matches[-1]
|
|
h, m, s = last_time.split(':')
|
|
return int(h) * 3600 + int(m) * 60 + int(s)
|
|
|
|
|
|
def create_workflow_json(input_file, output_dir, model, language, duration):
|
|
"""创建标准格式的 JSON 输出，供 KB Workflow 使用"""
|
|
input_path = Path(input_file)
|
|
basename = input_path.stem
|
|
|
|
txt_file = os.path.join(output_dir, f"{basename}.txt")
|
|
srt_file = os.path.join(output_dir, f"{basename}.srt")
|
|
json_file = os.path.join(output_dir, f"{basename}.json")
|
|
|
|
# 读取文本内容
|
|
text_content = ""
|
|
if os.path.exists(txt_file):
|
|
with open(txt_file, 'r', encoding='utf-8') as f:
|
|
text_content = f.read().strip()
|
|
|
|
# 解析 SRT 获取分段
|
|
segments = []
|
|
if os.path.exists(srt_file):
|
|
with open(srt_file, '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 enumerate(matches):
|
|
segments.append({
|
|
"id": idx,
|
|
"start": round(parse_time(start), 2),
|
|
"end": round(parse_time(end), 2),
|
|
"text": text.strip()
|
|
})
|
|
|
|
# 构建 JSON
|
|
workflow_json = {
|
|
"success": True,
|
|
"tool": "whisper",
|
|
"model": model,
|
|
"language": language,
|
|
"input_file": str(input_path),
|
|
"output_files": {
|
|
"txt": f"{basename}.txt",
|
|
"srt": f"{basename}.srt",
|
|
"json": f"{basename}.json"
|
|
},
|
|
"duration_seconds": duration,
|
|
"segments": segments,
|
|
"text": text_content
|
|
}
|
|
|
|
# 写入 JSON
|
|
with open(json_file, 'w', encoding='utf-8') as f:
|
|
json.dump(workflow_json, f, ensure_ascii=False, indent=2)
|
|
|
|
return workflow_json
|
|
|
|
|
|
def main():
|
|
"""主函数"""
|
|
args = parse_arguments()
|
|
|
|
# 检查输入文件
|
|
input_file = os.path.abspath(args.input)
|
|
if not os.path.exists(input_file):
|
|
print(f"错误: 输入文件不存在: {input_file}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# 获取输出目录
|
|
if args.output:
|
|
output_dir = os.path.abspath(args.output)
|
|
else:
|
|
output_dir = os.path.dirname(input_file)
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 获取 whisper.cpp 路径
|
|
try:
|
|
whisper_path = get_whisper_cpp_path()
|
|
print(f"✓ Whisper.cpp 路径: {whisper_path}")
|
|
except FileNotFoundError as e:
|
|
print(f"✗ 错误: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# 检查模型
|
|
try:
|
|
check_model(whisper_path, args.model)
|
|
print(f"✓ 使用模型: {args.model}")
|
|
except FileNotFoundError as e:
|
|
print(f"✗ 错误: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# 后台执行
|
|
if args.background:
|
|
print("启动后台处理...")
|
|
log_file = os.path.join(output_dir, f"{Path(input_file).stem}_whisper.log")
|
|
|
|
with open(log_file, 'w') as log:
|
|
subprocess.Popen(
|
|
[sys.executable, __file__,
|
|
'--input', input_file,
|
|
'--model', args.model,
|
|
'--output', output_dir,
|
|
'--language', args.language],
|
|
stdout=log,
|
|
stderr=log,
|
|
start_new_session=True
|
|
)
|
|
|
|
print(f"✓ 后台任务已启动")
|
|
print(f" 日志: {log_file}")
|
|
print(f" 输出: {output_dir}")
|
|
|
|
result = {
|
|
"success": True,
|
|
"background": True,
|
|
"status": "running",
|
|
"log_file": log_file,
|
|
"output_dir": output_dir
|
|
}
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
return
|
|
|
|
# 前台执行
|
|
print(f"\n📁 处理文件: {input_file}")
|
|
print(f"📁 输出目录: {output_dir}")
|
|
|
|
try:
|
|
# 1. 转换为 WAV
|
|
print(f"\n🔄 步骤1: 转换音频格式...")
|
|
audio_path = convert_to_wav(input_file, output_dir)
|
|
|
|
# 2. 运行 Whisper
|
|
print(f"\n🎙️ 步骤2: 语音识别 (语言: {args.language})...")
|
|
run_whisper(whisper_path, audio_path, args.model, output_dir, args.language)
|
|
|
|
# 3. 清理临时 WAV 文件
|
|
if audio_path != input_file and audio_path.endswith('_audio.wav'):
|
|
try:
|
|
os.remove(audio_path)
|
|
print("✓ 清理临时文件")
|
|
except Exception:
|
|
pass
|
|
|
|
# 4. 获取时长并创建 JSON
|
|
print(f"\n📝 步骤3: 生成输出文件...")
|
|
basename = Path(input_file).stem
|
|
srt_file = os.path.join(output_dir, f"{basename}.srt")
|
|
duration = parse_srt_for_duration(srt_file)
|
|
|
|
workflow_json = create_workflow_json(
|
|
input_file, output_dir, args.model, args.language, duration
|
|
)
|
|
|
|
# 输出结果
|
|
print(f"\n✅ 处理完成!")
|
|
print(f" 时长: {duration} 秒")
|
|
print(f" 分段数: {len(workflow_json['segments'])}")
|
|
print(f"\n📄 输出文件:")
|
|
for fmt in ['txt', 'srt', 'json']:
|
|
fpath = os.path.join(output_dir, f"{basename}.{fmt}")
|
|
if os.path.exists(fpath):
|
|
size = os.path.getsize(fpath)
|
|
print(f" • {fmt.upper()}: {fpath} ({size:,} bytes)")
|
|
|
|
print(f"\n📋 Workflow JSON:")
|
|
print(json.dumps({
|
|
"success": workflow_json["success"],
|
|
"tool": "whisper",
|
|
"text_length": len(workflow_json["text"]),
|
|
"segments_count": len(workflow_json["segments"])
|
|
}, ensure_ascii=False, indent=2))
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ 错误: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|