Trixy7/trixy_core/wakeword/detector.py

# -*- coding: utf-8 -*-
"""
Wakeword Detector - Wrapper fuer OpenWakeWord.

Laedt native OWW-Modelle (.onnx) — sowohl vortrainierte als auch
custom-trainierte Modelle aus models/wakeword/{name}/{name}.onnx.
"""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Callable

import numpy as np

try:
    from openwakeword.model import Model as OWWModel
    OPENWAKEWORD_AVAILABLE = True
except ImportError:
    OPENWAKEWORD_AVAILABLE = False
    OWWModel = None


class WakewordType(Enum):
    """Typ des erkannten Wakewords."""
    CUSTOM = "custom"
    SYSTEM_COMMAND = "system_command"
    UNKNOWN = "unknown"


@dataclass
class DetectorConfig:
    """Konfiguration fuer den Wakeword-Detector."""

    # Modell-Konfiguration
    model_directory: str = "models/wakeword"
    models: list[str] = field(default_factory=lambda: ["custom", "system_command"])
    use_onnx: bool = True

    # Erkennungs-Schwellenwerte
    threshold: float = 0.5
    threshold_per_model: dict[str, float] = field(default_factory=dict)

    # Audio-Konfiguration
    sample_rate: int = 16000
    frame_length_ms: int = 80  # OpenWakeWord empfiehlt 80ms Frames

    # Verhalten
    refractory_period_ms: int = 500  # Mindestzeit zwischen Erkennungen

    @property
    def frame_length_samples(self) -> int:
        """Anzahl Samples pro Frame."""
        return int(self.sample_rate * self.frame_length_ms / 1000)

    def get_threshold(self, model_name: str) -> float:
        """Gibt Threshold fuer spezifisches Modell zurueck."""
        return self.threshold_per_model.get(model_name, self.threshold)


