Update docstring

Author: mthrok

examples/source_separation/utils/__init__.py

@@ -1,3 +1,4 @@
 from . import (
+    dataset,
     metrics,
 )

examples/source_separation/utils/dataset/__init__.py

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

examples/source_separation/utils/dataset_utils.py renamed to examples/source_separation/utils/dataset/utils.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from utils.dataset import wsj0mix
+from . import wsj0mix
 
 Batch = namedtuple("Batch", ["mix", "src"])
 
@@ -35,13 +35,12 @@ def _fix_num_frames(waveform: torch.Tensor, target_num_frames: int):
     return torch.cat([waveform, pad], 1)
 
 
-def collate_fn_wsj0mix_train(samples: List[wsj0mix.Sample], sample_rate, duration):
+def collate_fn_wsj0mix_train(samples: List[wsj0mix.SampleType], sample_rate, duration):
     target_num_frames = int(duration * sample_rate)
 
     mixes, srcs = [], []
-    for sample in samples:
-        mix = sample.mix
-        src = torch.cat(sample.src, 0)
+    for (_, mix, src) in samples:
+        src = torch.cat(src, 0)
 
         num_frames = mix.shape[-1]
         if num_frames > target_num_frames:
@@ -58,7 +57,7 @@ def collate_fn_wsj0mix_train(samples: List[wsj0mix.Sample], sample_rate, duratio
     return Batch(torch.stack(mixes, 0), torch.stack(srcs, 0))
 
 
-def collate_fn_wsj0mix_test(samples: List[wsj0mix.Sample]):
+def collate_fn_wsj0mix_test(samples: List[wsj0mix.SampleType]) -> Batch:
     return [Batch(
         sample.mix.unsqueeze(0),
         torch.cat(sample.src, 0).unsqueeze(0),

examples/source_separation/utils/dataset/wsj0mix.py

@@ -6,17 +6,34 @@
 
 import torchaudio
 
+SampleType = Tuple[int, torch.Tensor, List[torch.Tensor]]
+
 
 class WSJ0Mix(Dataset):
+    """Create a Dataset for wsj0-mix.
+
+    Args:
+        root (str or Path): Path to the directory where the dataset is found.
+        num_speakers (int): The number of speakers, which determines the directories
+            to traverse. The Dataset will traverse ``s1`` to ``sN`` directories to collect
+            N source audios.
+        sample_rate (int): Expected sample rate of audio files. If any of the audio has a
+            different sample rate, raises ``ValueError``.
+        audio_ext (str): The extension of audio files to find. (default: ".wav")
+    """
     def __init__(
-        self, root: Union[str, Path], num_speakers, sample_rate, audio_ext="wav"
+        self,
+        root: Union[str, Path],
+        num_speakers: int,
+        sample_rate: int,
+        audio_ext: str = ".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 = [p.name for p in self.mix_dir.glob(f"*{audio_ext}")]
         self.files.sort()
 
     def _load_audio(self, path) -> torch.Tensor:
@@ -28,7 +45,7 @@ def _load_audio(self, path) -> torch.Tensor:
             )
         return waveform
 
-    def _load_sample(self, filename) -> Tuple[int, torch.Tensor, List[torch.Tensor]]:
+    def _load_sample(self, filename) -> SampleType:
         mixed = self._load_audio(str(self.mix_dir / filename))
         srcs = []
         for i, dir_ in enumerate(self.src_dirs):
@@ -43,5 +60,11 @@ def _load_sample(self, filename) -> Tuple[int, torch.Tensor, List[torch.Tensor]]
     def __len__(self) -> int:
         return len(self.files)
 
-    def __getitem__(self, key: int) -> Tuple[int, torch.Tensor, List[torch.Tensor]]:
+    def __getitem__(self, key: int) -> SampleType:
+        """Load the n-th sample from the dataset.
+        Args:
+            n (int): The index of the sample to be loaded
+        Returns:
+            tuple: ``(sample_rate, mix_waveform, list_of_source_waveforms)``
+        """
         return self._load_sample(self.files[key])