370 lines
14 KiB
Python
370 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
使用SenseVoice模型进行高精度中文语音识别
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import json
|
||
import argparse
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
import logging
|
||
|
||
# 设置日志
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
logger = logging.getLogger(__name__)
|
||
|
||
try:
|
||
from funasr import AutoModel
|
||
from funasr.utils.postprocess_utils import rich_transcription_postprocess
|
||
except ImportError:
|
||
logger.error("请先安装FunASR: pip install funasr")
|
||
raise
|
||
|
||
class SenseVoiceTranscriber:
|
||
"""使用SenseVoice模型进行高精度中文语音识别"""
|
||
|
||
def __init__(self, model_dir="/root/autodl-tmp/llm/sensevoice", device="cuda:0"):
|
||
"""
|
||
初始化SenseVoice转录器
|
||
|
||
Args:
|
||
model_dir: 模型路径或名称
|
||
device: 运行设备 (cuda:0, cpu)
|
||
"""
|
||
self.model_dir = model_dir
|
||
self.device = device
|
||
self.model = None
|
||
self.load_model()
|
||
|
||
def load_model(self):
|
||
"""加载SenseVoice模型"""
|
||
logger.info(f"正在加载SenseVoice模型: {self.model_dir}")
|
||
start_time = time.time()
|
||
|
||
try:
|
||
self.model = AutoModel(
|
||
model=self.model_dir,
|
||
trust_remote_code=True,
|
||
device=self.device,
|
||
disable_update=True
|
||
)
|
||
load_time = time.time() - start_time
|
||
logger.info(f"SenseVoice模型加载完成，耗时: {load_time:.2f} 秒")
|
||
except Exception as e:
|
||
logger.error(f"模型加载失败: {str(e)}")
|
||
raise
|
||
|
||
def transcribe_audio(self, audio_path, language="auto", use_itn=True):
|
||
"""
|
||
转录音频文件
|
||
|
||
Args:
|
||
audio_path: 音频文件路径
|
||
language: 语言设置 ("auto", "zh", "en", "yue", "ja", "ko")
|
||
use_itn: 是否使用逆文本标准化
|
||
|
||
Returns:
|
||
dict: 转录结果
|
||
"""
|
||
audio_path = Path(audio_path)
|
||
|
||
if not audio_path.exists():
|
||
raise FileNotFoundError(f"音频文件不存在: {audio_path}")
|
||
|
||
logger.info(f"开始转录音频文件: {audio_path}")
|
||
logger.info(f"语言设置: {language}, 使用ITN: {use_itn}")
|
||
|
||
start_time = time.time()
|
||
|
||
try:
|
||
# 执行转录
|
||
result = self.model.generate(
|
||
input=str(audio_path),
|
||
cache={},
|
||
language=language,
|
||
use_itn=use_itn,
|
||
batch_size_s=60, # 批处理大小（秒）
|
||
merge_vad=True, # 合并VAD结果
|
||
merge_length_s=15 # 合并长度（秒）
|
||
)
|
||
|
||
transcribe_time = time.time() - start_time
|
||
logger.info(f"转录完成，耗时: {transcribe_time:.2f} 秒")
|
||
|
||
# 直接返回原始结果
|
||
return {
|
||
"raw_result": result,
|
||
"transcribe_time": transcribe_time,
|
||
"model": "SenseVoice",
|
||
"timestamp": datetime.now().isoformat(),
|
||
"file_path": str(audio_path)
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"转录失败: {str(e)}")
|
||
raise
|
||
|
||
def segment_by_timestamp(self, result):
|
||
"""
|
||
根据时间戳将转录结果分段
|
||
|
||
Args:
|
||
result: 原始转录结果
|
||
|
||
Returns:
|
||
list: 分段后的结果列表
|
||
"""
|
||
if not isinstance(result, list):
|
||
return result
|
||
|
||
segmented_results = []
|
||
|
||
for item in result:
|
||
if isinstance(item, dict):
|
||
text = item.get("text", "")
|
||
timestamp = item.get("timestamp", [])
|
||
|
||
if timestamp and len(timestamp) >= 2:
|
||
# 有时间戳信息，进行分段
|
||
segment = {
|
||
"start_time": timestamp[0] / 1000.0, # 转换为秒
|
||
"end_time": timestamp[1] / 1000.0,
|
||
"duration": (timestamp[1] - timestamp[0]) / 1000.0,
|
||
"text": text,
|
||
"clean_text": self._clean_text(text)
|
||
}
|
||
segmented_results.append(segment)
|
||
else:
|
||
# 没有时间戳信息，直接添加
|
||
segment = {
|
||
"text": text,
|
||
"clean_text": self._clean_text(text)
|
||
}
|
||
segmented_results.append(segment)
|
||
|
||
return segmented_results
|
||
|
||
def _clean_text(self, text):
|
||
"""清理文本中的特殊标签"""
|
||
import re
|
||
# 移除SenseVoice的特殊标签
|
||
tags_pattern = r'<\|[^|]+\|>'
|
||
clean_text = re.sub(tags_pattern, '', text)
|
||
# 移除多余空格
|
||
clean_text = ' '.join(clean_text.split())
|
||
return clean_text.strip()
|
||
|
||
def transcribe_multiple_files(self, audio_files, output_dir="sensevoice_kuaishou", **kwargs):
|
||
"""批量转录多个音频文件"""
|
||
output_path = Path(output_dir)
|
||
output_path.mkdir(exist_ok=True)
|
||
|
||
results = []
|
||
|
||
for i, audio_file in enumerate(audio_files, 1):
|
||
logger.info(f"处理第 {i}/{len(audio_files)} 个文件")
|
||
|
||
try:
|
||
result = self.transcribe_audio(audio_file, **kwargs)
|
||
|
||
# 保存单个文件的结果
|
||
audio_name = Path(audio_file).stem
|
||
output_file = output_path / f"{audio_name}_sensevoice.json"
|
||
|
||
with open(output_file, 'w', encoding='utf-8') as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
|
||
logger.info(f"转录结果已保存: {output_file}")
|
||
results.append(result)
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理文件 {audio_file} 时出错: {str(e)}")
|
||
results.append(None)
|
||
|
||
return results
|
||
|
||
def save_transcript(self, result, output_path, format="json"):
|
||
"""保存转录结果"""
|
||
output_path = Path(output_path)
|
||
|
||
if format == "json":
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
|
||
elif format == "txt":
|
||
# 提取原始文本并清理
|
||
raw_result = result.get("raw_result", [])
|
||
if isinstance(raw_result, list):
|
||
texts = []
|
||
for item in raw_result:
|
||
if isinstance(item, dict):
|
||
text = item.get("text", "")
|
||
if text:
|
||
clean_text = self._clean_text(text)
|
||
texts.append(clean_text)
|
||
full_text = " ".join(texts)
|
||
else:
|
||
full_text = str(raw_result)
|
||
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
f.write(full_text)
|
||
|
||
elif format == "srt":
|
||
# 使用时间戳分段生成SRT
|
||
raw_result = result.get("raw_result", [])
|
||
segmented = self.segment_by_timestamp(raw_result)
|
||
self._save_srt_from_segments(segmented, output_path)
|
||
|
||
else:
|
||
raise ValueError(f"不支持的格式: {format}")
|
||
|
||
logger.info(f"转录结果已保存: {output_path}")
|
||
|
||
def _save_srt_from_segments(self, segments, output_path):
|
||
"""从分段结果保存为SRT字幕格式"""
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
for i, segment in enumerate(segments, 1):
|
||
if "start_time" in segment and "end_time" in segment:
|
||
start_time = self._format_time(segment["start_time"])
|
||
end_time = self._format_time(segment["end_time"])
|
||
text = segment.get("clean_text", segment.get("text", "")).strip()
|
||
|
||
f.write(f"{i}\n")
|
||
f.write(f"{start_time} --> {end_time}\n")
|
||
f.write(f"{text}\n\n")
|
||
|
||
def _format_time(self, seconds):
|
||
"""格式化时间为SRT格式"""
|
||
hours = int(seconds // 3600)
|
||
minutes = int((seconds % 3600) // 60)
|
||
seconds = seconds % 60
|
||
return f"{hours:02d}:{minutes:02d}:{seconds:06.3f}".replace('.', ',')
|
||
|
||
def main():
|
||
"""主函数"""
|
||
parser = argparse.ArgumentParser(description="使用SenseVoice模型进行高精度中文语音识别")
|
||
parser.add_argument("input", help="输入音频文件路径或目录")
|
||
parser.add_argument("-l", "--language", default="auto",
|
||
choices=["auto", "zh", "en", "yue", "ja", "ko"],
|
||
help="语言设置 (默认: auto)")
|
||
parser.add_argument("-o", "--output", default="/root/autodl-tmp/new_sensevoice",
|
||
help="输出目录 (默认: /root/autodl-tmp/sensevoice_kuaishou)")
|
||
parser.add_argument("-f", "--format", default="json",
|
||
choices=["json", "txt", "srt"],
|
||
help="输出格式 (默认: json)")
|
||
parser.add_argument("--device", default="cuda:0",
|
||
help="运行设备 (默认: cuda:0)")
|
||
parser.add_argument("--no-itn", action="store_true",
|
||
help="禁用逆文本标准化")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 创建转录器
|
||
transcriber = SenseVoiceTranscriber(device=args.device)
|
||
|
||
input_path = Path(args.input)
|
||
|
||
try:
|
||
if input_path.is_file():
|
||
# 处理单个文件
|
||
logger.info("处理单个音频文件")
|
||
result = transcriber.transcribe_audio(
|
||
args.input,
|
||
language=args.language,
|
||
use_itn=not args.no_itn
|
||
)
|
||
|
||
# 保存结果
|
||
output_dir = Path(args.output)
|
||
output_dir.mkdir(exist_ok=True)
|
||
|
||
file_stem = input_path.stem
|
||
output_file = output_dir / f"{file_stem}_sensevoice.{args.format}"
|
||
|
||
transcriber.save_transcript(result, output_file, args.format)
|
||
|
||
# 显示结果
|
||
print(f"\n转录完成!")
|
||
print(f"原文件: {args.input}")
|
||
print(f"输出文件: {output_file}")
|
||
print(f"转录时长: {result['transcribe_time']:.2f} 秒")
|
||
|
||
# 显示原始结果
|
||
raw_result = result.get("raw_result", [])
|
||
print(f"\n原始转录结果:")
|
||
print("-" * 50)
|
||
print(raw_result)
|
||
|
||
# 显示分段结果
|
||
if isinstance(raw_result, list):
|
||
segmented = transcriber.segment_by_timestamp(raw_result)
|
||
print(f"\n分段结果 (共 {len(segmented)} 段):")
|
||
print("-" * 50)
|
||
for i, segment in enumerate(segmented, 1):
|
||
if "start_time" in segment:
|
||
print(f"段 {i}: [{segment['start_time']:.2f}s - {segment['end_time']:.2f}s]")
|
||
else:
|
||
print(f"段 {i}: [无时间戳]")
|
||
print(f"原文: {segment.get('text', '')}")
|
||
print(f"清理: {segment.get('clean_text', '')}")
|
||
print("-" * 30)
|
||
|
||
# 显示完整清理文本
|
||
if isinstance(raw_result, list):
|
||
texts = []
|
||
for item in raw_result:
|
||
if isinstance(item, dict):
|
||
text = item.get("text", "")
|
||
if text:
|
||
clean_text = transcriber._clean_text(text)
|
||
texts.append(clean_text)
|
||
full_text = " ".join(texts)
|
||
print(f"\n完整清理文本:")
|
||
print("-" * 50)
|
||
print(full_text)
|
||
|
||
elif input_path.is_dir():
|
||
# 处理目录中的所有音频文件
|
||
logger.info("处理目录中的音频文件")
|
||
|
||
audio_extensions = ['.wav', '.mp3', '.m4a', '.flac', '.aac', '.ogg']
|
||
audio_files = []
|
||
|
||
for ext in audio_extensions:
|
||
audio_files.extend(input_path.glob(f"*{ext}"))
|
||
audio_files.extend(input_path.glob(f"*{ext.upper()}"))
|
||
|
||
if not audio_files:
|
||
logger.warning(f"在目录 {args.input} 中未找到音频文件")
|
||
return
|
||
|
||
logger.info(f"找到 {len(audio_files)} 个音频文件")
|
||
|
||
results = transcriber.transcribe_multiple_files(
|
||
[str(f) for f in audio_files],
|
||
output_dir=args.output,
|
||
language=args.language,
|
||
use_itn=not args.no_itn
|
||
)
|
||
|
||
# 统计结果
|
||
success_count = sum(1 for r in results if r is not None)
|
||
print(f"\n批量转录完成!")
|
||
print(f"总文件数: {len(audio_files)}")
|
||
print(f"成功转录: {success_count}")
|
||
print(f"失败: {len(audio_files) - success_count}")
|
||
print(f"输出目录: {args.output}")
|
||
|
||
else:
|
||
logger.error(f"输入路径无效: {args.input}")
|
||
return
|
||
|
||
except Exception as e:
|
||
logger.error(f"程序执行出错: {str(e)}")
|
||
return
|
||
|
||
if __name__ == "__main__":
|
||
main() |