Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/data/vap_dm.py

@@ -0,0 +1,100 @@
+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)