starmen1265
29dd236bb9
- 将步骤5从 OCR 改为 Whisper 语音识别
- 更新 kb-whisper 支持自动音频转换
- 默认使用中文语言 (zh) 和 base 模型
- 更新输出格式为 {platform}_{id}.txt/.srt/.json
- 添加 FunASR 作为备用方案
314 lines
9.8 KiB
Python
Executable File
314 lines
9.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
KB FunASR - 本地语音识别工具
|
|
基于阿里 FunASR 实现音视频转文字
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
|
|
|
|
class FunASRTranscriber:
|
|
"""FunASR 语音识别器"""
|
|
|
|
# 模型映射
|
|
MODEL_MAP = {
|
|
"paraformer-zh": "damo/speech_paraformer-large-vad_punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
|
|
"paraformer": "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
|
|
"sensevoice": "FunASR/SenseVoiceSmall"
|
|
}
|
|
|
|
def __init__(self, model: str = "paraformer-zh", cache_dir: Optional[str] = None):
|
|
"""初始化 FunASR 识别器
|
|
|
|
Args:
|
|
model: 模型名称
|
|
cache_dir: 模型缓存目录
|
|
"""
|
|
self.model_name = model
|
|
self.model_id = self.MODEL_MAP.get(model, self.MODEL_MAP["paraformer-zh"])
|
|
self.cache_dir = cache_dir or os.path.expanduser("~/.cache/modelscope/hub")
|
|
self._model = None
|
|
|
|
def _load_model(self):
|
|
"""延迟加载模型"""
|
|
if self._model is None:
|
|
try:
|
|
from funasr import AutoModel
|
|
print(f"正在加载模型: {self.model_name} ({self.model_id})")
|
|
print(f"缓存目录: {self.cache_dir}")
|
|
self._model = AutoModel(
|
|
model=self.model_id,
|
|
disable_pbar=False,
|
|
disable_log=False,
|
|
)
|
|
print("模型加载完成")
|
|
except ImportError:
|
|
print("错误: 未安装 funasr 库")
|
|
print("请运行: pip install funasr modelscope")
|
|
sys.exit(1)
|
|
|
|
def _convert_audio(self, input_path: str) -> str:
|
|
"""转换音频格式为 16kHz WAV
|
|
|
|
Args:
|
|
input_path: 输入文件路径
|
|
|
|
Returns:
|
|
转换后的文件路径（如果不需要转换则返回原路径）
|
|
"""
|
|
# 检查是否已经是 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:
|
|
sample_rate = int(result.stdout.strip())
|
|
if sample_rate == 16000:
|
|
return input_path
|
|
except Exception:
|
|
pass
|
|
|
|
# 需要转换
|
|
output_path = input_path.rsplit('.', 1)[0] + '_converted.wav'
|
|
print(f"正在转换音频格式: {input_path} -> {output_path}")
|
|
|
|
try:
|
|
subprocess.run([
|
|
'ffmpeg', '-y', '-i', input_path,
|
|
'-acodec', 'pcm_s16le',
|
|
'-ar', '16000',
|
|
'-ac', '1',
|
|
output_path
|
|
], check=True, capture_output=True)
|
|
print("音频转换完成")
|
|
return output_path
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"音频转换失败: {e}")
|
|
# 如果转换失败，尝试使用原文件
|
|
print("尝试使用原文件继续处理")
|
|
return input_path
|
|
except FileNotFoundError:
|
|
print("警告: 未找到 ffmpeg，跳过音频转换")
|
|
return input_path
|
|
|
|
def transcribe(self, input_path: str) -> Dict:
|
|
"""转录音频文件
|
|
|
|
Args:
|
|
input_path: 输入文件路径
|
|
|
|
Returns:
|
|
识别结果字典
|
|
"""
|
|
input_path = os.path.abspath(input_path)
|
|
|
|
if not os.path.exists(input_path):
|
|
return {
|
|
"success": False,
|
|
"error": f"文件不存在: {input_path}"
|
|
}
|
|
|
|
print(f"\n开始处理: {input_path}")
|
|
print(f"使用模型: {self.model_name}")
|
|
|
|
# 转换音频格式
|
|
audio_path = self._convert_audio(input_path)
|
|
|
|
# 加载模型
|
|
self._load_model()
|
|
|
|
# 执行识别
|
|
print("正在进行语音识别...")
|
|
try:
|
|
result = self._model.generate(
|
|
input=audio_path,
|
|
batch_size_s=300,
|
|
)
|
|
except Exception as e:
|
|
return {
|
|
"success": False,
|
|
"error": f"识别失败: {str(e)}"
|
|
}
|
|
|
|
# 解析结果
|
|
if not result or not isinstance(result, list):
|
|
return {
|
|
"success": False,
|
|
"error": "识别结果格式错误"
|
|
}
|
|
|
|
# FunASR 返回格式处理
|
|
segments = []
|
|
full_text = ""
|
|
|
|
for seg in result:
|
|
if isinstance(seg, dict):
|
|
text = seg.get("text", "").strip()
|
|
if "timestamp" in seg and seg["timestamp"]:
|
|
# 包含时间戳
|
|
timestamp = seg["timestamp"]
|
|
if len(timestamp) >= 2:
|
|
start_time = timestamp[0][0] if timestamp[0] else 0
|
|
end_time = timestamp[-1][1] if timestamp[-1] else start_time
|
|
else:
|
|
start_time = 0
|
|
end_time = 0
|
|
else:
|
|
start_time = 0
|
|
end_time = 0
|
|
|
|
if text:
|
|
segments.append({
|
|
"id": len(segments),
|
|
"start": round(start_time, 2),
|
|
"end": round(end_time, 2),
|
|
"text": text
|
|
})
|
|
full_text += text + "\n"
|
|
|
|
# 清理临时文件
|
|
if audio_path != input_path and "_converted.wav" in audio_path:
|
|
try:
|
|
os.remove(audio_path)
|
|
except Exception:
|
|
pass
|
|
|
|
# 计算时长
|
|
duration = segments[-1]["end"] if segments else 0
|
|
|
|
return {
|
|
"success": True,
|
|
"input_file": os.path.basename(input_path),
|
|
"model": self.model_name,
|
|
"output_files": {},
|
|
"duration_seconds": round(duration, 2),
|
|
"segments": segments,
|
|
"text": full_text.strip()
|
|
}
|
|
|
|
def save_outputs(self, result: Dict, output_dir: str):
|
|
"""保存输出文件
|
|
|
|
Args:
|
|
result: 识别结果
|
|
output_dir: 输出目录
|
|
"""
|
|
if not result.get("success"):
|
|
return
|
|
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
basename = Path(result["input_file"]).stem
|
|
|
|
# 保存 JSON
|
|
json_path = output_dir / f"{basename}.json"
|
|
with open(json_path, "w", encoding="utf-8") as f:
|
|
json.dump(result, f, ensure_ascii=False, indent=2)
|
|
print(f"已保存: {json_path}")
|
|
result["output_files"]["json"] = str(json_path)
|
|
|
|
# 保存 TXT
|
|
txt_path = output_dir / f"{basename}.txt"
|
|
with open(txt_path, "w", encoding="utf-8") as f:
|
|
f.write(result.get("text", ""))
|
|
print(f"已保存: {txt_path}")
|
|
result["output_files"]["txt"] = str(txt_path)
|
|
|
|
|
|
def run_as_background(input_path: str, model: str, output_dir: str):
|
|
"""后台运行识别
|
|
|
|
Args:
|
|
input_path: 输入文件路径
|
|
model: 模型名称
|
|
output_dir: 输出目录
|
|
"""
|
|
# 创建日志文件
|
|
log_file = os.path.join(output_dir, f"funasr_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
|
|
|
|
# 重新执行命令（后台模式）
|
|
cmd = [
|
|
sys.executable, __file__,
|
|
"--input", input_path,
|
|
"--model", model,
|
|
"--output", output_dir
|
|
]
|
|
|
|
with open(log_file, "w") as f:
|
|
process = subprocess.Popen(
|
|
cmd,
|
|
stdout=f,
|
|
stderr=subprocess.STDOUT,
|
|
start_new_session=True
|
|
)
|
|
|
|
print(f"\n✓ 后台任务已启动 (PID: {process.pid})")
|
|
print(f"日志文件: {log_file}")
|
|
print(f"输出目录: {output_dir}")
|
|
print("\n使用以下命令查看进度:")
|
|
print(f" tail -f {log_file}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="KB FunASR - 本地语音识别工具",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
示例:
|
|
python funasr.py --input video.mp4
|
|
python funasr.py --input podcast.mp3 --model paraformer-zh
|
|
python funasr.py --input video.mp4 --output ./transcripts --background
|
|
"""
|
|
)
|
|
|
|
parser.add_argument("--input", "-i", required=True, help="输入音视频文件路径")
|
|
parser.add_argument("--model", "-m", default="paraformer-zh",
|
|
choices=["paraformer-zh", "paraformer", "sensevoice"],
|
|
help="模型名称 (默认: paraformer-zh)")
|
|
parser.add_argument("--output", "-o", default=None, help="输出目录 (默认: 文件所在目录)")
|
|
parser.add_argument("--background", "-b", action="store_true", help="后台异步执行")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# 确定输出目录
|
|
if args.output:
|
|
output_dir = args.output
|
|
else:
|
|
output_dir = os.path.dirname(args.input) or "."
|
|
|
|
# 后台模式
|
|
if args.background:
|
|
run_as_background(args.input, args.model, output_dir)
|
|
return
|
|
|
|
# 前台模式
|
|
transcriber = FunASRTranscriber(model=args.model)
|
|
result = transcriber.transcribe(args.input)
|
|
|
|
if result.get("success"):
|
|
transcriber.save_outputs(result, output_dir)
|
|
print(f"\n✓ 处理完成!")
|
|
print(f" 时长: {result['duration_seconds']} 秒")
|
|
print(f" 分段数: {len(result['segments'])}")
|
|
else:
|
|
print(f"\n✗ 处理失败: {result.get('error')}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|