添加技能模块
This commit is contained in:
@@ -0,0 +1,403 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user