Trixy3/trixy_core/ml/voice_recognition/vad.py

"""
Voice Activity Detection (VAD) for Voice Recognition system.

This module provides comprehensive VAD capabilities including WebRTC VAD,
energy-based detection, and hybrid approaches optimized for speaker recognition.
"""

import logging
import warnings
from typing import Optional, List, Tuple, Union, Dict, Any
from dataclasses import dataclass
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from abc import ABC, abstractmethod

# Configure logging
logger = logging.getLogger(__name__)

# Try to import WebRTC VAD
try:
    import webrtcvad
    WEBRTC_AVAILABLE = True
except ImportError:
    WEBRTC_AVAILABLE = False
    logger.warning("WebRTC VAD not available. Install webrtcvad package for WebRTC support.")


@dataclass
class VADResult:
    """Result of voice activity detection."""
    is_speech: bool
    confidence: float
    start_time: float
    end_time: float
    frame_results: Optional[List[bool]] = None
    energy_levels: Optional[List[float]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            'is_speech': self.is_speech,
            'confidence': self.confidence,
            'start_time': self.start_time,
            'end_time': self.end_time,
            'frame_results': self.frame_results,
            'energy_levels': self.energy_levels
        }


@dataclass
class VADConfig:
    """Configuration for Voice Activity Detection."""
    # General settings
    sample_rate: int = 16000
    frame_duration_ms: int = 30  # Frame duration in milliseconds
    
    # WebRTC VAD settings
    webrtc_aggressiveness: int = 2  # 0-3, higher = more aggressive
    
    # Energy-based VAD settings
    energy_threshold: float = 0.01
    energy_min_length_ms: int = 100  # Minimum speech length
    energy_hangover_ms: int = 50     # Hangover time after speech
    energy_prethreshold: float = 0.005  # Pre-trigger threshold
    
    # Spectral VAD settings
    spectral_threshold: float = 0.5
    spectral_rolloff_threshold: float = 0.85
    spectral_centroid_threshold: float = 2000.0
    
    # Temporal smoothing
    use_temporal_smoothing: bool = True
    smoothing_window_size: int = 5
    
    # Combination settings (for hybrid VAD)
    combination_mode: str = "weighted"  # "majority", "weighted", "conservative", "aggressive"
    webrtc_weight: float = 0.4
    energy_weight: float = 0.3
    spectral_weight: float = 0.3
    
    def __post_init__(self):
        """Validate configuration parameters."""
        if self.frame_duration_ms not in [10, 20, 30]:
            raise ValueError("WebRTC VAD frame duration must be 10, 20, or 30 ms")
        
        if not 0 <= self.webrtc_aggressiveness <= 3:
            raise ValueError("WebRTC aggressiveness must be 0-3")
        
        # Validate combination weights
        total_weight = self.webrtc_weight + self.energy_weight + self.spectral_weight
        if abs(total_weight - 1.0) > 1e-6:
            logger.warning(f"Combination weights sum to {total_weight}, not 1.0")


