Spaces:
Running
Running
| import os | |
| import re | |
| from pathlib import Path | |
| from typing import Optional, Tuple, Union | |
| import torch | |
| from torch.utils.data import Dataset | |
| from torchaudio._internal import download_url_to_file | |
| from torchaudio.datasets.utils import _extract_tar, _load_waveform | |
| URL = "https://speech.fit.vutbr.cz/files/quesst14Database.tgz" | |
| SAMPLE_RATE = 8000 | |
| _CHECKSUM = "4f869e06bc066bbe9c5dde31dbd3909a0870d70291110ebbb38878dcbc2fc5e4" | |
| _LANGUAGES = [ | |
| "albanian", | |
| "basque", | |
| "czech", | |
| "nnenglish", | |
| "romanian", | |
| "slovak", | |
| ] | |
| class QUESST14(Dataset): | |
| """*QUESST14* :cite:`Mir2015QUESST2014EQ` dataset. | |
| Args: | |
| root (str or Path): Root directory where the dataset's top level directory is found | |
| subset (str): Subset of the dataset to use. Options: [``"docs"``, ``"dev"``, ``"eval"``]. | |
| language (str or None, optional): Language to get dataset for. | |
| Options: [``None``, ``albanian``, ``basque``, ``czech``, ``nnenglish``, ``romanian``, ``slovak``]. | |
| If ``None``, dataset consists of all languages. (default: ``"nnenglish"``) | |
| download (bool, optional): Whether to download the dataset if it is not found at root path. | |
| (default: ``False``) | |
| """ | |
| def __init__( | |
| self, | |
| root: Union[str, Path], | |
| subset: str, | |
| language: Optional[str] = "nnenglish", | |
| download: bool = False, | |
| ) -> None: | |
| if subset not in ["docs", "dev", "eval"]: | |
| raise ValueError("`subset` must be one of ['docs', 'dev', 'eval']") | |
| if language is not None and language not in _LANGUAGES: | |
| raise ValueError(f"`language` must be None or one of {str(_LANGUAGES)}") | |
| # Get string representation of 'root' | |
| root = os.fspath(root) | |
| basename = os.path.basename(URL) | |
| archive = os.path.join(root, basename) | |
| basename = basename.rsplit(".", 2)[0] | |
| self._path = os.path.join(root, basename) | |
| if not os.path.isdir(self._path): | |
| if not os.path.isfile(archive): | |
| if not download: | |
| raise RuntimeError("Dataset not found. Please use `download=True` to download") | |
| download_url_to_file(URL, archive, hash_prefix=_CHECKSUM) | |
| _extract_tar(archive, root) | |
| if subset == "docs": | |
| self.data = filter_audio_paths(self._path, language, "language_key_utterances.lst") | |
| elif subset == "dev": | |
| self.data = filter_audio_paths(self._path, language, "language_key_dev.lst") | |
| elif subset == "eval": | |
| self.data = filter_audio_paths(self._path, language, "language_key_eval.lst") | |
| def get_metadata(self, n: int) -> Tuple[str, int, str]: | |
| """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 audio | |
| int: | |
| Sample rate | |
| str: | |
| File name | |
| """ | |
| audio_path = self.data[n] | |
| relpath = os.path.relpath(audio_path, self._path) | |
| return relpath, SAMPLE_RATE, audio_path.with_suffix("").name | |
| def __getitem__(self, n: int) -> Tuple[torch.Tensor, int, str]: | |
| """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: | |
| File name | |
| """ | |
| metadata = self.get_metadata(n) | |
| waveform = _load_waveform(self._path, metadata[0], metadata[1]) | |
| return (waveform,) + metadata[1:] | |
| def __len__(self) -> int: | |
| return len(self.data) | |
| def filter_audio_paths( | |
| path: str, | |
| language: str, | |
| lst_name: str, | |
| ): | |
| """Extract audio paths for the given language.""" | |
| audio_paths = [] | |
| path = Path(path) | |
| with open(path / "scoring" / lst_name) as f: | |
| for line in f: | |
| audio_path, lang = line.strip().split() | |
| if language is not None and lang != language: | |
| continue | |
| audio_path = re.sub(r"^.*?\/", "", audio_path) | |
| audio_paths.append(path / audio_path) | |
| return audio_paths | |