cosmo-voice-assistant/cosmo/wake_word.py

"""
Wake word detector для Cosmo.
Слушает микрофон непрерывно, детектирует слово "cosmo" через openwakeword.
"""

import os
import glob
import time
import threading
import queue
import numpy as np
import sounddevice as sd
from openwakeword.model import Model
from loguru import logger


class WakeWordDetector:
    def __init__(self, config: dict, on_detected_callback):
        """
        config: секция audio + assistant из config.yaml
        on_detected_callback: вызывается без аргументов при детекте wake word
        """
        self.sample_rate = config["audio"].get("sample_rate", 16000)
        self.chunk_duration = config["audio"].get("chunk_duration", 0.08)
        self.chunk_size = int(self.sample_rate * self.chunk_duration)
        self.wake_word = config["assistant"]["wake_word"]
        self.on_detected = on_detected_callback

        self._audio_queue = queue.Queue()
        self._running = False
        self._paused = False
        self._thread = None

        # Порог уверенности для срабатывания (0.0 – 1.0)
        # 0.7 — баланс между надёжностью и защитой от ложных срабатываний/эха TTS
        self.threshold = 0.7

        logger.info("Загружаю wake word модель openwakeword...")

        # Ищем кастомную модель в папке models/
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        custom_models = glob.glob(os.path.join(project_root, "models", "*.onnx"))

        if custom_models:
            model_path = custom_models[0]
            model_name = os.path.basename(model_path)
            try:
                self.model = Model(
                    wakeword_models=[model_path],
                    inference_framework="onnx",
                )
                logger.info(f"Кастомная wake word модель загружена: {model_name}")
            except Exception as e:
                logger.error(f"Не удалось загрузить кастомную модель {model_name}: {e}")
                raise
        else:
            # Fallback на встроенную hey_jarvis
            try:
                self.model = Model(
                    wakeword_models=["hey_jarvis"],
                    inference_framework="onnx",
                )
                logger.info("Wake word модель 'hey_jarvis' загружена (onnx)")
                logger.warning(
                    "Кастомная модель не найдена в папке models/. "
                    "Активация по слову 'Hey Jarvis'. "
                    "Запусти train_wakeword/train.sh чтобы обучить модель на 'Hey Cosmo'."
                )
            except Exception as e:
                logger.error(f"Не удалось загрузить wake word модель: {e}")
                raise

    # ------------------------------------------------------------------
    # Публичный интерфейс
    # ------------------------------------------------------------------

    def start(self):
        """Запустить детектор в фоновом потоке."""
        self._running = True
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        logger.info(f"Wake word детектор запущен. Жду слово '{self.wake_word}'...")

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join(timeout=2)

    def pause(self):
        """Приостановить детект (пока идёт запись команды)."""
        self._paused = True

    def resume(self, cooldown: float = 1.5):
        """Возобновить детект после записи команды с защитой от эха."""
        # Ждём пока эхо от TTS затухнет
        time.sleep(cooldown)
        # Очищаем очередь — там буферизованный звук TTS
        while not self._audio_queue.empty():
            try:
                self._audio_queue.get_nowait()
            except queue.Empty:
                break
        # Сбрасываем внутреннее состояние модели (накопленные скоры)
        self.model.reset()
        self._paused = False
        logger.debug("Wake word детектор возобновлён")

    # ------------------------------------------------------------------
    # Внутренняя логика
    # ------------------------------------------------------------------

    def _audio_callback(self, indata, frames, time_info, status):
        if status:
            logger.debug(f"sounddevice статус: {status}")
        if not self._paused:
            # Копируем данные, т.к. буфер перезаписывается
            self._audio_queue.put(indata[:, 0].copy())

    def _run(self):
        with sd.InputStream(
            samplerate=self.sample_rate,
            channels=1,
            dtype="float32",
            blocksize=self.chunk_size,
            callback=self._audio_callback,
        ):
            while self._running:
                try:
                    chunk = self._audio_queue.get(timeout=0.5)
                except queue.Empty:
                    continue

                # openwakeword ожидает int16 PCM
                chunk_int16 = (chunk * 32767).astype(np.int16)
                prediction = self.model.predict(chunk_int16)

                for model_name, score in prediction.items():
                    if score >= self.threshold:
                        logger.info(
                            f"Wake word задетектирован! Модель={model_name}, "
                            f"уверенность={score:.2f}"
                        )
                        self.pause()
                        self.on_detected()
                        break