class VoiceActivityDetector(ABC):
    """Abstract base class for Voice Activity Detectors."""
    
    def __init__(self, config: VADConfig):
        self.config = config
        self.sample_rate = config.sample_rate
        self.frame_size = int(config.frame_duration_ms * config.sample_rate / 1000)
    
    @abstractmethod
    def detect(self, audio: Union[np.ndarray, torch.Tensor]) -> VADResult:
        """
        Detect voice activity in audio.
        
        Args:
            audio: Audio signal
            
        Returns:
            VAD result
        """
        pass
    
    def _ensure_numpy(self, audio: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
        """Convert audio to numpy array if needed."""
        if isinstance(audio, torch.Tensor):
            return audio.detach().cpu().numpy()
        return audio
    
    def _ensure_torch(self, audio: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
        """Convert audio to torch tensor if needed."""
        if isinstance(audio, np.ndarray):
            return torch.from_numpy(audio).float()
        return audio.float()


class WebRTCVAD(VoiceActivityDetector):
    """WebRTC-based Voice Activity Detector."""
    
    def __init__(self, config: VADConfig):
        super().__init__(config)
        
        if not WEBRTC_AVAILABLE:
            raise ImportError("WebRTC VAD not available. Install webrtcvad package.")
        
        self.vad = webrtcvad.Vad(config.webrtc_aggressiveness)
        
        # WebRTC VAD only supports 8kHz, 16kHz, 32kHz, 48kHz
        if config.sample_rate not in [8000, 16000, 32000, 48000]:
            logger.warning(f"WebRTC VAD may not work optimally with {config.sample_rate}Hz")
    
    def detect(self, audio: Union[np.ndarray, torch.Tensor]) -> VADResult:
        """Detect voice activity using WebRTC VAD."""
        audio_np = self._ensure_numpy(audio)
        
        # Ensure audio is in the right format (int16)
        if audio_np.dtype != np.int16:
            audio_np = (audio_np * 32767).astype(np.int16)
        
        # Process audio in frames
        frame_results = []
        num_frames = len(audio_np) // self.frame_size
        
        for i in range(num_frames):
            start_idx = i * self.frame_size
            end_idx = start_idx + self.frame_size
            frame = audio_np[start_idx:end_idx]
            
            # WebRTC VAD requires exact frame size
            if len(frame) == self.frame_size:
                try:
                    is_speech = self.vad.is_speech(frame.tobytes(), self.sample_rate)
                    frame_results.append(is_speech)
                except Exception as e:
                    logger.warning(f"WebRTC VAD error: {e}")
                    frame_results.append(False)
            else:
                frame_results.append(False)
        
        # Calculate overall result
        if not frame_results:
            return VADResult(False, 0.0, 0.0, len(audio_np) / self.sample_rate, frame_results)
        
        speech_frames = sum(frame_results)
        total_frames = len(frame_results)
        confidence = speech_frames / total_frames if total_frames > 0 else 0.0
        
        # Apply temporal smoothing if enabled
        if self.config.use_temporal_smoothing:
            frame_results = self._apply_temporal_smoothing(frame_results)
            is_speech = any(frame_results)
        else:
            is_speech = confidence > 0.3  # Threshold for overall speech detection
        
        duration = len(audio_np) / self.sample_rate
        return VADResult(is_speech, confidence, 0.0, duration, frame_results)
    
    def _apply_temporal_smoothing(self, frame_results: List[bool]) -> List[bool]:
        """Apply temporal smoothing to frame results."""
        if len(frame_results) < self.config.smoothing_window_size:
            return frame_results
        
        smoothed = []
        window_size = self.config.smoothing_window_size
        half_window = window_size // 2
        
        for i in range(len(frame_results)):
            start = max(0, i - half_window)
            end = min(len(frame_results), i + half_window + 1)
            window = frame_results[start:end]
            
            # Majority vote in window
            speech_count = sum(window)
            smoothed.append(speech_count > len(window) // 2)
        
        return smoothed


class EnergyBasedVAD(VoiceActivityDetector):
    """Energy-based Voice Activity Detector."""
    
    def __init__(self, config: VADConfig):
        super().__init__(config)
        self.energy_history = []
        self.background_energy = None
        self.adaptation_rate = 0.01
    
    def detect(self, audio: Union[np.ndarray, torch.Tensor]) -> VADResult:
        """Detect voice activity using energy-based method."""
        audio_tensor = self._ensure_torch(audio)
        
        # Calculate frame-wise energy
        frame_energies = []
        frame_results = []
        num_frames = len(audio_tensor) // self.frame_size
        
        for i in range(num_frames):
            start_idx = i * self.frame_size
            end_idx = start_idx + self.frame_size
            frame = audio_tensor[start_idx:end_idx]
            
            # Calculate RMS energy
            energy = torch.sqrt(torch.mean(frame ** 2)).item()
            frame_energies.append(energy)
            
            # Update background energy estimate
            self._update_background_energy(energy)
            
            # Determine if frame contains speech
            is_speech = self._is_speech_frame(energy)
            frame_results.append(is_speech)
        
        # Post-process results
        frame_results = self._post_process_frames(frame_results)
        
        # Calculate overall result
        if not frame_results:
            duration = len(audio_tensor) / self.sample_rate
            return VADResult(False, 0.0, 0.0, duration, frame_results, frame_energies)
        
        speech_frames = sum(frame_results)
        total_frames = len(frame_results)
        confidence = speech_frames / total_frames if total_frames > 0 else 0.0
        is_speech = confidence > 0.2
        
        duration = len(audio_tensor) / self.sample_rate
        return VADResult(is_speech, confidence, 0.0, duration, frame_results, frame_energies)
    
    def _update_background_energy(self, energy: float):
        """Update background energy estimate using adaptive method."""
        if self.background_energy is None:
            self.background_energy = energy
        else:
            # Only update if energy is low (likely background)
            if energy < self.background_energy * 2:
                self.background_energy = (1 - self.adaptation_rate) * self.background_energy + \
                                       self.adaptation_rate * energy
    
    def _is_speech_frame(self, energy: float) -> bool:
        """Determine if a frame contains speech based on energy."""
        if self.background_energy is None:
            return energy > self.config.energy_threshold
        
        # Adaptive threshold based on background energy
        adaptive_threshold = max(
            self.config.energy_threshold,
            self.background_energy * 3  # 3x background energy
        )
        
        return energy > adaptive_threshold
    
    def _post_process_frames(self, frame_results: List[bool]) -> List[bool]:
        """Post-process frame results to remove short segments and add hangover."""
        if not frame_results:
            return frame_results
        
        # Convert to frame indices
        min_length_frames = int(self.config.energy_min_length_ms / self.config.frame_duration_ms)
        hangover_frames = int(self.config.energy_hangover_ms / self.config.frame_duration_ms)
        
        # Remove short speech segments
        processed = self._remove_short_segments(frame_results, min_length_frames)
        
        # Add hangover
        processed = self._add_hangover(processed, hangover_frames)
        
        return processed
    
    def _remove_short_segments(self, results: List[bool], min_length: int) -> List[bool]:
        """Remove speech segments shorter than minimum length."""
        if min_length <= 1:
            return results
        
        processed = results.copy()
        i = 0
        
        while i < len(processed):
            if processed[i]:  # Start of speech segment
                segment_start = i
                while i < len(processed) and processed[i]:
                    i += 1
                segment_length = i - segment_start
                
                # Remove short segments
                if segment_length < min_length:
                    for j in range(segment_start, i):
                        processed[j] = False
            else:
                i += 1
        
        return processed
    
    def _add_hangover(self, results: List[bool], hangover_frames: int) -> List[bool]:
        """Add hangover time after speech segments."""
        if hangover_frames <= 0:
            return results
        
        processed = results.copy()
        
        for i in range(len(results) - 1):
            if results[i] and not results[i + 1]:  # End of speech segment
                # Add hangover
                end_hangover = min(len(processed), i + 1 + hangover_frames)
                for j in range(i + 1, end_hangover):
                    processed[j] = True
        
        return processed


class SpectralVAD(VoiceActivityDetector):
    """Spectral-based Voice Activity Detector."""
    
    def __init__(self, config: VADConfig):
        super().__init__(config)
        self.n_fft = 512
        self.hop_length = self.frame_size
    
    def detect(self, audio: Union[np.ndarray, torch.Tensor]) -> VADResult:
        """Detect voice activity using spectral features."""
        audio_tensor = self._ensure_torch(audio)
        
        # Compute STFT
        stft = torch.stft(
            audio_tensor,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            return_complex=True
        )
        
        magnitude = torch.abs(stft)
        power = magnitude ** 2
        
        # Calculate spectral features frame by frame
        frame_results = []
        num_frames = power.size(-1)
        
        for i in range(num_frames):
            frame_power = power[:, i]
            
            # Spectral rolloff
            rolloff = self._spectral_rolloff(frame_power)
            
            # Spectral centroid
            centroid = self._spectral_centroid(frame_power)
            
            # Zero crossing rate (approximate from phase)
            zcr = self._zero_crossing_rate(audio_tensor, i)
            
            # Combine features for decision
            is_speech = self._classify_frame(rolloff, centroid, zcr)
            frame_results.append(is_speech)
        
        # Calculate overall result
        if not frame_results:
            duration = len(audio_tensor) / self.sample_rate
            return VADResult(False, 0.0, 0.0, duration, frame_results)
        
        speech_frames = sum(frame_results)
        total_frames = len(frame_results)
        confidence = speech_frames / total_frames if total_frames > 0 else 0.0
        is_speech = confidence > 0.3
        
        duration = len(audio_tensor) / self.sample_rate
        return VADResult(is_speech, confidence, 0.0, duration, frame_results)
    
    def _spectral_rolloff(self, power_spectrum: torch.Tensor, rolloff_percent: float = 0.85) -> float:
        """Calculate spectral rolloff frequency."""
        total_energy = torch.sum(power_spectrum)
        if total_energy == 0:
            return 0.0
        
        cumulative_energy = torch.cumsum(power_spectrum, dim=0)
        rolloff_threshold = rolloff_percent * total_energy
        
        rolloff_bin = torch.where(cumulative_energy >= rolloff_threshold)[0]
        if len(rolloff_bin) > 0:
            return float(rolloff_bin[0]) / len(power_spectrum)
        return 1.0
    
    def _spectral_centroid(self, power_spectrum: torch.Tensor) -> float:
        """Calculate spectral centroid."""
        if torch.sum(power_spectrum) == 0:
            return 0.0
        
        frequencies = torch.arange(len(power_spectrum), dtype=torch.float32)
        centroid = torch.sum(frequencies * power_spectrum) / torch.sum(power_spectrum)
        
        # Convert to Hz
        return float(centroid * self.sample_rate / self.n_fft)
    
    def _zero_crossing_rate(self, audio: torch.Tensor, frame_idx: int) -> float:
        """Calculate zero crossing rate for a frame."""
        start_idx = frame_idx * self.hop_length
        end_idx = min(start_idx + self.frame_size, len(audio))
        
        if end_idx <= start_idx:
            return 0.0
        
        frame = audio[start_idx:end_idx]
        zero_crossings = torch.sum(torch.diff(torch.sign(frame)) != 0)
        return float(zero_crossings) / len(frame)
    
    def _classify_frame(self, rolloff: float, centroid: float, zcr: float) -> bool:
        """Classify frame as speech or non-speech based on spectral features."""
        # Speech typically has:
        # - Higher spectral rolloff (more high-frequency content)
        # - Moderate spectral centroid
        # - Moderate zero crossing rate
        
        rolloff_score = rolloff > self.config.spectral_rolloff_threshold
        centroid_score = centroid > self.config.spectral_centroid_threshold
        zcr_score = 0.01 < zcr < 0.5  # Typical range for speech
        
        # Simple voting
        votes = sum([rolloff_score, centroid_score, zcr_score])
        return votes >= 2


class HybridVAD(VoiceActivityDetector):
    """Hybrid Voice Activity Detector combining multiple methods."""
    
    def __init__(self, config: VADConfig):
        super().__init__(config)
        
        # Initialize component VADs
        self.vads = {}
        
        if WEBRTC_AVAILABLE:
            try:
                self.vads['webrtc'] = WebRTCVAD(config)
            except Exception as e:
                logger.warning(f"Failed to initialize WebRTC VAD: {e}")
        
        self.vads['energy'] = EnergyBasedVAD(config)
        self.vads['spectral'] = SpectralVAD(config)
    
    def detect(self, audio: Union[np.ndarray, torch.Tensor]) -> VADResult:
        """Detect voice activity using hybrid approach."""
        # Get results from all available VADs
        results = {}
        for name, vad in self.vads.items():
            try:
                results[name] = vad.detect(audio)
            except Exception as e:
                logger.warning(f"VAD {name} failed: {e}")
                continue
        
        if not results:
            duration = len(audio) / self.sample_rate
            return VADResult(False, 0.0, 0.0, duration)
        
        # Combine results
        return self._combine_results(results, len(audio) / self.sample_rate)
    
    def _combine_results(self, results: Dict[str, VADResult], duration: float) -> VADResult:
        """Combine results from multiple VADs."""
        if self.config.combination_mode == "majority":
            return self._majority_vote(results, duration)
        elif self.config.combination_mode == "weighted":
            return self._weighted_combination(results, duration)
        elif self.config.combination_mode == "conservative":
            return self._conservative_combination(results, duration)
        elif self.config.combination_mode == "aggressive":
            return self._aggressive_combination(results, duration)
        else:
            # Default to weighted
            return self._weighted_combination(results, duration)
    
    def _majority_vote(self, results: Dict[str, VADResult], duration: float) -> VADResult:
        """Combine using majority vote."""
        speech_votes = sum(1 for result in results.values() if result.is_speech)
        total_votes = len(results)
        
        is_speech = speech_votes > total_votes // 2
        confidence = sum(result.confidence for result in results.values()) / total_votes
        
        return VADResult(is_speech, confidence, 0.0, duration)
    
    def _weighted_combination(self, results: Dict[str, VADResult], duration: float) -> VADResult:
        """Combine using weighted average."""
        weights = {
            'webrtc': self.config.webrtc_weight,
            'energy': self.config.energy_weight,
            'spectral': self.config.spectral_weight
        }
        
        total_weight = 0.0
        weighted_confidence = 0.0
        weighted_speech = 0.0
        
        for name, result in results.items():
            weight = weights.get(name, 1.0 / len(results))
            total_weight += weight
            weighted_confidence += weight * result.confidence
            weighted_speech += weight * (1.0 if result.is_speech else 0.0)
        
        if total_weight > 0:
            final_confidence = weighted_confidence / total_weight
            final_speech = weighted_speech / total_weight
        else:
            final_confidence = 0.0
            final_speech = 0.0
        
        is_speech = final_speech > 0.5
        
        return VADResult(is_speech, final_confidence, 0.0, duration)
    
    def _conservative_combination(self, results: Dict[str, VADResult], duration: float) -> VADResult:
        """Conservative combination - all VADs must agree for speech."""
        is_speech = all(result.is_speech for result in results.values())
        confidence = min(result.confidence for result in results.values()) if results else 0.0
        
        return VADResult(is_speech, confidence, 0.0, duration)
    
    def _aggressive_combination(self, results: Dict[str, VADResult], duration: float) -> VADResult:
        """Aggressive combination - any VAD detecting speech triggers speech."""
        is_speech = any(result.is_speech for result in results.values())
        confidence = max(result.confidence for result in results.values()) if results else 0.0
        
        return VADResult(is_speech, confidence, 0.0, duration)


def create_vad_detector(vad_type: str, config: Optional[VADConfig] = None) -> VoiceActivityDetector:
    """
    Factory function to create VAD detector.
    
    Args:
        vad_type: Type of VAD ("webrtc", "energy", "spectral", "hybrid")
        config: VAD configuration
        
    Returns:
        Configured VAD detector
    """
    if config is None:
        config = VADConfig()
    
    if vad_type.lower() == "webrtc":
        return WebRTCVAD(config)
    elif vad_type.lower() == "energy":
        return EnergyBasedVAD(config)
    elif vad_type.lower() == "spectral":
        return SpectralVAD(config)
    elif vad_type.lower() == "hybrid":
        return HybridVAD(config)
    else:
        raise ValueError(f"Unsupported VAD type: {vad_type}")


def combine_vad_results(results: List[VADResult], method: str = "weighted") -> VADResult:
    """
    Combine multiple VAD results.
    
    Args:
        results: List of VAD results to combine
        method: Combination method ("majority", "weighted", "conservative", "aggressive")
        
    Returns:
        Combined VAD result
    """
    if not results:
        return VADResult(False, 0.0, 0.0, 0.0)
    
    if len(results) == 1:
        return results[0]
    
    # Calculate combined metrics
    confidences = [r.confidence for r in results]
    speech_flags = [r.is_speech for r in results]
    
    if method == "majority":
        is_speech = sum(speech_flags) > len(speech_flags) // 2
        confidence = sum(confidences) / len(confidences)
    elif method == "weighted":
        # Weight by confidence
        weights = [r.confidence for r in results]
        total_weight = sum(weights) if sum(weights) > 0 else len(weights)
        
        weighted_speech = sum(w * s for w, s in zip(weights, speech_flags)) / total_weight
        is_speech = weighted_speech > 0.5
        confidence = sum(confidences) / len(confidences)
    elif method == "conservative":
        is_speech = all(speech_flags)
        confidence = min(confidences)
    elif method == "aggressive":
        is_speech = any(speech_flags)
        confidence = max(confidences)
    else:
        raise ValueError(f"Unsupported combination method: {method}")
    
    # Use timing from first result
    start_time = results[0].start_time
    end_time = results[0].end_time
    
    return VADResult(is_speech, confidence, start_time, end_time)


# Example usage and testing
def test_vad_systems():
    """Test different VAD systems with sample audio."""
    # Generate test audio (sine wave + noise)
    sample_rate = 16000
    duration = 2.0
    t = np.linspace(0, duration, int(sample_rate * duration))
    
    # Create audio with speech-like characteristics
    audio = np.sin(2 * np.pi * 200 * t) * np.sin(2 * np.pi * 3 * t)  # AM modulation
    audio += 0.1 * np.random.randn(len(audio))  # Add noise
    
    # Test different VAD systems
    config = VADConfig()
    
    # Test energy-based VAD
    energy_vad = create_vad_detector("energy", config)
    energy_result = energy_vad.detect(audio)
    print(f"Energy VAD: speech={energy_result.is_speech}, confidence={energy_result.confidence:.3f}")
    
    # Test spectral VAD
    spectral_vad = create_vad_detector("spectral", config)
    spectral_result = spectral_vad.detect(audio)
    print(f"Spectral VAD: speech={spectral_result.is_speech}, confidence={spectral_result.confidence:.3f}")
    
    # Test WebRTC VAD if available
    if WEBRTC_AVAILABLE:
        try:
            webrtc_vad = create_vad_detector("webrtc", config)
            webrtc_result = webrtc_vad.detect(audio)
            print(f"WebRTC VAD: speech={webrtc_result.is_speech}, confidence={webrtc_result.confidence:.3f}")
        except Exception as e:
            print(f"WebRTC VAD failed: {e}")
    
    # Test hybrid VAD
    hybrid_vad = create_vad_detector("hybrid", config)
    hybrid_result = hybrid_vad.detect(audio)
    print(f"Hybrid VAD: speech={hybrid_result.is_speech}, confidence={hybrid_result.confidence:.3f}")


if __name__ == "__main__":
    test_vad_systems()