""" Wakeword-specific data handling and preprocessing. This module provides specialized data handling for wakeword detection tasks, including dataset creation, preprocessing, and augmentation specific to wakeword detection requirements. """ import os import logging import random import numpy as np from pathlib import Path from typing import Dict, List, Tuple, Optional, Any, Union from collections import defaultdict import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader import torchaudio import torchaudio.transforms as T from ..data_pipeline import AudioDataset, AudioProcessor, AudioAugmentation, AudioProcessingConfig class WakewordDataset(AudioDataset): """ Specialized dataset for wakeword detection. Extends the base AudioDataset with wakeword-specific functionality including class balancing, negative sample generation, and specialized augmentation strategies. """ def __init__(self, audio_files: List[str], labels: List[int], class_names: Dict[int, str], audio_processor: AudioProcessor, augmentation: Optional[AudioAugmentation] = None, augmentation_config: Optional[Dict[str, float]] = None, apply_augmentation: bool = True, balance_classes: bool = True, negative_augmentation_factor: float = 2.0): """ Initialize wakeword dataset. Args: audio_files: List of audio file paths labels: List of corresponding labels class_names: Mapping of class IDs to names audio_processor: Audio processor instance augmentation: Audio augmentation instance augmentation_config: Augmentation configuration apply_augmentation: Whether to apply augmentation balance_classes: Whether to balance classes using sampling negative_augmentation_factor: Extra augmentation factor for negative samples """ super().__init__( audio_files, labels, audio_processor, augmentation, augmentation_config, apply_augmentation ) self.class_names = class_names self.balance_classes = balance_classes self.negative_augmentation_factor = negative_augmentation_factor # Analyze class distribution self.class_distribution = self._analyze_class_distribution() # Create sampling weights for balanced training if balance_classes: self.sample_weights = self._calculate_sample_weights() else: self.sample_weights = None def _analyze_class_distribution(self) -> Dict[int, int]: """Analyze the distribution of classes in the dataset.""" distribution = defaultdict(int) for label in self.labels: distribution[label] += 1 return dict(distribution) def _calculate_sample_weights(self) -> torch.Tensor: """Calculate sample weights for balanced training.""" class_counts = [self.class_distribution.get(i, 1) for i in range(len(self.class_names))] total_samples = sum(class_counts) # Calculate weights inversely proportional to class frequency class_weights = [total_samples / (len(self.class_names) * count) for count in class_counts] # Assign weights to each sample sample_weights = torch.zeros(len(self.labels)) for i, label in enumerate(self.labels): sample_weights[i] = class_weights[label] return sample_weights def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: """Get item with wakeword-specific processing.""" audio_file = self.audio_files[idx] label = self.labels[idx] class_name = self.class_names.get(label, "unknown") try: # Load and process audio waveform, sr = self.audio_processor.load_audio(audio_file) # Convert to mono if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample waveform = self.audio_processor.resample(waveform, sr) # Apply augmentation if enabled if (self.apply_augmentation and self.augmentation is not None and self.augmentation_config): # Apply extra augmentation for negative samples if class_name in ["negative", "background"]: # More aggressive augmentation for negative samples enhanced_config = self.augmentation_config.copy() for key in enhanced_config: if key.endswith('_factor'): enhanced_config[key] *= self.negative_augmentation_factor waveform = self.augmentation.apply_augmentations(waveform, enhanced_config) else: waveform = self.augmentation.apply_augmentations(waveform, self.augmentation_config) # Normalize length and extract features waveform = self.audio_processor.normalize_length(waveform) features = self.audio_processor.extract_features(waveform) # Ensure consistent shape if features.dim() == 3: # (1, freq, time) features = features.squeeze(0) # (freq, time) # Apply wakeword-specific feature normalization features = self._normalize_features(features) return features, label except Exception as e: logging.warning(f"Error processing {audio_file}: {str(e)}") # Return zero tensor and label on error dummy_features = torch.zeros( self.audio_processor.config.n_mels, self.audio_processor.config.target_samples // self.audio_processor.config.hop_length + 1 ) return self._normalize_features(dummy_features), label def _normalize_features(self, features: torch.Tensor) -> torch.Tensor: """Apply wakeword-specific feature normalization.""" # Apply instance normalization to make model more robust # to volume variations and recording conditions # Calculate mean and std along time dimension mean = torch.mean(features, dim=-1, keepdim=True) std = torch.std(features, dim=-1, keepdim=True) + 1e-8 # Normalize features = (features - mean) / std # Clip extreme values features = torch.clamp(features, -3.0, 3.0) return features def get_class_distribution(self) -> Dict[str, int]: """Get class distribution with class names.""" return { self.class_names[class_id]: count for class_id, count in self.class_distribution.items() } class WakewordDataPreprocessor: """ Specialized preprocessor for wakeword detection data. Handles the specific requirements mentioned in the project description: - Trimming wakeword files to consistent length - Splitting negative/background files into chunks - Creating balanced train/val/test splits """ def __init__(self, audio_config: AudioProcessingConfig, logger: Optional[logging.Logger] = None): """Initialize wakeword data preprocessor.""" self.audio_config = audio_config self.logger = logger or logging.getLogger(__name__) self.audio_processor = AudioProcessor(audio_config) def preprocess_wakeword_data(self, raw_data_dir: str, output_dir: str, wakeword_classes: Dict[int, str], max_wakeword_length: Optional[float] = None) -> str: """ Preprocess raw wakeword data according to project requirements. This method implements the preprocessing pipeline described in the project description: 1. Find max wakeword length 2. Trim all files to consistent length with silence padding 3. Split negative/background files into chunks 4. Create balanced train/val/test splits Args: raw_data_dir: Directory containing raw wakeword data output_dir: Output directory for processed data wakeword_classes: Mapping of class IDs to names max_wakeword_length: Maximum wakeword length (auto-detect if None) Returns: Path to processed data directory """ self.logger.info(f"Preprocessing wakeword data from {raw_data_dir}") raw_path = Path(raw_data_dir) output_path = Path(output_dir) # Step 1: Analyze wakeword lengths to determine target length if max_wakeword_length is None: max_wakeword_length = self._find_max_wakeword_length(raw_path, wakeword_classes) self.logger.info(f"Auto-detected max wakeword length: {max_wakeword_length:.2f}s") # Update audio config with detected length self.audio_config.target_length = max_wakeword_length target_samples = int(self.audio_config.sample_rate * max_wakeword_length) # Step 2: Process each class processed_files_by_class = {} for class_id, class_name in wakeword_classes.items(): self.logger.info(f"Processing class: {class_name}") class_files = self._process_class_files( raw_path / class_name, class_name, target_samples ) if class_files: processed_files_by_class[class_id] = class_files self.logger.info(f"Processed {len(class_files)} files for {class_name}") else: self.logger.warning(f"No files processed for class: {class_name}") # Step 3: Save processed files chunked_dir = output_path / "chunked" self._save_processed_files(processed_files_by_class, wakeword_classes, chunked_dir) # Step 4: Create train/val/test splits self._create_balanced_splits(chunked_dir, output_path, wakeword_classes) self.logger.info(f"Wakeword preprocessing completed: {output_path}") return str(output_path) def _find_max_wakeword_length(self, raw_path: Path, wakeword_classes: Dict[int, str]) -> float: """Find the maximum wakeword length after trimming silence.""" max_length = 0.0 # Only analyze actual wakeword classes (not negative/background) wakeword_only_classes = [name for name in wakeword_classes.values() if name not in ["negative", "background"]] for class_name in wakeword_only_classes: class_dir = raw_path / class_name if not class_dir.exists(): continue # Get audio files audio_files = [] for ext in ['.wav', '.mp3', '.flac']: audio_files.extend(list(class_dir.glob(f"**/*{ext}"))) for audio_file in audio_files[:20]: # Sample first 20 files try: waveform, sr = self.audio_processor.load_audio(str(audio_file)) # Convert to mono and resample if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) waveform = self.audio_processor.resample(waveform, sr) # Trim silence trimmed = self.audio_processor.trim_silence(waveform, threshold=0.01) # Calculate length in seconds length_seconds = trimmed.shape[-1] / self.audio_config.sample_rate max_length = max(max_length, length_seconds) except Exception as e: self.logger.warning(f"Error analyzing {audio_file}: {str(e)}") continue # Add some padding and ensure minimum length max_length = max(max_length * 1.2, 1.0) # 20% padding, minimum 1 second return min(max_length, 3.0) # Cap at 3 seconds for practical reasons def _process_class_files(self, class_dir: Path, class_name: str, target_samples: int) -> List[Tuple[torch.Tensor, str]]: """Process all files for a specific class.""" if not class_dir.exists(): self.logger.warning(f"Class directory not found: {class_dir}") return [] # Get audio files audio_files = [] for ext in ['.wav', '.mp3', '.flac']: audio_files.extend(list(class_dir.glob(f"**/*{ext}"))) processed_files = [] for audio_file in audio_files: try: # Load and process audio waveform, sr = self.audio_processor.load_audio(str(audio_file)) # Convert to mono and resample if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) waveform = self.audio_processor.resample(waveform, sr) if class_name in ["negative", "background"]: # Split long files into chunks chunks = self._split_into_chunks(waveform, target_samples) for i, chunk in enumerate(chunks): filename = f"{audio_file.stem}_chunk_{i:03d}.wav" processed_files.append((chunk, filename)) else: # For wakeword files: trim silence and normalize length # Trim silence first trimmed = self.audio_processor.trim_silence(waveform, threshold=0.01) # Normalize to target length with smart padding normalized = self._smart_normalize_length(trimmed, target_samples) processed_files.append((normalized, audio_file.name)) except Exception as e: self.logger.error(f"Error processing {audio_file}: {str(e)}") continue return processed_files def _smart_normalize_length(self, waveform: torch.Tensor, target_samples: int) -> torch.Tensor: """ Smart length normalization that preserves wakeword positioning. For wakewords, we want to: 1. Center the wakeword in the target duration 2. Add silence padding around it 3. Ensure the full wakeword is preserved """ current_samples = waveform.shape[-1] if current_samples >= target_samples: # If too long, crop from center start = (current_samples - target_samples) // 2 return waveform[..., start:start + target_samples] else: # If too short, pad with silence # Add random positioning to create variety max_pad_start = target_samples - current_samples pad_start = random.randint(0, max_pad_start // 2) # Bias towards center pad_end = target_samples - current_samples - pad_start return F.pad(waveform, (pad_start, pad_end)) def _split_into_chunks(self, waveform: torch.Tensor, chunk_size: int, overlap_ratio: float = 0.1) -> List[torch.Tensor]: """Split long audio into overlapping chunks.""" chunks = [] total_samples = waveform.shape[-1] if total_samples <= chunk_size: # If shorter than chunk size, pad to chunk size padded = F.pad(waveform, (0, chunk_size - total_samples)) chunks.append(padded) else: # Split into overlapping chunks step_size = int(chunk_size * (1 - overlap_ratio)) start = 0 while start + chunk_size <= total_samples: chunk = waveform[..., start:start + chunk_size] chunks.append(chunk) start += step_size # Add final chunk if there's remaining audio if start < total_samples and total_samples - start > chunk_size // 2: final_chunk = waveform[..., -chunk_size:] chunks.append(final_chunk) return chunks def _save_processed_files(self, processed_files_by_class: Dict[int, List[Tuple[torch.Tensor, str]]], wakeword_classes: Dict[int, str], output_dir: Path): """Save processed files to disk.""" for class_id, files in processed_files_by_class.items(): class_name = wakeword_classes[class_id] class_dir = output_dir / class_name class_dir.mkdir(parents=True, exist_ok=True) for waveform, filename in files: output_file = class_dir / filename torchaudio.save( str(output_file), waveform, self.audio_config.sample_rate ) def _create_balanced_splits(self, chunked_dir: Path, output_dir: Path, wakeword_classes: Dict[int, str]): """Create balanced train/validation/test splits.""" # Define split ratios train_ratio = 0.7 val_ratio = 0.2 test_ratio = 0.1 # Create output directories for split in ['train', 'val', 'test']: for class_name in wakeword_classes.values(): (output_dir / split / class_name).mkdir(parents=True, exist_ok=True) # Process each class for class_id, class_name in wakeword_classes.items(): class_dir = chunked_dir / class_name if not class_dir.exists(): continue # Get all files for this class files = list(class_dir.glob("*.wav")) if len(files) < 3: self.logger.warning(f"Too few files for {class_name} - copying to train only") # Copy all to train for file in files: import shutil shutil.copy2(file, output_dir / "train" / class_name / file.name) continue # Shuffle files random.shuffle(files) # Calculate split indices total_files = len(files) train_end = int(total_files * train_ratio) val_end = int(total_files * (train_ratio + val_ratio)) # Split files train_files = files[:train_end] val_files = files[train_end:val_end] test_files = files[val_end:] # Ensure minimum samples in each split if len(val_files) == 0 and len(files) > 1: val_files = [train_files.pop()] if len(test_files) == 0 and len(files) > 2: test_files = [train_files.pop()] # Copy files to respective directories splits = [ ("train", train_files), ("val", val_files), ("test", test_files) ] for split_name, split_files in splits: for file in split_files: import shutil dest_dir = output_dir / split_name / class_name shutil.copy2(file, dest_dir / file.name) self.logger.info(f"Split {class_name}: {len(train_files)} train, " f"{len(val_files)} val, {len(test_files)} test") def create_dataset_from_directory(self, data_dir: str, split: str, wakeword_classes: Dict[int, str], augmentation_config: Dict[str, float] = None, balance_classes: bool = True) -> WakewordDataset: """ Create wakeword dataset from processed directory. Args: data_dir: Directory containing processed data split: Data split ('train', 'val', 'test') wakeword_classes: Mapping of class IDs to names augmentation_config: Augmentation configuration balance_classes: Whether to balance classes Returns: WakewordDataset instance """ data_path = Path(data_dir) / split # Collect all files and labels audio_files = [] labels = [] class_name_to_id = {name: class_id for class_id, name in wakeword_classes.items()} for class_name, class_id in class_name_to_id.items(): class_dir = data_path / class_name if not class_dir.exists(): self.logger.warning(f"Class directory not found: {class_dir}") continue # Get all audio files for audio_file in class_dir.glob("*.wav"): audio_files.append(str(audio_file)) labels.append(class_id) if not audio_files: raise ValueError(f"No audio files found in {data_path}") # Create augmentation augmentation = None if split == "train" and augmentation_config: augmentation = AudioAugmentation(self.audio_config.sample_rate) # Create dataset dataset = WakewordDataset( audio_files=audio_files, labels=labels, class_names=wakeword_classes, audio_processor=self.audio_processor, augmentation=augmentation, augmentation_config=augmentation_config, apply_augmentation=(split == "train"), balance_classes=balance_classes ) self.logger.info(f"Created {split} dataset with {len(dataset)} samples") self.logger.info(f"Class distribution: {dataset.get_class_distribution()}") return dataset def create_wakeword_dataloaders(data_dir: str, wakeword_classes: Dict[int, str], audio_config: AudioProcessingConfig, batch_size: int = 32, num_workers: int = 4, augmentation_config: Dict[str, float] = None, balance_training: bool = True) -> Tuple[DataLoader, DataLoader, DataLoader]: """ Create train, validation, and test data loaders for wakeword detection. Args: data_dir: Directory containing processed wakeword data wakeword_classes: Mapping of class IDs to names audio_config: Audio processing configuration batch_size: Batch size for data loaders num_workers: Number of worker processes augmentation_config: Augmentation configuration balance_training: Whether to balance training data Returns: Tuple of (train_loader, val_loader, test_loader) """ preprocessor = WakewordDataPreprocessor(audio_config) # Create datasets train_dataset = preprocessor.create_dataset_from_directory( data_dir, "train", wakeword_classes, augmentation_config, balance_training ) val_dataset = preprocessor.create_dataset_from_directory( data_dir, "val", wakeword_classes, None, False ) test_dataset = preprocessor.create_dataset_from_directory( data_dir, "test", wakeword_classes, None, False ) # Create data loaders train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True ) val_loader = DataLoader( val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True ) test_loader = DataLoader( test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True ) return train_loader, val_loader, test_loader