Trixy3/trixy_core/ml/wakeword/audio_features.py

"""
Audio feature extraction for wakeword detection.

This module provides efficient feature extraction from raw audio data,
specifically optimized for wakeword detection using Log-Mel spectrograms.
"""

import math
import warnings
from dataclasses import dataclass
from typing import Tuple, Optional, Union

import torch
import torch.nn.functional as F
import torchaudio
import torchaudio.functional as AF
import torchaudio.transforms as AT
import numpy as np


@dataclass 
class SpectrogramConfig:
    """Configuration for spectrogram generation."""
    
    # Audio parameters
    sample_rate: int = 16000
    n_fft: int = 512
    hop_length: int = 160  # 10ms at 16kHz
    win_length: int = 400  # 25ms at 16kHz
    window_fn: str = "hann"
    
    # Mel filterbank parameters  
    n_mels: int = 40
    f_min: float = 80.0
    f_max: Optional[float] = None  # Will default to sample_rate/2
    
    # Normalization and preprocessing
    power: float = 2.0
    normalized: bool = True
    center: bool = True
    pad_mode: str = "constant"
    
    # Log transformation
    log_offset: float = 1e-8
    
    # Output parameters
    time_frames: int = 151  # Target number of time frames
    
    def __post_init__(self):
        """Validate and set derived parameters."""
        if self.f_max is None:
            self.f_max = self.sample_rate / 2.0
            
        # Calculate expected number of frames for validation
        # This helps ensure consistent input sizes to the model
        self.expected_frames = self.time_frames


