Spaces:
Running on Zero
Running on Zero
| """PyTorch Dataset and DataLoader with Audio Augmentation for Bangla Speech Emotion Recognition.""" | |
| import os | |
| import random | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| import numpy as np | |
| import soundfile as sf | |
| import librosa | |
| import pandas as pd | |
| from src.utils.config import ( | |
| METADATA_DIR, | |
| TARGET_SAMPLE_RATE, | |
| MAX_DURATION_SECONDS, | |
| EMOTION_TO_ID, | |
| EMOTION_CLASSES | |
| ) | |
| class AudioAugmenter: | |
| """Acoustic augmentations tailored for speech emotion recognition without distorting emotion.""" | |
| def __init__(self, p: float = 0.5): | |
| self.p = p | |
| def apply(self, y: np.ndarray, sr: int = TARGET_SAMPLE_RATE) -> np.ndarray: | |
| if random.random() > self.p: | |
| return y | |
| aug_type = random.choice(["noise", "gain", "shift", "pitch"]) | |
| if aug_type == "noise": | |
| noise_amp = random.uniform(0.002, 0.015) | |
| noise = np.random.randn(len(y)) * noise_amp | |
| y = y + noise | |
| elif aug_type == "gain": | |
| gain_factor = random.uniform(0.8, 1.2) | |
| y = y * gain_factor | |
| elif aug_type == "shift": | |
| shift_samples = int(random.uniform(-0.15, 0.15) * sr) | |
| y = np.roll(y, shift_samples) | |
| elif aug_type == "pitch": | |
| # Mild pitch variation (-1 to +1 semitone) | |
| n_steps = random.choice([-1.0, -0.5, 0.5, 1.0]) | |
| y = librosa.effects.pitch_shift(y, sr=sr, n_steps=n_steps) | |
| # Clip to prevent digital clipping | |
| y = np.clip(y, -1.0, 1.0) | |
| return y | |
| class BanglaSpeechDataset(Dataset): | |
| """Dataset delivering log Mel-spectrograms or raw waveforms with labels.""" | |
| def __init__( | |
| self, | |
| split: str = "train", | |
| feature_type: str = "melspec", | |
| n_mels: int = 128, | |
| augment: bool = False, | |
| cached_mel_dict: dict = None | |
| ): | |
| self.split = split | |
| self.feature_type = feature_type | |
| self.n_mels = n_mels | |
| self.augment = augment | |
| self.augmenter = AudioAugmenter(p=0.5) if augment else None | |
| self.target_len = int(MAX_DURATION_SECONDS * TARGET_SAMPLE_RATE) # 160,000 samples | |
| meta_path = METADATA_DIR / "metadata.csv" | |
| df = pd.read_csv(meta_path) | |
| self.df = df[df["split"] == split].reset_index(drop=True) | |
| # In-memory cached mel spectrograms if provided | |
| self.cached_mel = cached_mel_dict.get(split) if cached_mel_dict else None | |
| def __len__(self): | |
| return len(self.df) | |
| def __getitem__(self, idx): | |
| row = self.df.iloc[idx] | |
| emotion = row["emotion"] | |
| label_id = EMOTION_TO_ID[emotion] | |
| audio_id = row["audio_id"] | |
| if self.feature_type == "melspec" and self.cached_mel is not None: | |
| # Ultra-fast indexed from RAM | |
| mel = self.cached_mel[idx] | |
| tensor_mel = torch.tensor(mel, dtype=torch.float32).unsqueeze(0) | |
| return tensor_mel, label_id, audio_id | |
| file_path = row["processed_file_path"] | |
| audio, sr = sf.read(file_path, dtype="float32") | |
| if self.augment and self.augmenter: | |
| audio = self.augmenter.apply(audio, sr=sr) | |
| # Pad or truncate to fixed length (10 seconds) | |
| if len(audio) < self.target_len: | |
| audio = np.pad(audio, (0, self.target_len - len(audio)), mode="constant") | |
| else: | |
| audio = audio[:self.target_len] | |
| if self.feature_type == "raw": | |
| return torch.tensor(audio, dtype=torch.float32), label_id, audio_id | |
| elif self.feature_type == "melspec": | |
| # Compute 128 Mel bands | |
| mel_spec = librosa.feature.melspectrogram( | |
| y=audio, sr=sr, n_fft=1024, hop_length=512, n_mels=self.n_mels, fmin=20, fmax=8000 | |
| ) | |
| log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max) | |
| log_mel_spec = (log_mel_spec - np.mean(log_mel_spec)) / (np.std(log_mel_spec) + 1e-6) | |
| tensor_mel = torch.tensor(log_mel_spec, dtype=torch.float32).unsqueeze(0) | |
| return tensor_mel, label_id, audio_id | |
| elif self.feature_type == "mfcc": | |
| mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=40, n_fft=1024, hop_length=512) | |
| mfcc = (mfcc - np.mean(mfcc)) / (np.std(mfcc) + 1e-6) | |
| # Shape: (313, 40) | |
| tensor_mfcc = torch.tensor(mfcc.T, dtype=torch.float32) | |
| return tensor_mfcc, label_id, audio_id | |
| def get_dataloader( | |
| split: str = "train", | |
| feature_type: str = "melspec", | |
| batch_size: int = 16, | |
| augment: bool = False, | |
| num_workers: int = 0, | |
| cached_mel_dict: dict = None | |
| ) -> DataLoader: | |
| dataset = BanglaSpeechDataset( | |
| split=split, | |
| feature_type=feature_type, | |
| augment=augment, | |
| cached_mel_dict=cached_mel_dict | |
| ) | |
| shuffle = (split == "train") | |
| return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers) | |