generated from dellevin/template
339 lines
10 KiB
Python
339 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
STT 语音识别蓝图
|
||
使用 faster-whisper 实现音频转文字
|
||
"""
|
||
import os
|
||
import json
|
||
import uuid
|
||
import subprocess
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from datetime import timedelta
|
||
from flask import Blueprint, render_template, request, jsonify
|
||
|
||
bp = Blueprint('stt', __name__, url_prefix='/stt')
|
||
|
||
# HuggingFace 镜像
|
||
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
|
||
|
||
# OpenCC 繁体转简体
|
||
try:
|
||
from opencc import OpenCC
|
||
_cc_t2s = OpenCC('t2s')
|
||
except Exception:
|
||
_cc_t2s = None
|
||
|
||
try:
|
||
from config import BASE_DIR
|
||
except ImportError:
|
||
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
CONFIG_PATH = os.path.join(BASE_DIR, 'config', 'stt_config.json')
|
||
|
||
# 模型缓存
|
||
_model_cache = {}
|
||
_task_queue = []
|
||
_task_progress = {}
|
||
_task_results = {}
|
||
_worker_started = False
|
||
_worker_lock = threading.Lock()
|
||
|
||
LANGUAGES = {
|
||
'auto': '自动检测',
|
||
'zh': '中文', 'en': '英语', 'ja': '日语', 'ko': '韩语',
|
||
'fr': '法语', 'de': '德语', 'es': '西班牙语', 'ru': '俄语',
|
||
'th': '泰语', 'it': '意大利语', 'pt': '葡萄牙语', 'vi': '越南语',
|
||
'ar': '阿拉伯语', 'tr': '土耳其语',
|
||
}
|
||
|
||
|
||
def _load_config():
|
||
cfg = {'model_dir': ''}
|
||
if os.path.exists(CONFIG_PATH):
|
||
with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
|
||
cfg.update(json.load(f))
|
||
return cfg
|
||
|
||
|
||
def _save_config(cfg):
|
||
with open(CONFIG_PATH, 'w', encoding='utf-8') as f:
|
||
json.dump(cfg, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def _scan_models(directory):
|
||
"""扫描模型目录，提取已下载的模型名称"""
|
||
models = []
|
||
if not directory or not os.path.isdir(directory):
|
||
return models
|
||
prefix = 'models--Systran--faster-whisper-'
|
||
for name in os.listdir(directory):
|
||
full = os.path.join(directory, name)
|
||
if os.path.isdir(full) and name.startswith(prefix):
|
||
model_name = name[len(prefix):]
|
||
if model_name:
|
||
models.append(model_name)
|
||
return sorted(models)
|
||
|
||
|
||
@bp.route('/')
|
||
def page():
|
||
cfg = _load_config()
|
||
return render_template('stt.html', models=_scan_models(cfg['model_dir']), languages=LANGUAGES)
|
||
|
||
|
||
@bp.route('/config', methods=['GET'])
|
||
def get_config():
|
||
cfg = _load_config()
|
||
cfg['models'] = _scan_models(cfg['model_dir'])
|
||
return jsonify(cfg)
|
||
|
||
|
||
@bp.route('/config', methods=['POST'])
|
||
def save_config():
|
||
data = request.get_json()
|
||
cfg = _load_config()
|
||
if 'model_dir' in data:
|
||
cfg['model_dir'] = data['model_dir'].strip()
|
||
_save_config(cfg)
|
||
return jsonify({'success': True})
|
||
|
||
|
||
@bp.route('/models')
|
||
def list_models():
|
||
cfg = _load_config()
|
||
return jsonify({'models': _scan_models(cfg['model_dir'])})
|
||
|
||
|
||
@bp.route('/cuda-check')
|
||
def cuda_check():
|
||
"""检查 CUDA 是否可用（ctranslate2 优先，PyTorch 兜底）"""
|
||
# 方法1: ctranslate2（faster-whisper 的实际后端）
|
||
try:
|
||
import ctranslate2
|
||
count = ctranslate2.get_cuda_device_count()
|
||
if count > 0:
|
||
name = 'CUDA Device'
|
||
try:
|
||
name = ctranslate2.get_cuda_device_name(0) or name
|
||
except Exception:
|
||
pass
|
||
return jsonify({'cuda': True, 'device_count': count, 'name': name})
|
||
except Exception:
|
||
pass
|
||
# 方法2: PyTorch
|
||
try:
|
||
import torch
|
||
if torch.cuda.is_available():
|
||
return jsonify({'cuda': True, 'device_count': torch.cuda.device_count(),
|
||
'name': torch.cuda.get_device_name(0)})
|
||
except Exception:
|
||
pass
|
||
return jsonify({'cuda': False})
|
||
|
||
|
||
def _ms_to_srt_time(ms):
|
||
td = timedelta(milliseconds=ms)
|
||
h, rem = divmod(td.seconds, 3600)
|
||
m, s = divmod(rem, 60)
|
||
ms_part = td.microseconds // 1000
|
||
return f'{h:02d}:{m:02d}:{s:02d},{ms_part:03d}'
|
||
|
||
|
||
def _convert_to_wav(input_path):
|
||
"""用 FFmpeg 转为 16kHz 单声道 WAV，返回 wav 路径或 None"""
|
||
wav_path = os.path.join(tempfile.gettempdir(), uuid.uuid4().hex + '.wav')
|
||
cmd = ['ffmpeg', '-y', '-i', input_path, '-ar', '16000', '-ac', '1', wav_path]
|
||
try:
|
||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300,
|
||
creationflags=0x08000000 if os.name == 'nt' else 0)
|
||
if result.returncode == 0 and os.path.exists(wav_path):
|
||
return wav_path
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
def _get_model(model_name, device='cpu'):
|
||
"""获取或加载模型，CUDA 失败自动回退 CPU"""
|
||
cfg = _load_config()
|
||
model_dir = cfg['model_dir']
|
||
cache_key = f'{model_name}_{device}_{model_dir}'
|
||
if cache_key not in _model_cache:
|
||
from faster_whisper import WhisperModel
|
||
try:
|
||
_model_cache[cache_key] = WhisperModel(
|
||
model_name, device=device, download_root=model_dir
|
||
)
|
||
except Exception as e:
|
||
if device == 'cuda':
|
||
# CUDA 加载失败，回退 CPU
|
||
cpu_key = f'{model_name}_cpu_{model_dir}'
|
||
if cpu_key not in _model_cache:
|
||
_model_cache[cpu_key] = WhisperModel(
|
||
model_name, device='cpu', download_root=model_dir
|
||
)
|
||
return _model_cache[cpu_key], f'CUDA 加载失败({e})，已回退到 CPU'
|
||
raise
|
||
return _model_cache[cache_key], None
|
||
|
||
|
||
def _process_task(task):
|
||
"""处理单个转录任务"""
|
||
task_id = task['task_id']
|
||
fmt = task.get('format', 'text')
|
||
device = task.get('device', 'cpu')
|
||
try:
|
||
_task_progress[task_id] = {'percent': 0, 'status': 'loading', 'format': fmt}
|
||
|
||
model, warn = _get_model(task['model'], device)
|
||
if warn:
|
||
_task_progress[task_id] = {'percent': 0, 'status': 'transcribing', 'format': fmt, 'warning': warn}
|
||
else:
|
||
_task_progress[task_id] = {'percent': 0, 'status': 'transcribing', 'format': fmt}
|
||
|
||
lang = task['language'] if task['language'] != 'auto' else None
|
||
segments, info = model.transcribe(
|
||
task['wav_path'],
|
||
beam_size=5, best_of=5,
|
||
vad_filter=True,
|
||
language=lang,
|
||
)
|
||
|
||
total_duration = max(info.duration, 0.01)
|
||
results = []
|
||
|
||
for seg in segments:
|
||
_task_progress[task_id] = {
|
||
'percent': round(seg.end / total_duration, 2),
|
||
'status': 'transcribing',
|
||
'format': fmt,
|
||
}
|
||
text = seg.text.strip()
|
||
if not text or len(text) <= 1:
|
||
continue
|
||
if _cc_t2s:
|
||
text = _cc_t2s.convert(text)
|
||
|
||
start_ms = int(seg.start * 1000)
|
||
end_ms = int(seg.end * 1000)
|
||
|
||
results.append({
|
||
'start': start_ms,
|
||
'end': end_ms,
|
||
'start_time': _ms_to_srt_time(start_ms),
|
||
'end_time': _ms_to_srt_time(end_ms),
|
||
'text': text,
|
||
})
|
||
|
||
_task_results[task_id] = results
|
||
_task_progress[task_id] = {'percent': 1, 'status': 'done', 'format': fmt}
|
||
|
||
except Exception as e:
|
||
_task_progress[task_id] = {'percent': 0, 'status': 'error', 'error': str(e), 'format': fmt}
|
||
finally:
|
||
# 清理临时文件
|
||
try:
|
||
if os.path.exists(task['wav_path']):
|
||
os.remove(task['wav_path'])
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _worker():
|
||
"""后台 worker 线程"""
|
||
while True:
|
||
if not _task_queue:
|
||
time.sleep(1)
|
||
continue
|
||
task = _task_queue.pop(0)
|
||
_process_task(task)
|
||
|
||
|
||
def _ensure_worker():
|
||
global _worker_started
|
||
if not _worker_started:
|
||
with _worker_lock:
|
||
if not _worker_started:
|
||
t = threading.Thread(target=_worker, daemon=True)
|
||
t.start()
|
||
_worker_started = True
|
||
|
||
|
||
@bp.route('/transcribe', methods=['POST'])
|
||
def transcribe():
|
||
if 'audio' not in request.files:
|
||
return jsonify({'success': False, 'error': '请上传音频文件'}), 400
|
||
|
||
audio_file = request.files['audio']
|
||
if not audio_file.filename:
|
||
return jsonify({'success': False, 'error': '未选择文件'}), 400
|
||
|
||
model_name = request.form.get('model', 'base')
|
||
language = request.form.get('language', 'auto')
|
||
device = request.form.get('device', 'cpu')
|
||
output_format = request.form.get('format', 'text')
|
||
|
||
if model_name not in _scan_models(_load_config()['model_dir']):
|
||
return jsonify({'success': False, 'error': f'模型 {model_name} 不存在，请先下载到 stt_models 目录'}), 400
|
||
|
||
# 保存原始文件
|
||
ext = os.path.splitext(audio_file.filename)[1].lower()
|
||
original_path = os.path.join(tempfile.gettempdir(), uuid.uuid4().hex + ext)
|
||
audio_file.save(original_path)
|
||
|
||
# 转 WAV
|
||
wav_path = _convert_to_wav(original_path)
|
||
try:
|
||
os.remove(original_path)
|
||
except Exception:
|
||
pass
|
||
|
||
if not wav_path:
|
||
return jsonify({'success': False, 'error': '音频转换失败，请确保 FFmpeg 已安装'}), 500
|
||
|
||
# 创建任务
|
||
task_id = uuid.uuid4().hex
|
||
_task_progress[task_id] = {'percent': 0, 'status': 'queued', 'format': output_format}
|
||
_task_queue.append({
|
||
'task_id': task_id,
|
||
'wav_path': wav_path,
|
||
'model': model_name,
|
||
'language': language,
|
||
'device': device,
|
||
'format': output_format,
|
||
})
|
||
|
||
_ensure_worker()
|
||
return jsonify({'success': True, 'task_id': task_id})
|
||
|
||
|
||
@bp.route('/status/<task_id>')
|
||
def task_status(task_id):
|
||
progress = _task_progress.get(task_id)
|
||
if not progress:
|
||
return jsonify({'success': False, 'error': '任务不存在'}), 404
|
||
|
||
resp = {'success': True, 'progress': progress}
|
||
|
||
if progress['status'] == 'done':
|
||
results = _task_results.get(task_id, [])
|
||
fmt = progress.get('format', 'text')
|
||
|
||
if fmt == 'srt':
|
||
lines = []
|
||
for i, r in enumerate(results):
|
||
lines.append(f"{i+1}\n{r['start_time']} --> {r['end_time']}\n{r['text']}\n")
|
||
resp['result'] = '\n'.join(lines)
|
||
elif fmt == 'json':
|
||
resp['result'] = results
|
||
else:
|
||
resp['result'] = '\n'.join(r['text'] for r in results)
|
||
|
||
# 清理
|
||
_task_progress.pop(task_id, None)
|
||
_task_results.pop(task_id, None)
|
||
|
||
return jsonify(resp)
|