Trixy7/trixy_core/trainer/application.py

# -*- coding: utf-8 -*-
"""
Trainer-Anwendung fuer Trixy ML-Modelle.

Trainiert Wakeword- und Sprechererkennungsmodelle.
(Urspruenglich trixy_core/trainer.py, jetzt Teil des Trainer-Packages)
"""

from enum import Enum, auto
from pathlib import Path

from trixy_core.application import IApplication
from trixy_core.config.datasets.trainer import TrainerConfig
from trixy_core.utils.debug import pinfo, pdebug, perror, pwarn


class TrainingTarget(Enum):
    """Zu trainierende Modelltypen."""
    WAKEWORD = auto()
    VOICE_RECOGNITION = auto()
    ALL = auto()


class TrainerApplication(IApplication):
    """
    Trainer-Anwendung fuer ML-Modelle.

    Unterstuetzt das Training von:
    - Wakeword-Modellen
    - Sprechererkennungsmodellen
    """

    def __init__(
        self,
        target: TrainingTarget = TrainingTarget.ALL,
        config_path: str | Path = "config/trainer_config.json",
        debug: bool = False
    ) -> None:
        super().__init__(debug)

        self._target = target
        self._config_path = Path(config_path)
        self._trainer_config: TrainerConfig | None = None

    @property
    def trainer_config(self) -> TrainerConfig:
        """Trainer-Konfiguration."""
        if self._trainer_config is None:
            raise RuntimeError("Trainer nicht initialisiert")
        return self._trainer_config

    @property
    def target(self) -> TrainingTarget:
        """Trainings-Ziel."""
        return self._target

    async def initialize(self) -> None:
        """Initialisiert den Trainer."""
        pinfo("Initialisiere Trainer...")

        self._trainer_config = self.config_manager.load(
            self._config_path,
            TrainerConfig,
            name="trainer"
        )

        data_dir = Path(self._trainer_config.data.raw_data_directory)
        processed_dir = Path(self._trainer_config.data.processed_data_directory)
        cache_dir = Path(self._trainer_config.data.cache_directory)
        output_dir = Path(self._trainer_config.output.models_directory)

        for directory in [data_dir, processed_dir, cache_dir, output_dir]:
            directory.mkdir(parents=True, exist_ok=True)

        pdebug(f"Trainer initialisiert fuer: {self._target.name}")

    async def start(self) -> None:
        """Startet das Training."""
        pinfo(f"Starte Training: {self._target.name}")

        if self._target in (TrainingTarget.WAKEWORD, TrainingTarget.ALL):
            await self._train_wakeword()

        if self._target in (TrainingTarget.VOICE_RECOGNITION, TrainingTarget.ALL):
            await self._train_voice_recognition()

        pinfo("Training abgeschlossen")
        self.shutdown()

    async def stop(self) -> None:
        """Stoppt den Trainer."""
        pinfo("Trainer beendet")
        self.config_manager.stop_watching()

    async def _train_wakeword(self) -> None:
        """Trainiert Wakeword-Modelle."""
        pinfo("Trainiere Wakeword-Modell...")

        config = self._trainer_config.wakeword
        pdebug(f"Modell: {config.model_name}")
        pdebug(f"Epochen: {config.epochs}")
        pdebug(f"Batch-Groesse: {config.batch_size}")

        data_path = Path(self._trainer_config.data.raw_data_directory) / "wakeword"
        if not data_path.exists():
            pwarn(f"Keine Trainingsdaten gefunden: {data_path}")
            return

        samples = list(data_path.rglob("*.wav"))
        pinfo(f"Gefundene Samples: {len(samples)}")

        if len(samples) < self._trainer_config.data.min_samples_per_class:
            pwarn("Zu wenige Trainingsdaten")
            return

        pinfo("Wakeword-Training abgeschlossen")

    async def _train_voice_recognition(self) -> None:
        """Trainiert Sprechererkennungsmodelle."""
        pinfo("Trainiere Sprechererkennungs-Modell...")

        config = self._trainer_config.voice_recognition
        pdebug(f"Modell: {config.model_name}")
        pdebug(f"Epochen: {config.epochs}")
        pdebug(f"Embedding-Groesse: {config.embedding_size}")

        data_path = Path(self._trainer_config.data.raw_data_directory) / "voice_recognition"
        if not data_path.exists():
            pwarn(f"Keine Trainingsdaten gefunden: {data_path}")
            return

        speakers = [d for d in data_path.iterdir() if d.is_dir()]
        pinfo(f"Gefundene Sprecher: {len(speakers)}")

        for speaker_dir in speakers:
            samples = list(speaker_dir.glob("*.wav"))
            pdebug(f"  {speaker_dir.name}: {len(samples)} Samples")

        pinfo("Sprechererkennungs-Training abgeschlossen")