| """PyTorch Dataset over a manifest.jsonl produced by build_dataset.py."""
|
|
|
| from __future__ import annotations
|
|
|
| import json
|
| from pathlib import Path
|
|
|
| import soundfile as sf
|
| import torch
|
| from torch.utils.data import Dataset
|
|
|
|
|
| class VantaDataset(Dataset):
|
| """Reads (mixture, target, enrollment) triples from a manifest.
|
|
|
| All clips are fixed-length at build time, so there is no padding logic
|
| here and no need for a custom collate_fn beyond the default stacker.
|
| """
|
|
|
| def __init__(self, manifest_path: str | Path, base_dir: str | Path | None = None):
|
| manifest_path = Path(manifest_path)
|
|
|
|
|
| self.base = Path(base_dir) if base_dir else manifest_path.parent.parent
|
| with manifest_path.open() as f:
|
| self.entries = [json.loads(line) for line in f]
|
|
|
| def __len__(self) -> int:
|
| return len(self.entries)
|
|
|
| def _read(self, rel: str) -> torch.Tensor:
|
|
|
| p = self.base / rel.replace("\\", "/")
|
| wav, _ = sf.read(p, dtype="float32")
|
| return torch.from_numpy(wav)
|
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor | str]:
|
| e = self.entries[idx]
|
| return {
|
| "mixture": self._read(e["mixture"]),
|
| "target": self._read(e["target"]),
|
| "enrollment": self._read(e["enrollment"]),
|
| "target_speaker": e["target_speaker"],
|
| }
|
|
|
