190 lines
6.8 KiB
Python
190 lines
6.8 KiB
Python
import os
|
|
from pathlib import Path
|
|
from typing import List, 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
|
|
|
|
URL = "train-clean-100"
|
|
FOLDER_IN_ARCHIVE = "LibriSpeech"
|
|
SAMPLE_RATE = 16000
|
|
_DATA_SUBSETS = [
|
|
"dev-clean",
|
|
"dev-other",
|
|
"test-clean",
|
|
"test-other",
|
|
"train-clean-100",
|
|
"train-clean-360",
|
|
"train-other-500",
|
|
]
|
|
_CHECKSUMS = {
|
|
"http://www.openslr.org/resources/12/dev-clean.tar.gz": "76f87d090650617fca0cac8f88b9416e0ebf80350acb97b343a85fa903728ab3", # noqa: E501
|
|
"http://www.openslr.org/resources/12/dev-other.tar.gz": "12661c48e8c3fe1de2c1caa4c3e135193bfb1811584f11f569dd12645aa84365", # noqa: E501
|
|
"http://www.openslr.org/resources/12/test-clean.tar.gz": "39fde525e59672dc6d1551919b1478f724438a95aa55f874b576be21967e6c23", # noqa: E501
|
|
"http://www.openslr.org/resources/12/test-other.tar.gz": "d09c181bba5cf717b3dee7d4d592af11a3ee3a09e08ae025c5506f6ebe961c29", # noqa: E501
|
|
"http://www.openslr.org/resources/12/train-clean-100.tar.gz": "d4ddd1d5a6ab303066f14971d768ee43278a5f2a0aa43dc716b0e64ecbbbf6e2", # noqa: E501
|
|
"http://www.openslr.org/resources/12/train-clean-360.tar.gz": "146a56496217e96c14334a160df97fffedd6e0a04e66b9c5af0d40be3c792ecf", # noqa: E501
|
|
"http://www.openslr.org/resources/12/train-other-500.tar.gz": "ddb22f27f96ec163645d53215559df6aa36515f26e01dd70798188350adcb6d2", # noqa: E501
|
|
}
|
|
|
|
|
|
def _download_librispeech(root, url):
|
|
base_url = "http://www.openslr.org/resources/12/"
|
|
ext_archive = ".tar.gz"
|
|
|
|
filename = url + ext_archive
|
|
archive = os.path.join(root, filename)
|
|
download_url = os.path.join(base_url, filename)
|
|
if not os.path.isfile(archive):
|
|
checksum = _CHECKSUMS.get(download_url, None)
|
|
download_url_to_file(download_url, archive, hash_prefix=checksum)
|
|
_extract_tar(archive)
|
|
|
|
|
|
def _get_librispeech_metadata(
|
|
fileid: str, root: str, folder: str, ext_audio: str, ext_txt: str, blist: List[str]
|
|
) -> Tuple[str, int, str, int, int, int]:
|
|
blist = blist or []
|
|
speaker_id, chapter_id, utterance_id = fileid.split("-")
|
|
|
|
# Get audio path and sample rate
|
|
fileid_audio = f"{speaker_id}-{chapter_id}-{utterance_id}"
|
|
filepath = os.path.join(folder, speaker_id, chapter_id, f"{fileid_audio}{ext_audio}")
|
|
|
|
# Load text
|
|
file_text = f"{speaker_id}-{chapter_id}{ext_txt}"
|
|
file_text = os.path.join(root, folder, speaker_id, chapter_id, file_text)
|
|
uttblist = []
|
|
with open(file_text) as ft:
|
|
for line in ft:
|
|
fileid_text, transcript = line.strip().split(" ", 1)
|
|
if fileid_audio == fileid_text:
|
|
# get utterance biasing list
|
|
for word in transcript.split():
|
|
if word in blist and word not in uttblist:
|
|
uttblist.append(word)
|
|
break
|
|
else:
|
|
# Translation not found
|
|
raise FileNotFoundError(f"Translation not found for {fileid_audio}")
|
|
|
|
return (
|
|
filepath,
|
|
SAMPLE_RATE,
|
|
transcript,
|
|
int(speaker_id),
|
|
int(chapter_id),
|
|
int(utterance_id),
|
|
uttblist,
|
|
)
|
|
|
|
|
|
class LibriSpeechBiasing(Dataset):
|
|
"""*LibriSpeech* :cite:`7178964` dataset with prefix-tree construction and biasing support.
|
|
|
|
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 ``"dev-clean"``, ``"dev-other"``, ``"test-clean"``,
|
|
``"test-other"``, ``"train-clean-100"``, ``"train-clean-360"`` and
|
|
``"train-other-500"``. (default: ``"train-clean-100"``)
|
|
folder_in_archive (str, optional):
|
|
The top-level directory of the dataset. (default: ``"LibriSpeech"``)
|
|
download (bool, optional):
|
|
Whether to download the dataset if it is not found at root path. (default: ``False``).
|
|
blist (list, optional):
|
|
The list of biasing words (default: ``[]``).
|
|
"""
|
|
|
|
_ext_txt = ".trans.txt"
|
|
_ext_audio = ".flac"
|
|
|
|
def __init__(
|
|
self,
|
|
root: Union[str, Path],
|
|
url: str = URL,
|
|
folder_in_archive: str = FOLDER_IN_ARCHIVE,
|
|
download: bool = False,
|
|
blist: List[str] = None,
|
|
) -> None:
|
|
self._url = url
|
|
if url not in _DATA_SUBSETS:
|
|
raise ValueError(f"Invalid url '{url}' given; please provide one of {_DATA_SUBSETS}.")
|
|
|
|
root = os.fspath(root)
|
|
self._archive = os.path.join(root, folder_in_archive)
|
|
self._path = os.path.join(root, folder_in_archive, url)
|
|
|
|
if not os.path.isdir(self._path):
|
|
if download:
|
|
_download_librispeech(root, url)
|
|
else:
|
|
raise RuntimeError(
|
|
f"Dataset not found at {self._path}. Please set `download=True` to download the dataset."
|
|
)
|
|
|
|
self._walker = sorted(str(p.stem) for p in Path(self._path).glob("*/*/*" + self._ext_audio))
|
|
self.blist = blist
|
|
|
|
def get_metadata(self, n: int) -> Tuple[str, int, str, int, int, 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 audio
|
|
int:
|
|
Sample rate
|
|
str:
|
|
Transcript
|
|
int:
|
|
Speaker ID
|
|
int:
|
|
Chapter ID
|
|
int:
|
|
Utterance ID
|
|
list:
|
|
List of biasing words in the utterance
|
|
"""
|
|
fileid = self._walker[n]
|
|
return _get_librispeech_metadata(fileid, self._archive, self._url, self._ext_audio, self._ext_txt, self.blist)
|
|
|
|
def __getitem__(self, n: int) -> Tuple[Tensor, int, str, int, int, 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:
|
|
Transcript
|
|
int:
|
|
Speaker ID
|
|
int:
|
|
Chapter ID
|
|
int:
|
|
Utterance ID
|
|
list:
|
|
List of biasing words in the utterance
|
|
"""
|
|
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)
|