Trixy3/trainer/voice_recognition/data.py

"""
Data handling and preprocessing for voice recognition training.

This module provides specialized data handling for speaker recognition tasks,
including speaker enrollment, data augmentation, and balanced sampling.
"""

import os
import random
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import logging

import torch
from torch.utils.data import DataLoader

from ..data_pipeline import (
    DatasetBuilder, DataPreprocessor, AudioDataset, 
    AudioProcessor, AudioAugmentation, DataLoaderFactory
)


class VoiceRecognitionDataPreprocessor(DataPreprocessor):
    """Specialized data preprocessor for voice recognition tasks."""
    
    def __init__(self, audio_config, logger: Optional[logging.Logger] = None):
        """Initialize voice recognition data preprocessor."""
        super().__init__(audio_config, logger)
    
    def preprocess_voice_recognition_data(self, raw_data_dir: str, output_dir: str,
                                        target_chunk_length: float = 3.0) -> str:
        """
        Preprocess raw voice recognition data with speaker enrollment focus.
        
        Args:
            raw_data_dir: Directory containing raw speaker audio files
            output_dir: Output directory for processed data
            target_chunk_length: Target length for audio chunks (seconds)
            
        Returns:
            Path to processed data directory
        """
        return super().preprocess_voice_recognition_data(
            raw_data_dir, output_dir, target_chunk_length
        )


def create_voice_recognition_dataloaders(data_dir: str,
                                       speaker_mapping: Dict[str, str],
                                       audio_config,
                                       batch_size: int = 32,
                                       num_workers: int = 4,
                                       augmentation_config: Optional[Dict[str, float]] = None,
                                       balance_speakers: bool = True) -> Tuple[DataLoader, DataLoader, DataLoader, Dict[str, str]]:
    """
    Create data loaders for voice recognition training.
    
    Args:
        data_dir: Directory containing processed voice recognition data
        speaker_mapping: Mapping of directory names to speaker names
        audio_config: Audio processing configuration
        batch_size: Batch size for training
        num_workers: Number of worker processes
        augmentation_config: Configuration for data augmentation
        balance_speakers: Whether to balance speakers in training batches
        
    Returns:
        Tuple of (train_loader, val_loader, test_loader, final_speaker_mapping)
    """
    logger = logging.getLogger(__name__)
    
    # Initialize components
    audio_processor = AudioProcessor(audio_config)
    dataset_builder = DatasetBuilder(logger)
    
    # Create datasets
    try:
        train_dataset, val_dataset, test_dataset, final_speaker_mapping = dataset_builder.create_voice_recognition_dataset(
            data_dir, audio_processor, speaker_mapping,
            augmentation_config=augmentation_config
        )
    except Exception as e:
        logger.error(f"Failed to create voice recognition datasets: {str(e)}")
        raise
    
    # Create data loaders with appropriate sampling
    train_loader = DataLoaderFactory.create_dataloader(
        train_dataset, batch_size, shuffle=True, num_workers=num_workers,
        balanced_sampling=balance_speakers
    )
    
    val_loader = DataLoaderFactory.create_dataloader(
        val_dataset, batch_size, shuffle=False, num_workers=num_workers
    )
    
    test_loader = DataLoaderFactory.create_dataloader(
        test_dataset, batch_size, shuffle=False, num_workers=num_workers
    )
    
    return train_loader, val_loader, test_loader, final_speaker_mapping