101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
|
from pathlib import Path
|
||
|
|
|
||
|
|
from lightning.pytorch import LightningDataModule
|
||
|
|
from torch.utils.data import DataLoader
|
||
|
|
|
||
|
|
from data.dataloaders.fisher import Fisher, FisherEvent
|
||
|
|
from models.vap import audio_channels, needs_vad
|
||
|
|
|
||
|
|
|
||
|
|
class VAPDataModule(LightningDataModule):
|
||
|
|
"""Fisher segments for fitting, annotated turn-taking events for testing."""
|
||
|
|
|
||
|
|
def __init__(self, cfg, projection):
|
||
|
|
super().__init__()
|
||
|
|
vap_cfg = cfg["vap"]
|
||
|
|
self.root = Path(cfg["fisher_path"])
|
||
|
|
self.projection = projection
|
||
|
|
self.batch_size = vap_cfg["batch_size"]
|
||
|
|
self.num_workers = vap_cfg["num_workers"]
|
||
|
|
self.channels = audio_channels(vap_cfg)
|
||
|
|
self.with_vad = needs_vad(vap_cfg)
|
||
|
|
self.swap_channels = bool(vap_cfg.get("swap_channels", False))
|
||
|
|
self.source = str(vap_cfg.get("source", "npy")).lower()
|
||
|
|
|
||
|
|
def _split(self, name):
|
||
|
|
path = self.root / f"{name}.txt"
|
||
|
|
if not path.exists():
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"missing Fisher split {path}; run preprocess/vap/00_prep_fisher.py first"
|
||
|
|
)
|
||
|
|
return path
|
||
|
|
|
||
|
|
def _event_splits(self, names):
|
||
|
|
paths = [self.root / f"{name}_events.txt" for name in names]
|
||
|
|
found = [path for path in paths if path.exists()]
|
||
|
|
if not found:
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"none of {[p.name for p in paths]} exist; run preprocess/vap/02_prep_event.py first"
|
||
|
|
)
|
||
|
|
return found
|
||
|
|
|
||
|
|
def setup(self, stage=None):
|
||
|
|
if stage in ("fit", None):
|
||
|
|
self.train_dataset = self._segments("train", swap=self.swap_channels)
|
||
|
|
self.val_dataset = self._segments("val")
|
||
|
|
|
||
|
|
if stage in ("test", None):
|
||
|
|
self.tune_dataset = FisherEvent(
|
||
|
|
self.root,
|
||
|
|
self._event_splits(["train", "val"]),
|
||
|
|
channels=self.channels,
|
||
|
|
source=self.source,
|
||
|
|
with_vad=self.with_vad,
|
||
|
|
frame_hz=self.projection.frame_hz,
|
||
|
|
)
|
||
|
|
self.test_dataset = FisherEvent(
|
||
|
|
self.root,
|
||
|
|
self._event_splits(["test"]),
|
||
|
|
channels=self.channels,
|
||
|
|
source=self.source,
|
||
|
|
with_vad=self.with_vad,
|
||
|
|
frame_hz=self.projection.frame_hz,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _segments(self, name, swap=False):
|
||
|
|
return Fisher(
|
||
|
|
self.root,
|
||
|
|
self._split(name),
|
||
|
|
self.projection,
|
||
|
|
channels=self.channels,
|
||
|
|
frame_hz=self.projection.frame_hz,
|
||
|
|
swap_channels=swap,
|
||
|
|
source=self.source,
|
||
|
|
with_vad=self.with_vad,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _loader(self, dataset, shuffle=False):
|
||
|
|
workers = self.num_workers
|
||
|
|
return DataLoader(
|
||
|
|
dataset,
|
||
|
|
batch_size=self.batch_size,
|
||
|
|
shuffle=shuffle,
|
||
|
|
num_workers=workers,
|
||
|
|
pin_memory=True,
|
||
|
|
persistent_workers=bool(workers),
|
||
|
|
prefetch_factor=2 if workers else None,
|
||
|
|
drop_last=shuffle,
|
||
|
|
)
|
||
|
|
|
||
|
|
def train_dataloader(self):
|
||
|
|
return self._loader(self.train_dataset, shuffle=True)
|
||
|
|
|
||
|
|
def val_dataloader(self):
|
||
|
|
return self._loader(self.val_dataset)
|
||
|
|
|
||
|
|
def tune_dataloader(self):
|
||
|
|
return self._loader(self.tune_dataset)
|
||
|
|
|
||
|
|
def test_dataloader(self):
|
||
|
|
return self._loader(self.test_dataset)
|