Spaces:
Sleeping
Sleeping
| """ | |
| backend/src/models/dataset.py | |
| PyTorch Dataset and Collation implementations for Speech Emotion Recognition. | |
| """ | |
| import soundfile as sf | |
| import pandas as pd | |
| import torch | |
| from torch.utils.data import Dataset | |
| from pathlib import Path | |
| from typing import Dict, Any, List | |
| from src.models.config import EMOTION_LABELS | |
| class SpeechEmotionDataset(Dataset): | |
| """ | |
| Dataset representing preprocessed audio files mapped to their target labels. | |
| """ | |
| def __init__(self, manifest_path: Path): | |
| if not manifest_path.exists(): | |
| raise FileNotFoundError(f"Manifest not found at {manifest_path}") | |
| self.df = pd.read_csv(manifest_path) | |
| self.label_map = {emotion: idx for idx, emotion in enumerate(EMOTION_LABELS)} | |
| def __len__(self) -> int: | |
| return len(self.df) | |
| def __getitem__(self, idx: int) -> Dict[str, Any]: | |
| row = self.df.iloc[idx] | |
| file_path = row["processed_file_path"] | |
| emotion = row["emotion"] | |
| # Load audio (already preprocessed to 16kHz mono PCM) | |
| speech, sr = sf.read(file_path) | |
| return { | |
| "input_values": speech, | |
| "label": self.label_map[emotion] | |
| } | |
| class SpeechDataCollator: | |
| """ | |
| Data collator that dynamically pads waveforms in a batch to the maximum sequence length. | |
| """ | |
| def __init__(self, processor: Any): | |
| self.processor = processor | |
| def __call__(self, batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: | |
| input_values = [item["input_values"] for item in batch] | |
| labels = [item["label"] for item in batch] | |
| # Standard huggingface processor pads waveforms and converts to torch tensors | |
| features = self.processor( | |
| input_values, | |
| sampling_rate=16000, | |
| padding=True, | |
| return_tensors="pt" | |
| ) | |
| # Map labels into torch tensors | |
| features["labels"] = torch.tensor(labels, dtype=torch.long) | |
| return features | |