Add wsj0mix dataset

Author: mthrok

examples/source_separation/test/__init__.py

@@ -0,0 +1,7 @@
+import os
+import sys
+
+_THIS_DIR = os.path.abspath(os.path.dirname(__file__))
+
+
+sys.path.append(os.path.join(_THIS_DIR, "..", "..", "..", "test"))

examples/source_separation/test/dataset/wsj0mix_test.py

@@ -0,0 +1,109 @@
+import os
+
+from torchaudio_unittest.common_utils import (
+    TempDirMixin,
+    TorchaudioTestCase,
+    get_whitenoise,
+    save_wav,
+    normalize_wav,
+)
+
+from utils.dataset import wsj0mix
+
+
+_FILENAMES = [
+    "012c0207_1.9952_01cc0202_-1.9952.wav",
+    "01co0302_1.63_014c020q_-1.63.wav",
+    "01do0316_0.24011_205a0104_-0.24011.wav",
+    "01lc020x_1.1301_027o030r_-1.1301.wav",
+    "01mc0202_0.34056_205o0106_-0.34056.wav",
+    "01nc020t_0.53821_018o030w_-0.53821.wav",
+    "01po030f_2.2136_40ko031a_-2.2136.wav",
+    "01ra010o_2.4098_403a010f_-2.4098.wav",
+    "01xo030b_0.22377_016o031a_-0.22377.wav",
+    "02ac020x_0.68566_01ec020b_-0.68566.wav",
+    "20co010m_0.82801_019c0212_-0.82801.wav",
+    "20da010u_1.2483_017c0211_-1.2483.wav",
+    "20oo010d_1.0631_01ic020s_-1.0631.wav",
+    "20sc0107_2.0222_20fo010h_-2.0222.wav",
+    "20tc010f_0.051456_404a0110_-0.051456.wav",
+    "407c0214_1.1712_02ca0113_-1.1712.wav",
+    "40ao030w_2.4697_20vc010a_-2.4697.wav",
+    "40pa0101_1.1087_40ea0107_-1.1087.wav",
+]
+
+
+def _mock_dataset(root_dir, num_speaker):
+    dirnames = ["mix"] + [f"s{i+1}" for i in range(num_speaker)]
+    for dirname in dirnames:
+        os.makedirs(os.path.join(root_dir, dirname), exist_ok=True)
+
+    seed = 0
+    sample_rate = 8000
+    expected = []
+    for filename in _FILENAMES:
+        mix = None
+        src = []
+        for dirname in dirnames:
+            waveform = get_whitenoise(
+                sample_rate=8000, duration=1, n_channels=1, dtype="int16", seed=seed
+            )
+            seed += 1
+
+            path = os.path.join(root_dir, dirname, filename)
+            save_wav(path, waveform, sample_rate)
+            waveform = normalize_wav(waveform)
+
+            if dirname == "mix":
+                mix = waveform
+            else:
+                src.append(waveform)
+        expected.append(wsj0mix.Sample(sample_rate, mix, src))
+    return expected
+
+
+class TestWSJ0Mix2(TempDirMixin, TorchaudioTestCase):
+    backend = "default"
+    root_dir = None
+    expected = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.root_dir = cls.get_base_temp_dir()
+        cls.expected = _mock_dataset(cls.root_dir, 2)
+
+    def test_wsj0mix(self):
+        dataset = wsj0mix.WSJ0Mix(self.root_dir, num_speakers=2, sample_rate=8000)
+
+        n_ite = 0
+        for i, sample in enumerate(dataset):
+            expected = self.expected[i]
+            self.assertEqual(sample.mix, expected.mix, atol=5e-5, rtol=1e-8)
+            self.assertEqual(sample.src[0], expected.src[0], atol=5e-5, rtol=1e-8)
+            self.assertEqual(sample.src[1], expected.src[1], atol=5e-5, rtol=1e-8)
+            n_ite += 1
+        assert n_ite == len(self.expected)
+
+
+class TestWSJ0Mix3(TempDirMixin, TorchaudioTestCase):
+    backend = "default"
+    root_dir = None
+    expected = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.root_dir = cls.get_base_temp_dir()
+        cls.expected = _mock_dataset(cls.root_dir, 3)
+
+    def test_wsj0mix(self):
+        dataset = wsj0mix.WSJ0Mix(self.root_dir, num_speakers=3, sample_rate=8000)
+
+        n_ite = 0
+        for i, sample in enumerate(dataset):
+            expected = self.expected[i]
+            self.assertEqual(sample.mix, expected.mix, atol=5e-5, rtol=1e-8)
+            self.assertEqual(sample.src[0], expected.src[0], atol=5e-5, rtol=1e-8)
+            self.assertEqual(sample.src[1], expected.src[1], atol=5e-5, rtol=1e-8)
+            self.assertEqual(sample.src[2], expected.src[2], atol=5e-5, rtol=1e-8)
+            n_ite += 1
+        assert n_ite == len(self.expected)

