| import json |
| import logging |
| import os |
| import shutil |
| from pathlib import Path |
|
|
| import torchaudio |
| from speechbrain.utils.data_utils import download_file |
| from tqdm import tqdm |
|
|
| from hw_asr.base.base_dataset import BaseDataset |
| from hw_asr.utils import ROOT_PATH |
|
|
| logger = logging.getLogger(__name__) |
|
|
| URL_LINKS = { |
| "dev-clean": "https://www.openslr.org/resources/12/dev-clean.tar.gz", |
| "dev-other": "https://www.openslr.org/resources/12/dev-other.tar.gz", |
| "test-clean": "https://www.openslr.org/resources/12/test-clean.tar.gz", |
| "test-other": "https://www.openslr.org/resources/12/test-other.tar.gz", |
| "train-clean-100": "https://www.openslr.org/resources/12/train-clean-100.tar.gz", |
| "train-clean-360": "https://www.openslr.org/resources/12/train-clean-360.tar.gz", |
| "train-other-500": "https://www.openslr.org/resources/12/train-other-500.tar.gz", |
| } |
|
|
|
|
| class LibrispeechDataset(BaseDataset): |
| def __init__(self, part, data_dir=None, *args, **kwargs): |
| assert part in URL_LINKS or part == 'train_all' |
|
|
| if data_dir is None: |
| data_dir = ROOT_PATH / "data" / "datasets" / "librispeech" |
| data_dir.mkdir(exist_ok=True, parents=True) |
| self._data_dir = data_dir |
| if part == 'train_all': |
| index = sum([self._get_or_load_index(part) |
| for part in URL_LINKS if 'train' in part], []) |
| else: |
| index = self._get_or_load_index(part) |
|
|
| super().__init__(index, *args, **kwargs) |
|
|
| def _load_part(self, part): |
| arch_path = self._data_dir / f"{part}.tar.gz" |
| print(f"Loading part {part}") |
| download_file(URL_LINKS[part], arch_path) |
| shutil.unpack_archive(arch_path, self._data_dir) |
| for fpath in (self._data_dir / "LibriSpeech").iterdir(): |
| shutil.move(str(fpath), str(self._data_dir / fpath.name)) |
| os.remove(str(arch_path)) |
| shutil.rmtree(str(self._data_dir / "LibriSpeech")) |
|
|
| def _get_or_load_index(self, part): |
| index_path = self._data_dir / f"{part}_index.json" |
| if index_path.exists(): |
| with index_path.open() as f: |
| index = json.load(f) |
| else: |
| index = self._create_index(part) |
| with index_path.open("w") as f: |
| json.dump(index, f, indent=2) |
| return index |
|
|
| def _create_index(self, part): |
| index = [] |
| split_dir = self._data_dir / part |
| if not split_dir.exists(): |
| self._load_part(part) |
|
|
| flac_dirs = set() |
| for dirpath, dirnames, filenames in os.walk(str(split_dir)): |
| if any([f.endswith(".flac") for f in filenames]): |
| flac_dirs.add(dirpath) |
| for flac_dir in tqdm( |
| list(flac_dirs), desc=f"Preparing librispeech folders: {part}" |
| ): |
| flac_dir = Path(flac_dir) |
| trans_path = list(flac_dir.glob("*.trans.txt"))[0] |
| with trans_path.open() as f: |
| for line in f: |
| f_id = line.split()[0] |
| f_text = " ".join(line.split()[1:]).strip() |
| flac_path = flac_dir / f"{f_id}.flac" |
| t_info = torchaudio.info(str(flac_path)) |
| length = t_info.num_frames / t_info.sample_rate |
| index.append( |
| { |
| "path": str(flac_path.absolute().resolve()), |
| "text": f_text.lower(), |
| "audio_len": length, |
| } |
| ) |
| return index |
|
|
