Spaces:
Running on Zero
Running on Zero
File size: 4,942 Bytes
15ff271 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """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)
|