@dataclass
class WakewordDetection:
    """Ergebnis einer Wakeword-Erkennung."""

    wakeword_type: WakewordType
    model_name: str
    confidence: float
    timestamp: datetime
    audio_level: float  # RMS-Level fuer Arbitration
    audio_chunks: list[bytes] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Konvertiert zu Dictionary."""
        return {
            "wakeword_type": self.wakeword_type.value,
            "model_name": self.model_name,
            "confidence": self.confidence,
            "timestamp": self.timestamp.isoformat(),
            "audio_level": self.audio_level,
            "audio_chunks_count": len(self.audio_chunks),
        }


class WakewordDetector:
    """
    Wrapper fuer OpenWakeWord.

    Verarbeitet Audio-Frames und erkennt Wakewords.
    Laedt ONNX-Modelle aus dem konfigurierten Verzeichnis.
    """

    def __init__(self, config: DetectorConfig | None = None):
        self._config = config or DetectorConfig()
        self._model: OWWModel | None = None
        self._loaded_models: dict[str, str] = {}  # model_name -> file_path
        self._last_detection_time: datetime | None = None

        # Audio-Ring-Buffer fuer Wakeword-Chunks
        self._chunk_buffer: list[bytes] = []
        self._chunk_buffer_max = 10  # Letzte 10 Chunks (~800ms bei 80ms Frames)

        # Callbacks
        self._on_detection_callbacks: list[Callable[[WakewordDetection], None]] = []

        # Statistiken
        self._total_frames_processed = 0
        self._total_detections = 0

    @property
    def config(self) -> DetectorConfig:
        """Gibt Konfiguration zurueck."""
        return self._config

    @property
    def is_loaded(self) -> bool:
        """Prueft ob Modelle geladen sind."""
        return self._model is not None

    @property
    def loaded_models(self) -> list[str]:
        """Liste geladener Modelle."""
        return list(self._loaded_models.keys())

    @property
    def stats(self) -> dict:
        """Statistiken."""
        return {
            "frames_processed": self._total_frames_processed,
            "detections": self._total_detections,
            "is_loaded": self.is_loaded,
            "loaded_models": self.loaded_models,
        }

    def load_models(self) -> bool:
        """
        Laedt die konfigurierten Wakeword-Modelle.

        Sucht nach ONNX-Dateien in models/wakeword/{name}/.
        """
        if not OPENWAKEWORD_AVAILABLE:
            raise RuntimeError(
                "OpenWakeWord nicht installiert. "
                "Installiere mit: pip install openwakeword"
            )

        model_dir = Path(self._config.model_directory)
        model_paths = []

        for model_name in self._config.models:
            model_path = self._find_model_file(model_dir / model_name)
            if model_path:
                model_paths.append(str(model_path))
                self._loaded_models[model_name] = str(model_path)
            else:
                # Pre-trained OWW-Modell (wird von OWW selbst geladen)
                self._loaded_models[model_name] = model_name

        try:
            if model_paths:
                self._model = OWWModel(wakeword_model_paths=model_paths)
            else:
                self._model = OWWModel()
            return True

        except Exception as e:
            raise RuntimeError(f"Fehler beim Laden der Modelle: {e}")

    def _find_model_file(self, model_dir: Path) -> Path | None:
        """Sucht nach Modell-Datei (.onnx/.tflite/.pt) im Verzeichnis."""
        if not model_dir.exists():
            return None

        extensions = [".onnx", ".tflite", ".pt"]

        for ext in extensions:
            model_file = model_dir / f"model{ext}"
            if model_file.exists():
                return model_file

            files = list(model_dir.glob(f"*{ext}"))
            if files:
                return files[0]

        return None

    def unload_models(self) -> None:
        """Entlaedt alle Modelle."""
        self._model = None
        self._loaded_models.clear()

    def process_frame(self, audio_frame: bytes | np.ndarray) -> WakewordDetection | None:
        """
        Verarbeitet einen Audio-Frame.

        Args:
            audio_frame: Audio-Daten (16-bit PCM oder numpy array)

        Returns:
            WakewordDetection wenn erkannt, sonst None
        """
        if not self.is_loaded:
            raise RuntimeError("Modelle nicht geladen. Rufe load_models() auf.")

        # Konvertiere zu numpy wenn noetig
        if isinstance(audio_frame, bytes):
            audio_array = np.frombuffer(audio_frame, dtype=np.int16)
            self._add_to_chunk_buffer(audio_frame)
        else:
            audio_array = audio_frame
            self._add_to_chunk_buffer(audio_array.astype(np.int16).tobytes())

        # Audio-Level berechnen (RMS)
        audio_level = self._calculate_rms(audio_array)

        # OpenWakeWord Prediction
        predictions = self._model.predict(audio_array)

        self._total_frames_processed += 1

        # Pruefe auf Erkennungen
        detection = self._check_predictions(predictions, audio_level)

        if detection:
            self._total_detections += 1
            self._last_detection_time = detection.timestamp

            for callback in self._on_detection_callbacks:
                try:
                    callback(detection)
                except Exception:
                    pass

        return detection

    def _add_to_chunk_buffer(self, audio_bytes: bytes) -> None:
        """Fuegt Audio-Chunk zum Ring-Buffer hinzu."""
        self._chunk_buffer.append(audio_bytes)
        while len(self._chunk_buffer) > self._chunk_buffer_max:
            self._chunk_buffer.pop(0)

    def _get_wakeword_chunks(self) -> list[bytes]:
        """Gibt Audio-Chunks zurueck, die das Wakeword enthalten."""
        return list(self._chunk_buffer)

    def _calculate_rms(self, audio: np.ndarray) -> float:
        """Berechnet RMS-Level."""
        if len(audio) == 0:
            return 0.0
        audio_float = audio.astype(np.float32) / 32768.0
        return float(np.sqrt(np.mean(audio_float ** 2)))

    def _check_predictions(
        self,
        predictions: dict,
        audio_level: float,
    ) -> WakewordDetection | None:
        """Prueft Predictions auf Erkennungen."""
        now = datetime.now()

        # Refractory Period pruefen
        if self._last_detection_time:
            elapsed_ms = (now - self._last_detection_time).total_seconds() * 1000
            if elapsed_ms < self._config.refractory_period_ms:
                return None

        # Hoechste Confidence ueber Threshold suchen
        best_model = None
        best_confidence = 0.0

        for model_name, confidence in predictions.items():
            threshold = self._config.get_threshold(model_name)
            if confidence >= threshold and confidence > best_confidence:
                best_model = model_name
                best_confidence = confidence

        if best_model is None:
            return None

        wakeword_type = self._determine_wakeword_type(best_model)

        return WakewordDetection(
            wakeword_type=wakeword_type,
            model_name=best_model,
            confidence=best_confidence,
            timestamp=now,
            audio_level=audio_level,
            audio_chunks=self._get_wakeword_chunks(),
        )

    def _determine_wakeword_type(self, model_name: str) -> WakewordType:
        """Bestimmt Wakeword-Typ basierend auf Modellname."""
        model_lower = model_name.lower()

        if "custom" in model_lower or model_lower in ["hey_jarvis", "alexa", "hey_mycroft"]:
            return WakewordType.CUSTOM
        elif "system" in model_lower or "command" in model_lower:
            return WakewordType.SYSTEM_COMMAND
        else:
            return WakewordType.CUSTOM

    def on_detection(self, callback: Callable[[WakewordDetection], None]) -> None:
        """Registriert Callback fuer Wakeword-Erkennungen."""
        self._on_detection_callbacks.append(callback)

    def remove_callback(self, callback: Callable[[WakewordDetection], None]) -> bool:
        """Entfernt Callback."""
        try:
            self._on_detection_callbacks.remove(callback)
            return True
        except ValueError:
            return False

    def reset(self) -> None:
        """Setzt Detector-State zurueck."""
        self._chunk_buffer.clear()
        self._last_detection_time = None

        if self._model and hasattr(self._model, "reset"):
            self._model.reset()

    def reset_stats(self) -> None:
        """Setzt Statistiken zurueck."""
        self._total_frames_processed = 0
        self._total_detections = 0