""" Wakeword detection engine with confidence scoring and inference optimization. This module provides the core detection engine that processes audio features and generates wakeword predictions with confidence scores and threshold management. """ import time import threading from typing import Optional, Dict, List, Tuple, Callable, Union from dataclasses import dataclass from enum import Enum import warnings import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from .audio_features import AudioFeatureExtractor, SpectrogramConfig from .audio_buffer import AudioChunk from .model_loader import ModelMetadata class WakewordType(Enum): """Supported wakeword types.""" CUSTOM = "custom" SYSTEM_COMMAND = "system_command" NEGATIVE = "negative" # Not a wakeword @dataclass class DetectionResult: """Result of wakeword detection.""" wakeword_type: WakewordType confidence: float raw_scores: Dict[str, float] processing_time_ms: float chunk_id: int timestamp: float features_shape: Tuple[int, ...] @property def is_wakeword_detected(self) -> bool: """Check if a valid wakeword was detected.""" return self.wakeword_type in [WakewordType.CUSTOM, WakewordType.SYSTEM_COMMAND] def to_dict(self) -> Dict: """Convert result to dictionary.""" return { 'wakeword_type': self.wakeword_type.value, 'confidence': self.confidence, 'raw_scores': self.raw_scores, 'processing_time_ms': self.processing_time_ms, 'chunk_id': self.chunk_id, 'timestamp': self.timestamp, 'features_shape': list(self.features_shape), 'is_wakeword_detected': self.is_wakeword_detected } @dataclass class DetectionConfig: """Configuration for wakeword detection.""" # Confidence thresholds custom_threshold: float = 0.7 system_command_threshold: float = 0.8 # Higher threshold for admin commands negative_threshold: float = 0.3 # Temporal filtering use_temporal_smoothing: bool = True smoothing_window_size: int = 5 min_detections_in_window: int = 3 # Performance optimization use_half_precision: bool = False enable_model_optimization: bool = True batch_inference: bool = False max_batch_size: int = 8 # Debugging log_all_predictions: bool = False save_detection_features: bool = False def validate(self) -> bool: """Validate configuration parameters.""" if not (0.0 <= self.custom_threshold <= 1.0): return False if not (0.0 <= self.system_command_threshold <= 1.0): return False if not (0.0 <= self.negative_threshold <= 1.0): return False if self.smoothing_window_size < 1: return False if self.min_detections_in_window < 1: return False if self.max_batch_size < 1: return False return True class ConfidenceCalibrator: """ Calibrate model outputs to well-calibrated probability scores. Uses temperature scaling and Platt scaling for confidence calibration. """ def __init__(self): self.temperature = 1.0 self.platt_a = 1.0 self.platt_b = 0.0 self.is_calibrated = False def calibrate_with_temperature(self, logits: torch.Tensor) -> torch.Tensor: """Apply temperature scaling to logits.""" return logits / self.temperature def calibrate_with_platt(self, scores: torch.Tensor) -> torch.Tensor: """Apply Platt scaling to probability scores.""" return torch.sigmoid(self.platt_a * scores + self.platt_b) def set_temperature(self, temperature: float): """Set temperature parameter.""" self.temperature = max(0.1, temperature) self.is_calibrated = True def set_platt_parameters(self, a: float, b: float): """Set Platt scaling parameters.""" self.platt_a = a self.platt_b = b self.is_calibrated = True class TemporalFilter: """ Temporal filtering for reducing false positives. Implements various temporal smoothing strategies. """ def __init__(self, window_size: int, min_detections: int): self.window_size = window_size self.min_detections = min_detections self.detection_history = [] self._lock = threading.Lock() def add_detection(self, result: DetectionResult) -> bool: """ Add detection result and check if temporal condition is met. Args: result: Detection result Returns: True if temporal filtering criteria are met """ with self._lock: # Add current detection self.detection_history.append(result) # Keep only recent detections if len(self.detection_history) > self.window_size: self.detection_history.pop(0) # Count positive detections in window positive_detections = sum( 1 for r in self.detection_history if r.is_wakeword_detected and r.wakeword_type == result.wakeword_type ) return positive_detections >= self.min_detections def clear(self): """Clear detection history.""" with self._lock: self.detection_history.clear() def get_history(self) -> List[DetectionResult]: """Get detection history.""" with self._lock: return self.detection_history.copy() class WakewordDetectionEngine: """ Core wakeword detection engine with optimized inference and confidence scoring. This class handles the complete detection pipeline from audio features to wakeword classification with confidence estimation. """ def __init__(self, model: nn.Module, metadata: ModelMetadata, feature_extractor: AudioFeatureExtractor, config: DetectionConfig, device: Optional[torch.device] = None): """ Initialize wakeword detection engine. Args: model: Trained PyTorch model metadata: Model metadata feature_extractor: Audio feature extractor config: Detection configuration device: PyTorch device """ self.model = model self.metadata = metadata self.feature_extractor = feature_extractor self.config = config self.device = device or torch.device("cpu") # Validate configuration if not config.validate(): raise ValueError("Invalid detection configuration") # Move model to device and set to eval mode self.model.to(self.device) self.model.eval() # Optimize model if requested if config.enable_model_optimization: self._optimize_model() # Set up half precision if requested if config.use_half_precision and self.device.type == "cuda": self.model = self.model.half() self._use_half_precision = True else: self._use_half_precision = False # Initialize confidence calibrator self.confidence_calibrator = ConfidenceCalibrator() # Initialize temporal filter if enabled if config.use_temporal_smoothing: self.temporal_filter = TemporalFilter( config.smoothing_window_size, config.min_detections_in_window ) else: self.temporal_filter = None # Class mapping self.class_labels = metadata.class_labels self.class_to_wakeword = { "custom": WakewordType.CUSTOM, "system_command": WakewordType.SYSTEM_COMMAND, "negative": WakewordType.NEGATIVE } # Performance tracking self._inference_stats = { 'total_inferences': 0, 'total_time_ms': 0.0, 'avg_time_ms': 0.0, 'min_time_ms': float('inf'), 'max_time_ms': 0.0, 'detections_count': {wt.value: 0 for wt in WakewordType} } self._lock = threading.Lock() def _optimize_model(self): """Apply model optimizations for inference.""" try: # Try to use TorchScript if possible if hasattr(self.model, 'reparameterize'): # RepCNN-specific optimization self.model.reparameterize() # Compile model for better performance (PyTorch 2.0+) if hasattr(torch, 'compile'): self.model = torch.compile(self.model, mode="reduce-overhead") except Exception as e: warnings.warn(f"Model optimization failed: {e}") def detect_wakeword(self, audio_chunk: AudioChunk) -> DetectionResult: """ Detect wakeword in audio chunk. Args: audio_chunk: Audio chunk to process Returns: Detection result with confidence scores """ start_time = time.time() try: # Extract features features = self.feature_extractor.extract_features( audio_chunk.to_tensor(self.device), normalize_audio=True ) # Run inference with torch.no_grad(): if self._use_half_precision: features = features.half() # Forward pass logits = self.model(features) # Apply softmax to get probabilities probabilities = F.softmax(logits, dim=-1) # Apply confidence calibration if available if self.confidence_calibrator.is_calibrated: probabilities = self.confidence_calibrator.calibrate_with_platt(probabilities) # Convert to numpy for processing probs = probabilities.cpu().float().numpy().squeeze() # Create raw scores dictionary raw_scores = {} for i, label in enumerate(self.class_labels): raw_scores[label] = float(probs[i]) # Determine predicted class and confidence predicted_class_idx = np.argmax(probs) predicted_label = self.class_labels[predicted_class_idx] confidence = float(probs[predicted_class_idx]) # Map to wakeword type wakeword_type = self.class_to_wakeword.get(predicted_label, WakewordType.NEGATIVE) # Apply threshold-based filtering is_above_threshold = self._check_confidence_threshold(wakeword_type, confidence) if not is_above_threshold: wakeword_type = WakewordType.NEGATIVE confidence = max(raw_scores.get("negative", 0.0), 1.0 - confidence) # Calculate processing time processing_time_ms = (time.time() - start_time) * 1000 # Create result result = DetectionResult( wakeword_type=wakeword_type, confidence=confidence, raw_scores=raw_scores, processing_time_ms=processing_time_ms, chunk_id=audio_chunk.chunk_id, timestamp=audio_chunk.timestamp, features_shape=features.shape ) # Apply temporal filtering if enabled if self.temporal_filter: temporal_pass = self.temporal_filter.add_detection(result) if not temporal_pass and result.is_wakeword_detected: # Override detection if temporal filter doesn't pass result.wakeword_type = WakewordType.NEGATIVE result.confidence = raw_scores.get("negative", 0.5) # Update statistics self._update_stats(result) return result except Exception as e: # Return negative result on error processing_time_ms = (time.time() - start_time) * 1000 return DetectionResult( wakeword_type=WakewordType.NEGATIVE, confidence=0.0, raw_scores={"error": 1.0}, processing_time_ms=processing_time_ms, chunk_id=audio_chunk.chunk_id, timestamp=audio_chunk.timestamp, features_shape=(0,) ) def detect_batch(self, audio_chunks: List[AudioChunk]) -> List[DetectionResult]: """ Batch detection for multiple audio chunks. Args: audio_chunks: List of audio chunks Returns: List of detection results """ if not self.config.batch_inference or len(audio_chunks) == 1: # Process individually return [self.detect_wakeword(chunk) for chunk in audio_chunks] # Batch processing batch_size = min(len(audio_chunks), self.config.max_batch_size) results = [] for i in range(0, len(audio_chunks), batch_size): batch = audio_chunks[i:i + batch_size] batch_results = self._process_batch(batch) results.extend(batch_results) return results def _process_batch(self, audio_chunks: List[AudioChunk]) -> List[DetectionResult]: """Process a batch of audio chunks.""" start_time = time.time() try: # Extract features for all chunks batch_features = [] for chunk in audio_chunks: features = self.feature_extractor.extract_features( chunk.to_tensor(self.device), normalize_audio=True ) batch_features.append(features) # Stack into batch tensor batch_tensor = torch.cat(batch_features, dim=0) # Run batch inference with torch.no_grad(): if self._use_half_precision: batch_tensor = batch_tensor.half() batch_logits = self.model(batch_tensor) batch_probabilities = F.softmax(batch_logits, dim=-1) if self.confidence_calibrator.is_calibrated: batch_probabilities = self.confidence_calibrator.calibrate_with_platt(batch_probabilities) batch_probs = batch_probabilities.cpu().float().numpy() # Process results for each chunk results = [] total_processing_time = (time.time() - start_time) * 1000 processing_time_per_chunk = total_processing_time / len(audio_chunks) for i, (chunk, probs) in enumerate(zip(audio_chunks, batch_probs)): # Create raw scores raw_scores = {} for j, label in enumerate(self.class_labels): raw_scores[label] = float(probs[j]) # Determine prediction predicted_class_idx = np.argmax(probs) predicted_label = self.class_labels[predicted_class_idx] confidence = float(probs[predicted_class_idx]) wakeword_type = self.class_to_wakeword.get(predicted_label, WakewordType.NEGATIVE) # Apply threshold filtering if not self._check_confidence_threshold(wakeword_type, confidence): wakeword_type = WakewordType.NEGATIVE confidence = max(raw_scores.get("negative", 0.0), 1.0 - confidence) result = DetectionResult( wakeword_type=wakeword_type, confidence=confidence, raw_scores=raw_scores, processing_time_ms=processing_time_per_chunk, chunk_id=chunk.chunk_id, timestamp=chunk.timestamp, features_shape=batch_features[i].shape ) results.append(result) self._update_stats(result) return results except Exception as e: # Return negative results on error processing_time_per_chunk = (time.time() - start_time) * 1000 / len(audio_chunks) return [ DetectionResult( wakeword_type=WakewordType.NEGATIVE, confidence=0.0, raw_scores={"error": 1.0}, processing_time_ms=processing_time_per_chunk, chunk_id=chunk.chunk_id, timestamp=chunk.timestamp, features_shape=(0,) ) for chunk in audio_chunks ] def _check_confidence_threshold(self, wakeword_type: WakewordType, confidence: float) -> bool: """Check if confidence meets threshold for wakeword type.""" if wakeword_type == WakewordType.CUSTOM: return confidence >= self.config.custom_threshold elif wakeword_type == WakewordType.SYSTEM_COMMAND: return confidence >= self.config.system_command_threshold elif wakeword_type == WakewordType.NEGATIVE: return confidence >= self.config.negative_threshold return False def _update_stats(self, result: DetectionResult): """Update performance statistics.""" with self._lock: self._inference_stats['total_inferences'] += 1 self._inference_stats['total_time_ms'] += result.processing_time_ms self._inference_stats['avg_time_ms'] = ( self._inference_stats['total_time_ms'] / self._inference_stats['total_inferences'] ) self._inference_stats['min_time_ms'] = min( self._inference_stats['min_time_ms'], result.processing_time_ms ) self._inference_stats['max_time_ms'] = max( self._inference_stats['max_time_ms'], result.processing_time_ms ) self._inference_stats['detections_count'][result.wakeword_type.value] += 1 def update_thresholds(self, custom_threshold: Optional[float] = None, system_command_threshold: Optional[float] = None, negative_threshold: Optional[float] = None): """Update confidence thresholds.""" if custom_threshold is not None: self.config.custom_threshold = custom_threshold if system_command_threshold is not None: self.config.system_command_threshold = system_command_threshold if negative_threshold is not None: self.config.negative_threshold = negative_threshold def calibrate_confidence(self, validation_data: List[Tuple[torch.Tensor, int]]): """ Calibrate confidence scores using validation data. Args: validation_data: List of (features, label) tuples """ if not validation_data: return # Extract features and labels features_list = [] labels_list = [] for features, label in validation_data: features_list.append(features) labels_list.append(label) # Stack tensors all_features = torch.stack(features_list).to(self.device) all_labels = torch.tensor(labels_list).to(self.device) # Get model predictions with torch.no_grad(): logits = self.model(all_features) probabilities = F.softmax(logits, dim=-1) # Implement temperature scaling (simplified) # This is a basic implementation - for production use, implement proper optimization temperatures = [0.5, 1.0, 1.5, 2.0, 2.5] best_temperature = 1.0 best_score = float('inf') for temp in temperatures: scaled_probs = F.softmax(logits / temp, dim=-1) loss = F.cross_entropy(scaled_probs, all_labels) if loss < best_score: best_score = loss best_temperature = temp self.confidence_calibrator.set_temperature(best_temperature) def get_performance_stats(self) -> Dict: """Get performance statistics.""" with self._lock: stats = self._inference_stats.copy() # Add additional metrics stats.update({ 'model_name': self.metadata.model_name, 'model_type': self.metadata.model_type, 'device': str(self.device), 'use_half_precision': self._use_half_precision, 'temporal_filtering_enabled': self.temporal_filter is not None, 'confidence_calibrated': self.confidence_calibrator.is_calibrated }) return stats def reset_stats(self): """Reset performance statistics.""" with self._lock: self._inference_stats = { 'total_inferences': 0, 'total_time_ms': 0.0, 'avg_time_ms': 0.0, 'min_time_ms': float('inf'), 'max_time_ms': 0.0, 'detections_count': {wt.value: 0 for wt in WakewordType} } if self.temporal_filter: self.temporal_filter.clear() def warm_up(self, num_warmup_iterations: int = 5): """ Warm up the model for consistent performance measurements. Args: num_warmup_iterations: Number of warmup iterations """ # Create dummy input matching expected shape dummy_input = torch.randn( 1, self.metadata.input_shape[0], self.metadata.input_shape[1], self.metadata.input_shape[2] ).to(self.device) if self._use_half_precision: dummy_input = dummy_input.half() # Run warmup iterations with torch.no_grad(): for _ in range(num_warmup_iterations): _ = self.model(dummy_input) # Synchronize CUDA if using GPU if self.device.type == "cuda": torch.cuda.synchronize() def create_detection_engine(model: nn.Module, metadata: ModelMetadata, feature_extractor: AudioFeatureExtractor, config: Optional[DetectionConfig] = None, device: Optional[torch.device] = None) -> WakewordDetectionEngine: """ Factory function to create detection engine. Args: model: Trained PyTorch model metadata: Model metadata feature_extractor: Audio feature extractor config: Detection configuration (uses default if None) device: PyTorch device Returns: Detection engine instance """ if config is None: config = DetectionConfig() return WakewordDetectionEngine(model, metadata, feature_extractor, config, device)