349 lines
9.8 KiB
Python
Executable File
349 lines
9.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
KB Whisper - 本地语音识别工具
|
|
基于 Whisper.cpp 的音视频转文字
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
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 audio.mp3 --output ./output --background
|
|
"""
|
|
)
|
|
parser.add_argument(
|
|
'--input',
|
|
required=True,
|
|
help='输入音视频文件路径'
|
|
)
|
|
parser.add_argument(
|
|
'--model',
|
|
default='tiny',
|
|
choices=['tiny', 'base', 'small', 'medium', 'large'],
|
|
help='Whisper 模型类型 (默认: tiny)'
|
|
)
|
|
parser.add_argument(
|
|
'--output',
|
|
help='输出目录 (默认: 输入文件所在目录)'
|
|
)
|
|
parser.add_argument(
|
|
'--background',
|
|
action='store_true',
|
|
help='后台异步执行'
|
|
)
|
|
parser.add_argument(
|
|
'--language',
|
|
help='指定语言 (默认: 自动检测)，例如: 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 run_whisper_cpp(whisper_path, input_file, model, output_dir, language=None):
|
|
"""运行 whisper.cpp 命令行工具"""
|
|
|
|
# 构建命令
|
|
cmd = [
|
|
os.path.join(whisper_path, 'main'),
|
|
'-m', os.path.join(whisper_path, 'models', f'ggml-{model}.bin'),
|
|
'-f', input_file,
|
|
'-of', os.path.join(output_dir, Path(input_file).stem),
|
|
'-otxt', # 输出 txt
|
|
'-osrt', # 输出 srt
|
|
'-oj', # 输出 json
|
|
]
|
|
|
|
# 添加语言参数（如果指定）
|
|
if language:
|
|
cmd.extend(['-l', language])
|
|
|
|
print(f"执行命令: {' '.join(cmd)}")
|
|
|
|
# 运行命令
|
|
result = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
cwd=whisper_path
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
raise RuntimeError(
|
|
f"Whisper.cpp 执行失败:\n"
|
|
f"标准输出: {result.stdout}\n"
|
|
f"错误输出: {result.stderr}"
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def parse_whisper_json(json_file):
|
|
"""解析 whisper.cpp 输出的 JSON 文件"""
|
|
try:
|
|
with open(json_file, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
return data
|
|
except json.JSONDecodeError:
|
|
# whisper.cpp 的 JSON 格式可能不同，尝试解析标准格式
|
|
return None
|
|
|
|
|
|
def create_standard_json(input_file, output_dir, srt_file, txt_file, model):
|
|
"""创建标准格式的 JSON 输出"""
|
|
basename = Path(input_file).stem
|
|
|
|
# 解析 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)
|
|
|
|
for idx, (num, start, end, text) in enumerate(matches):
|
|
# 转换时间戳
|
|
def parse_time(t):
|
|
h, m, s_ms = t.split(':')
|
|
s, ms = s_ms.split(',')
|
|
return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000
|
|
|
|
segments.append({
|
|
"id": idx,
|
|
"start": parse_time(start),
|
|
"end": parse_time(end),
|
|
"text": text.strip()
|
|
})
|
|
|
|
# 读取纯文本
|
|
txt_content = ""
|
|
if os.path.exists(txt_file):
|
|
with open(txt_file, 'r', encoding='utf-8') as f:
|
|
txt_content = f.read().strip()
|
|
|
|
# 获取音频时长（使用 ffprobe）
|
|
duration = 0
|
|
try:
|
|
result = subprocess.run(
|
|
['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
|
|
'-of', 'default=noprint_wrappers=1:nokey=1', input_file],
|
|
capture_output=True,
|
|
text=True
|
|
)
|
|
if result.returncode == 0:
|
|
duration = float(result.stdout.strip())
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
# 构建标准 JSON
|
|
json_data = {
|
|
"success": True,
|
|
"input_file": os.path.basename(input_file),
|
|
"input_path": input_file,
|
|
"model": model,
|
|
"output_files": {
|
|
"json": f"{basename}.json",
|
|
"srt": f"{basename}.srt",
|
|
"txt": f"{basename}.txt"
|
|
},
|
|
"duration_seconds": duration,
|
|
"segments": segments
|
|
}
|
|
|
|
return json_data
|
|
|
|
|
|
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("启动后台处理...")
|
|
|
|
# 使用 nohup 在后台运行
|
|
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] if args.language else []),
|
|
stdout=log,
|
|
stderr=log,
|
|
start_new_session=True
|
|
)
|
|
|
|
print(f"后台任务已启动，日志文件: {log_file}")
|
|
print(f"输出文件将生成到: {output_dir}")
|
|
|
|
# 返回任务信息
|
|
result = {
|
|
"success": True,
|
|
"background": True,
|
|
"status": "running",
|
|
"log_file": log_file,
|
|
"output_dir": output_dir,
|
|
"message": "后台处理中，请稍候检查输出文件"
|
|
}
|
|
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
return
|
|
|
|
# 前台执行
|
|
print(f"处理文件: {input_file}")
|
|
print(f"输出目录: {output_dir}")
|
|
|
|
try:
|
|
# 运行 whisper.cpp
|
|
result = run_whisper_cpp(
|
|
whisper_path,
|
|
input_file,
|
|
args.model,
|
|
output_dir,
|
|
args.language
|
|
)
|
|
|
|
# 输出文件路径
|
|
basename = Path(input_file).stem
|
|
json_file = os.path.join(output_dir, f"{basename}.json")
|
|
srt_file = os.path.join(output_dir, f"{basename}.srt")
|
|
txt_file = os.path.join(output_dir, f"{basename}.txt")
|
|
|
|
# 检查输出文件
|
|
output_files = {}
|
|
if os.path.exists(srt_file):
|
|
output_files['srt'] = srt_file
|
|
if os.path.exists(txt_file):
|
|
output_files['txt'] = txt_file
|
|
|
|
# 创建标准 JSON 格式
|
|
standard_json = create_standard_json(
|
|
input_file,
|
|
output_dir,
|
|
srt_file,
|
|
txt_file,
|
|
args.model
|
|
)
|
|
|
|
# 写入标准 JSON
|
|
with open(json_file, 'w', encoding='utf-8') as f:
|
|
json.dump(standard_json, f, ensure_ascii=False, indent=2)
|
|
|
|
output_files['json'] = json_file
|
|
|
|
# 输出结果
|
|
result = {
|
|
"success": True,
|
|
"input_file": input_file,
|
|
"model": args.model,
|
|
"output_files": {
|
|
k: os.path.basename(v) for k, v in output_files.items()
|
|
},
|
|
"output_dir": output_dir,
|
|
"total_segments": len(standard_json.get('segments', [])),
|
|
"duration": standard_json.get('duration_seconds', 0)
|
|
}
|
|
|
|
print("\n✓ 处理完成!")
|
|
print(f"\n输出文件:")
|
|
for fmt, fname in output_files.items():
|
|
print(f" - {fname.upper()}: {fname}")
|
|
|
|
print(f"\n完整结果:")
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ 错误: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|