examples/source_separation/utils/dataset/__init__.py

@@ -0,0 +1 @@
+from . import wsj0mix

examples/source_separation/utils/dataset/wsj0mix.py

@@ -0,0 +1,50 @@
+from typing import Union
+from pathlib import Path
+from collections import namedtuple
+
+import torch
+from torch.utils.data import Dataset
+
+import torchaudio
+
+Sample = namedtuple("Sample", ["sample_rate", "mix", "src"])
+
+
+class WSJ0Mix(Dataset):
+    def __init__(
+        self, root: Union[str, Path], num_speakers, sample_rate, audio_ext="wav"
+    ):
+        self.root = Path(root)
+        self.sample_rate = sample_rate
+        self.mix_dir = (self.root / "mix").resolve()
+        self.src_dirs = [(self.root / f"s{i+1}").resolve() for i in range(num_speakers)]
+
+        self.files = [p.name for p in self.mix_dir.glob(f"*.{audio_ext}")]
+        self.files.sort()
+
+    def _load_audio(self, path) -> torch.Tensor:
+        waveform, sample_rate = torchaudio.load(path)
+        if sample_rate != self.sample_rate:
+            raise ValueError(
+                f"The dataset contains audio file of sample rate {sample_rate}. "
+                "Where the requested sample rate is {self.sample_rate}."
+            )
+        return waveform
+
+    def _load_sample(self, filename) -> Sample:
+        mixed = self._load_audio(str(self.mix_dir / filename))
+        srcs = []
+        for i, dir_ in enumerate(self.src_dirs):
+            src = self._load_audio(str(dir_ / filename))
+            if mixed.shape != src.shape:
+                raise ValueError(
+                    f"Different waveform shapes. mixed: {mixed.shape}, src[{i}]: {src.shape}"
+                )
+            srcs.append(src)
+        return Sample(self.sample_rate, mixed, srcs)
+
+    def __len__(self) -> int:
+        return len(self.files)
+
+    def __getitem__(self, key: int) -> Sample:
+        return self._load_sample(self.files[key])

examples/source_separation/utils/dataset_utils.py

@@ -0,0 +1,52 @@
+from typing import List
+from functools import partial
+from collections import namedtuple
+
+import torch
+
+from utils.dataset import wsj0mix
+
+Batch = namedtuple("Batch", ["mix", "src"])
+
+
+def get_dataset(dataset_type, root_dir, num_speakers, sample_rate):
+    if dataset_type == "wsj0mix":
+        train = wsj0mix.WSJ0Mix(root_dir / "tr", num_speakers, sample_rate)
+        validation = wsj0mix.WSJ0Mix(root_dir / "cv", num_speakers, sample_rate)
+        evaluation = wsj0mix.WSJ0Mix(root_dir / "tt", num_speakers, sample_rate)
+    else:
+        raise ValueError(f"Unexpected dataset: {dataset_type}")
+    return train, validation, evaluation
+
+
+def _fix_num_frames(waveform: torch.Tensor, target_num_frames: int):
+    """Ensure waveform has exact number of frames by slicing or padding"""
+    num_channels, num_frames = waveform.shape
+    if num_frames == target_num_frames:
+        return waveform
+    if num_frames > target_num_frames:
+        return waveform[..., :target_num_frames]
+    pad = torch.zeros(
+        num_channels,
+        target_num_frames - num_frames,
+        dtype=waveform.dtype,
+        device=waveform.device,
+    )
+    return torch.cat([waveform, pad], 1)
+
+
+def collate_fn_wsj0mix(samples: List[wsj0mix.Sample], sample_rate, duration):
+    target_num_frames = int(duration * sample_rate)
+
+    mixed = [_fix_num_frames(s.mix, target_num_frames) for s in samples]
+    mixed = torch.stack(mixed, 0)
+
+    src = [_fix_num_frames(torch.cat(s.src, 0), target_num_frames) for s in samples]
+    src = torch.stack(src, 0)
+    return Batch(mixed, src)
+
+
+def get_collate_fn(dataset_type, sample_rate, duration=4):
+    if dataset_type == "wsj0mix":
+        return partial(collate_fn_wsj0mix, sample_rate=sample_rate, duration=duration)
+    raise ValueError(f"Unexpected dataset: {dataset_type}")