import os from pathlib import Path from typing import Optional, Tuple, Union from torch import Tensor from torch.utils.data import Dataset from torchaudio._internal import download_url_to_file from torchaudio.datasets.utils import _extract_tar, _load_waveform FOLDER_IN_ARCHIVE = "SpeechCommands" URL = "speech_commands_v0.02" HASH_DIVIDER = "_nohash_" EXCEPT_FOLDER = "_background_noise_" SAMPLE_RATE = 16000 _CHECKSUMS = { "http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz": "743935421bb51cccdb6bdd152e04c5c70274e935c82119ad7faeec31780d811d", # noqa: E501 "http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz": "af14739ee7dc311471de98f5f9d2c9191b18aedfe957f4a6ff791c709868ff58", # noqa: E501 } def _load_list(root, *filenames): output = [] for filename in filenames: filepath = os.path.join(root, filename) with open(filepath) as fileobj: output += [os.path.normpath(os.path.join(root, line.strip())) for line in fileobj] return output def _get_speechcommands_metadata(filepath: str, path: str) -> Tuple[str, int, str, str, int]: relpath = os.path.relpath(filepath, path) reldir, filename = os.path.split(relpath) _, label = os.path.split(reldir) # Besides the officially supported split method for datasets defined by "validation_list.txt" # and "testing_list.txt" over "speech_commands_v0.0x.tar.gz" archives, an alternative split # method referred to in paragraph 2-3 of Section 7.1, references 13 and 14 of the original # paper, and the checksums file from the tensorflow_datasets package [1] is also supported. # Some filenames in those "speech_commands_test_set_v0.0x.tar.gz" archives have the form # "xxx.wav.wav", so file extensions twice needs to be stripped twice. # [1] https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/url_checksums/speech_commands.txt speaker, _ = os.path.splitext(filename) speaker, _ = os.path.splitext(speaker) speaker_id, utterance_number = speaker.split(HASH_DIVIDER) utterance_number = int(utterance_number) return relpath, SAMPLE_RATE, label, speaker_id, utterance_number class SPEECHCOMMANDS(Dataset): """*Speech Commands* :cite:`speechcommandsv2` dataset. Args: root (str or Path): Path to the directory where the dataset is found or downloaded. url (str, optional): The URL to download the dataset from, or the type of the dataset to dowload. Allowed type values are ``"speech_commands_v0.01"`` and ``"speech_commands_v0.02"`` (default: ``"speech_commands_v0.02"``) folder_in_archive (str, optional): The top-level directory of the dataset. (default: ``"SpeechCommands"``) download (bool, optional): Whether to download the dataset if it is not found at root path. (default: ``False``). subset (str or None, optional): Select a subset of the dataset [None, "training", "validation", "testing"]. None means the whole dataset. "validation" and "testing" are defined in "validation_list.txt" and "testing_list.txt", respectively, and "training" is the rest. Details for the files "validation_list.txt" and "testing_list.txt" are explained in the README of the dataset and in the introduction of Section 7 of the original paper and its reference 12. The original paper can be found `here `_. (Default: ``None``) """ def __init__( self, root: Union[str, Path], url: str = URL, folder_in_archive: str = FOLDER_IN_ARCHIVE, download: bool = False, subset: Optional[str] = None, ) -> None: if subset is not None and subset not in ["training", "validation", "testing"]: raise ValueError("When `subset` is not None, it must be one of ['training', 'validation', 'testing'].") if url in [ "speech_commands_v0.01", "speech_commands_v0.02", ]: base_url = "http://download.tensorflow.org/data/" ext_archive = ".tar.gz" url = os.path.join(base_url, url + ext_archive) # Get string representation of 'root' in case Path object is passed root = os.fspath(root) self._archive = os.path.join(root, folder_in_archive) basename = os.path.basename(url) archive = os.path.join(root, basename) basename = basename.rsplit(".", 2)[0] folder_in_archive = os.path.join(folder_in_archive, basename) self._path = os.path.join(root, folder_in_archive) if download: if not os.path.isdir(self._path): if not os.path.isfile(archive): checksum = _CHECKSUMS.get(url, None) download_url_to_file(url, archive, hash_prefix=checksum) _extract_tar(archive, self._path) else: if not os.path.exists(self._path): raise RuntimeError( f"The path {self._path} doesn't exist. " "Please check the ``root`` path or set `download=True` to download it" ) if subset == "validation": self._walker = _load_list(self._path, "validation_list.txt") elif subset == "testing": self._walker = _load_list(self._path, "testing_list.txt") elif subset == "training": excludes = set(_load_list(self._path, "validation_list.txt", "testing_list.txt")) walker = sorted(str(p) for p in Path(self._path).glob("*/*.wav")) self._walker = [ w for w in walker if HASH_DIVIDER in w and EXCEPT_FOLDER not in w and os.path.normpath(w) not in excludes ] else: walker = sorted(str(p) for p in Path(self._path).glob("*/*.wav")) self._walker = [w for w in walker if HASH_DIVIDER in w and EXCEPT_FOLDER not in w] def get_metadata(self, n: int) -> Tuple[str, int, str, str, int]: """Get metadata for the n-th sample from the dataset. Returns filepath instead of waveform, but otherwise returns the same fields as :py:func:`__getitem__`. Args: n (int): The index of the sample to be loaded Returns: Tuple of the following items; str: Path to the audio int: Sample rate str: Label str: Speaker ID int: Utterance number """ fileid = self._walker[n] return _get_speechcommands_metadata(fileid, self._archive) def __getitem__(self, n: int) -> Tuple[Tensor, int, str, str, int]: """Load the n-th sample from the dataset. Args: n (int): The index of the sample to be loaded Returns: Tuple of the following items; Tensor: Waveform int: Sample rate str: Label str: Speaker ID int: Utterance number """ metadata = self.get_metadata(n) waveform = _load_waveform(self._archive, metadata[0], metadata[1]) return (waveform,) + metadata[1:] def __len__(self) -> int: return len(self._walker)