class AudioFeatureExtractor:
    """
    Efficient audio feature extraction for wakeword detection.
    
    This class provides optimized Log-Mel spectrogram extraction with
    proper normalization and windowing for real-time inference.
    """
    
    def __init__(self, config: SpectrogramConfig, device: Optional[torch.device] = None):
        """
        Initialize the feature extractor.
        
        Args:
            config: Spectrogram configuration
            device: PyTorch device for computation
        """
        self.config = config
        self.device = device or torch.device("cpu")
        
        # Pre-compute mel filterbank for efficiency
        self.mel_scale = AT.MelScale(
            n_mels=config.n_mels,
            sample_rate=config.sample_rate,
            f_min=config.f_min,
            f_max=config.f_max,
            n_stft=config.n_fft // 2 + 1,
            norm="slaney",
            mel_scale="htk"
        ).to(self.device)
        
        # Pre-compute window function
        if config.window_fn == "hann":
            self.window = torch.hann_window(config.win_length).to(self.device)
        elif config.window_fn == "hamming":
            self.window = torch.hamming_window(config.win_length).to(self.device)
        elif config.window_fn == "blackman":
            self.window = torch.blackman_window(config.win_length).to(self.device)
        else:
            raise ValueError(f"Unsupported window function: {config.window_fn}")
        
        # Cache for efficiency metrics
        self._cache_stats = {
            "total_extractions": 0,
            "total_time_ms": 0.0,
            "avg_time_ms": 0.0
        }
    
    def extract_features(self, waveform: torch.Tensor, 
                        normalize_audio: bool = True) -> torch.Tensor:
        """
        Extract Log-Mel spectrogram features from audio waveform.
        
        Args:
            waveform: Input audio tensor of shape (batch_size, samples) or (samples,)
            normalize_audio: Whether to normalize audio amplitude
            
        Returns:
            Log-Mel spectrogram tensor of shape (batch_size, n_mels, time_frames)
        """
        start_time = torch.cuda.Event(enable_timing=True) if self.device.type == "cuda" else None
        
        # Ensure proper tensor format
        if waveform.dim() == 1:
            waveform = waveform.unsqueeze(0)  # Add batch dimension
        
        waveform = waveform.to(self.device)
        
        # Audio preprocessing
        if normalize_audio:
            waveform = self._normalize_audio(waveform)
        
        # Apply pre-emphasis filter (optional but common for speech)
        waveform = self._apply_preemphasis(waveform, alpha=0.97)
        
        # Compute STFT
        stft = torch.stft(
            waveform.squeeze(0) if waveform.shape[0] == 1 else waveform,
            n_fft=self.config.n_fft,
            hop_length=self.config.hop_length,
            win_length=self.config.win_length,
            window=self.window,
            center=self.config.center,
            pad_mode=self.config.pad_mode,
            normalized=self.config.normalized,
            onesided=True,
            return_complex=True
        )
        
        # Compute power spectrogram
        power_spec = stft.abs().pow(self.config.power)
        
        # Apply mel filterbank
        mel_spec = self.mel_scale(power_spec)
        
        # Apply log transformation
        log_mel_spec = torch.log(mel_spec + self.config.log_offset)
        
        # Ensure correct shape and size
        log_mel_spec = self._ensure_correct_frames(log_mel_spec)
        
        # Add batch dimension if needed
        if log_mel_spec.dim() == 2:
            log_mel_spec = log_mel_spec.unsqueeze(0)
        
        # Update performance stats
        if start_time and self.device.type == "cuda":
            end_time = torch.cuda.Event(enable_timing=True)
            end_time.record()
            torch.cuda.synchronize()
            elapsed_ms = start_time.elapsed_time(end_time)
            self._update_stats(elapsed_ms)
        
        return log_mel_spec
    
    def extract_features_streaming(self, audio_chunk: torch.Tensor,
                                 previous_context: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Extract features for streaming audio with context management.
        
        Args:
            audio_chunk: Current audio chunk
            previous_context: Previous audio context for overlap processing
            
        Returns:
            Tuple of (features, new_context) for next iteration
        """
        # Combine with previous context if available
        if previous_context is not None:
            combined_audio = torch.cat([previous_context, audio_chunk], dim=-1)
        else:
            combined_audio = audio_chunk
        
        # Extract features
        features = self.extract_features(combined_audio, normalize_audio=True)
        
        # Calculate context for next iteration (50% overlap)
        context_samples = len(audio_chunk) // 2
        new_context = combined_audio[..., -context_samples:] if len(combined_audio) > context_samples else combined_audio
        
        return features, new_context
    
    def _normalize_audio(self, waveform: torch.Tensor) -> torch.Tensor:
        """Normalize audio waveform."""
        # RMS normalization
        rms = torch.sqrt(torch.mean(waveform ** 2, dim=-1, keepdim=True))
        # Avoid division by zero
        rms = torch.clamp(rms, min=1e-8)
        normalized = waveform / rms
        
        # Scale to reasonable range
        normalized = normalized * 0.1
        
        return normalized
    
    def _apply_preemphasis(self, waveform: torch.Tensor, alpha: float = 0.97) -> torch.Tensor:
        """Apply pre-emphasis filter to enhance high frequencies."""
        return torch.cat([
            waveform[..., :1],
            waveform[..., 1:] - alpha * waveform[..., :-1]
        ], dim=-1)
    
    def _ensure_correct_frames(self, mel_spec: torch.Tensor) -> torch.Tensor:
        """Ensure spectrogram has the correct number of time frames."""
        current_frames = mel_spec.shape[-1]
        target_frames = self.config.time_frames
        
        if current_frames == target_frames:
            return mel_spec
        elif current_frames > target_frames:
            # Truncate from center
            start_idx = (current_frames - target_frames) // 2
            return mel_spec[..., start_idx:start_idx + target_frames]
        else:
            # Pad to reach target frames
            pad_total = target_frames - current_frames
            pad_left = pad_total // 2
            pad_right = pad_total - pad_left
            return F.pad(mel_spec, (pad_left, pad_right), mode='constant', value=0)
    
    def _update_stats(self, elapsed_ms: float):
        """Update performance statistics."""
        self._cache_stats["total_extractions"] += 1
        self._cache_stats["total_time_ms"] += elapsed_ms
        self._cache_stats["avg_time_ms"] = (
            self._cache_stats["total_time_ms"] / self._cache_stats["total_extractions"]
        )
    
    def get_performance_stats(self) -> dict:
        """Get performance statistics."""
        return self._cache_stats.copy()
    
    def reset_stats(self):
        """Reset performance statistics."""
        self._cache_stats = {
            "total_extractions": 0,
            "total_time_ms": 0.0,
            "avg_time_ms": 0.0
        }
    
    def calculate_audio_duration(self, n_samples: int) -> float:
        """Calculate audio duration in seconds from number of samples."""
        return n_samples / self.config.sample_rate
    
    def calculate_required_samples(self, duration_seconds: float) -> int:
        """Calculate number of samples needed for given duration."""
        return int(duration_seconds * self.config.sample_rate)
    
    def validate_audio_format(self, waveform: torch.Tensor) -> bool:
        """
        Validate that audio format matches expected configuration.
        
        Args:
            waveform: Input audio tensor
            
        Returns:
            True if format is valid, False otherwise
        """
        if waveform.dim() not in [1, 2]:
            return False
            
        if waveform.dim() == 2 and waveform.shape[0] > 1:
            # Multi-channel audio - should be mono for wakeword detection
            warnings.warn("Multi-channel audio detected. Consider converting to mono.")
            return False
            
        return True


class BatchAudioFeatureExtractor(AudioFeatureExtractor):
    """
    Batch-optimized feature extractor for processing multiple audio samples.
    
    Useful for training and batch inference scenarios.
    """
    
    def __init__(self, config: SpectrogramConfig, device: Optional[torch.device] = None):
        super().__init__(config, device)
        self.batch_size_limit = 32  # Prevent OOM for large batches
    
    def extract_features_batch(self, waveforms: torch.Tensor,
                             lengths: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Extract features from a batch of waveforms.
        
        Args:
            waveforms: Batch of waveforms (batch_size, max_length)
            lengths: Actual lengths of each waveform (optional)
            
        Returns:
            Batch of Log-Mel spectrograms (batch_size, n_mels, time_frames)
        """
        batch_size = waveforms.shape[0]
        
        if batch_size > self.batch_size_limit:
            # Process in chunks to avoid OOM
            all_features = []
            for i in range(0, batch_size, self.batch_size_limit):
                chunk = waveforms[i:i + self.batch_size_limit]
                chunk_lengths = lengths[i:i + self.batch_size_limit] if lengths is not None else None
                chunk_features = self._process_batch_chunk(chunk, chunk_lengths)
                all_features.append(chunk_features)
            
            return torch.cat(all_features, dim=0)
        else:
            return self._process_batch_chunk(waveforms, lengths)
    
    def _process_batch_chunk(self, waveforms: torch.Tensor,
                           lengths: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Process a chunk of the batch."""
        features_list = []
        
        for i in range(waveforms.shape[0]):
            waveform = waveforms[i]
            
            # Trim to actual length if provided
            if lengths is not None:
                actual_length = int(lengths[i].item())
                waveform = waveform[:actual_length]
            
            # Extract features for this sample
            features = self.extract_features(waveform.unsqueeze(0), normalize_audio=True)
            features_list.append(features.squeeze(0))
        
        return torch.stack(features_list, dim=0)


def create_feature_extractor(config: Optional[SpectrogramConfig] = None,
                           device: Optional[torch.device] = None,
                           batch_mode: bool = False) -> AudioFeatureExtractor:
    """
    Factory function to create feature extractors.
    
    Args:
        config: Spectrogram configuration (uses default if None)
        device: PyTorch device
        batch_mode: Whether to create batch-optimized extractor
        
    Returns:
        Feature extractor instance
    """
    if config is None:
        config = SpectrogramConfig()
    
    if batch_mode:
        return BatchAudioFeatureExtractor(config, device)
    else:
        return AudioFeatureExtractor(config, device)


def validate_spectrogram_config(config: SpectrogramConfig) -> bool:
    """
    Validate spectrogram configuration parameters.
    
    Args:
        config: Configuration to validate
        
    Returns:
        True if valid, False otherwise
    """
    try:
        # Check basic constraints
        if config.sample_rate <= 0:
            return False
        
        if config.n_fft <= 0 or config.n_fft > config.sample_rate:
            return False
            
        if config.hop_length <= 0 or config.hop_length > config.n_fft:
            return False
            
        if config.win_length <= 0 or config.win_length > config.n_fft:
            return False
            
        if config.n_mels <= 0:
            return False
            
        if config.f_min < 0:
            return False
            
        if config.f_max is not None and config.f_max <= config.f_min:
            return False
            
        if config.time_frames <= 0:
            return False
        
        return True
        
    except Exception:
        return False