103 lines
4.9 KiB
Python
103 lines
4.9 KiB
Python
# vosk_stt.py
|
|
import json
|
|
import yaml
|
|
import logging
|
|
from datetime import datetime
|
|
from vosk import Model, KaldiRecognizer
|
|
|
|
class VoskSTT:
|
|
def __init__(self, model_path="/home/sirmrak/models/vosk-model-small-ru-0.22", commands_file="/mnt/Project/my_speaker/commands/commands.yaml", sample_rate=16000):
|
|
"""
|
|
Инициализирует Vosk STT с моделью и словарем команд из YAML файла.
|
|
|
|
:param model_path: Путь к модели Vosk (vosk-model-small-ru-0.22)
|
|
:param commands_file: Путь к файлу commands.yaml
|
|
:param sample_rate: Частота дискретизации аудио (по умолчанию 16000)
|
|
"""
|
|
self.model_path = model_path
|
|
self.commands_file = commands_file
|
|
self.sample_rate = sample_rate
|
|
self.model = None
|
|
self.recognizer = None
|
|
self._load_model()
|
|
self._load_commands()
|
|
|
|
def _load_model(self):
|
|
"""Загружает модель Vosk."""
|
|
logger = logging.getLogger('VADApp')
|
|
try:
|
|
self.model = Model(self.model_path)
|
|
logger.info(f"Модель Vosk загружена: {self.model_path}")
|
|
except Exception as e:
|
|
logger.error(f"Ошибка загрузки модели: {e}")
|
|
raise
|
|
|
|
def _load_commands(self):
|
|
"""Формирует словарь (список фраз) из commands.yaml и инициализирует recognizer."""
|
|
logger = logging.getLogger('VADApp')
|
|
try:
|
|
with open(self.commands_file, 'r', encoding='utf-8') as f:
|
|
commands = yaml.safe_load(f)
|
|
|
|
phrases = []
|
|
|
|
# Собираем все фразы из всех разделов
|
|
for section, content in commands.items():
|
|
if section in ['Wakewords', 'Session']:
|
|
# Обрабатываем списки с вложенными phrases
|
|
for item in content:
|
|
if isinstance(item, dict) and 'phrases' in item:
|
|
phrases.extend(item['phrases'])
|
|
elif section == 'Volume':
|
|
# Обрабатываем новую структуру Volume с приоритетом
|
|
for item in content:
|
|
if isinstance(item, dict) and 'commands' in item:
|
|
phrases.extend(list(item['commands'].keys()))
|
|
elif section == 'Text':
|
|
# Обрабатываем новую структуру Text с приоритетом
|
|
for item in content:
|
|
if isinstance(item, dict) and 'commands' in item:
|
|
phrases.extend(list(item['commands'].keys()))
|
|
|
|
# Удаляем дубликаты и сортируем
|
|
phrases = sorted(set(phrases))
|
|
|
|
# Инициализируем recognizer с грамматикой (списком фраз)
|
|
grammar = json.dumps(phrases, ensure_ascii=False)
|
|
self.recognizer = KaldiRecognizer(self.model, self.sample_rate, grammar)
|
|
|
|
logger.info(f"Словарь команд загружен: {len(phrases)} фраз")
|
|
logger.debug(f"Фразы: {phrases}")
|
|
except Exception as e:
|
|
logger.error(f"Ошибка загрузки команд: {e}")
|
|
raise
|
|
|
|
def recognize(self, audio_data):
|
|
"""
|
|
Распознает речь из аудио данных (bytes).
|
|
|
|
:param audio_data: Bytes аудио (16-bit PCM, mono)
|
|
:return: Распознанный текст или None если ничего не распознано
|
|
"""
|
|
logger = logging.getLogger('VADApp')
|
|
try:
|
|
if not self.recognizer.AcceptWaveform(audio_data):
|
|
# Частичный результат (если нужно, но для коротких команд используем Final)
|
|
partial = json.loads(self.recognizer.PartialResult()).get("partial", "")
|
|
if partial:
|
|
logger.debug(f"Частичный результат: {partial}")
|
|
result = json.loads(self.recognizer.FinalResult()).get("text", "")
|
|
if result:
|
|
logger.info(f"Распознанная команда: {result}")
|
|
return result
|
|
else:
|
|
logger.info("Ничего не распознано")
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"Ошибка распознавания: {e}")
|
|
return None
|
|
|
|
def reload_commands(self):
|
|
"""Перезагружает словарь команд из файла."""
|
|
self._load_commands()
|