Spaces:
Running on Zero
Running on Zero
| # -*- coding: utf-8 -*- | |
| # Copyright 2019 Tomoki Hayashi | |
| # MIT License (https://opensource.org/licenses/MIT) | |
| # Portions of this code are a modification of the ParallelWaveGAN project. | |
| # The original source code can be found at: | |
| # https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/datasets/scp_dataset.py | |
| # | |
| # Copyright (c) 2025 Yiwei Guo | |
| # Licensed under the MIT license. | |
| """Dataset modules based on kaldi-style scp files.""" | |
| import torch | |
| import logging | |
| import random | |
| import copy | |
| import sox | |
| from multiprocessing import Manager | |
| import kaldiio | |
| import librosa | |
| import numpy as np | |
| from torch.utils.data import Dataset | |
| from tqdm import tqdm | |
| from lscodec.utils import HDF5ScpLoader | |
| from lscodec.utils import NpyScpLoader | |
| from lscodec.datasets.perturbs import Perturbs | |
| logging.getLogger("sox").setLevel(logging.ERROR) | |
| def _get_feats_scp_loader(feats_scp): | |
| # read the first line of feats.scp file | |
| with open(feats_scp) as f: | |
| key, value = f.readlines()[0].replace("\n", "").split() | |
| # check scp type | |
| if ":" in value: | |
| value_1, value_2 = value.split(":") | |
| if value_1.endswith(".ark"): | |
| # kaldi-ark case: utt_id_1 /path/to/utt_id_1.ark:index | |
| return kaldiio.load_scp(feats_scp) | |
| elif value_1.endswith(".h5"): | |
| # hdf5 case with path in hdf5: utt_id_1 /path/to/utt_id_1.h5:feats | |
| return HDF5ScpLoader(feats_scp) | |
| else: | |
| raise ValueError("Not supported feats.scp type.") | |
| else: | |
| if value.endswith(".h5"): | |
| # hdf5 case without path in hdf5: utt_id_1 /path/to/utt_id_1.h5 | |
| return HDF5ScpLoader(feats_scp) | |
| elif value.endswith(".npy"): | |
| # npy case: utt_id_1 /path/to/utt_id_1.npy | |
| return NpyScpLoader(feats_scp) | |
| else: | |
| raise ValueError("Not supported feats.scp type.") | |
| class AudioMelSCPDataset(Dataset): | |
| """PyTorch compatible audio and feat dataset based on kaldi-stype scp files.""" | |
| def __init__( | |
| self, | |
| wav_scp, | |
| prompt_scp, | |
| target_scp, | |
| aux_scp, | |
| idx_scp, | |
| utt2num_frames=None, | |
| segments=None, | |
| batch_frames=None, | |
| batch_size=None, | |
| min_num_frames=None, | |
| max_num_frames=None, | |
| return_utt_id=False, | |
| return_sampling_rate=False, | |
| allow_cache=False, | |
| length_tolerance=2, | |
| prompt_min_length=500, # in terms of frames | |
| extractor_fs=16000, | |
| # speed_alpha_range=(0.8, 1.2), | |
| perturber=None, | |
| utt2spk=None, | |
| spk2gender=None | |
| ): | |
| """Initialize dataset. | |
| Args: | |
| wav_scp (str): Kaldi-style wav.scp file. | |
| prompt_scp (str): Kaldi-style fests.scp file. | |
| aux_scp (str): Kaldi-style fests.scp file. | |
| segments (str): Kaldi-style segments file. | |
| min_num_frames (int): Threshold to remove short feature files. | |
| max_num_frames (int): Threshold to remove long feature files. | |
| return_utt_id (bool): Whether to return utterance id. | |
| return_sampling_rate (bool): Whether to return sampling rate. | |
| allow_cache (bool): Whether to allow cache of the loaded files. | |
| """ | |
| # load scp as lazy dict | |
| self.audio_loader = kaldiio.load_scp(wav_scp, segments=segments) | |
| self.prompt_loader = _get_feats_scp_loader(prompt_scp) | |
| # with open(prompt_scp, 'r') as fr: | |
| # self.prompt_loader = dict() | |
| # for line in fr.readlines(): | |
| # terms = line.strip().split() | |
| # utt, path = terms | |
| # self.prompt_loader[utt] = path | |
| self.target_loader = _get_feats_scp_loader(target_scp) | |
| self.aux_loader = _get_feats_scp_loader(aux_scp) | |
| self.idx_loader = _get_feats_scp_loader(idx_scp) | |
| self.utt_ids = list(self.prompt_loader.keys()) | |
| self.return_utt_id = return_utt_id | |
| self.return_sampling_rate = return_sampling_rate | |
| self.allow_cache = allow_cache | |
| self.extractor_fs = extractor_fs | |
| utt2num_frames_loader = None | |
| if utt2num_frames is not None: | |
| with open(utt2num_frames, 'r') as f: | |
| utt2num_frames_loader = dict([(x.split()[0], int(x.split()[1])) for x in f.readlines()]) | |
| else: | |
| utt2num_frames_loader = dict([(k, mel.shape[0]) for k, mel in self.prompt_loader.items()]) | |
| self.utt2num_frames_loader = utt2num_frames_loader | |
| # filter by threshold | |
| if (min_num_frames or max_num_frames) is not None: | |
| mel_lengths = [utt2num_frames_loader[key] for key in self.utt_ids] | |
| idxs = [ | |
| idx | |
| for idx in range(len(self.utt_ids)) | |
| if (min_num_frames and mel_lengths[idx] >= min_num_frames) and (max_num_frames and mel_lengths[idx] <= max_num_frames) | |
| ] | |
| if len(self.utt_ids) != len(idxs): | |
| logging.warning( | |
| f"Some files are filtered by mel length threshold " | |
| f"({len(self.utt_ids)} -> {len(idxs)})." | |
| ) | |
| self.utt_ids = [self.utt_ids[idx] for idx in idxs] | |
| # batchify | |
| if batch_frames is not None: | |
| self.batches = self.batchify(utt2num_frames_loader, batch_frames=batch_frames) | |
| elif batch_size is not None: | |
| self.batches = self.batchify(utt2num_frames_loader, batch_size=batch_size) | |
| else: | |
| self.batches = [[utt_id] for utt_id in self.utt_ids] | |
| if allow_cache: | |
| # NOTE(kan-bayashi): Manager is need to share memory in dataloader with num_workers > 0 | |
| self.manager = Manager() | |
| self.caches = self.manager.dict() | |
| self.length_tolerance = length_tolerance | |
| self.prompt_min_length = prompt_min_length | |
| # assert len(speed_alpha_range)==2, speed_alpha_range | |
| # self.alpha_range = speed_alpha_range | |
| self.perturber = perturber | |
| if utt2spk is not None: | |
| self.utt2spk = dict() | |
| with open(utt2spk, 'r') as f: | |
| for line in f.readlines(): | |
| terms = line.strip().split() | |
| self.utt2spk[terms[0]] = terms[1] | |
| assert spk2gender is not None | |
| self.spk2gender = dict() | |
| with open(spk2gender, 'r') as f: | |
| for line in f.readlines(): | |
| terms = line.strip().split() | |
| assert terms[1] in {"F", "M"} | |
| self.spk2gender[terms[0]] = terms[1] | |
| else: | |
| self.utt2spk = None | |
| def batchify(self, utt2num_frames_loader, batch_frames=None, batch_size=None, min_batch_size=1, drop_last=True): | |
| assert batch_size is None or batch_size > min_batch_size | |
| batches = [] | |
| batch = [] | |
| accum_num_frames = 0 | |
| utt_ids_set = set(self.utt_ids) | |
| for utt_id, mel_length in tqdm(sorted(list(utt2num_frames_loader.items()), key=lambda x: x[1], reverse=True)): | |
| if utt_id not in utt_ids_set: | |
| continue | |
| if (batch_frames is not None and accum_num_frames + mel_length > batch_frames and len(batch) > min_batch_size) or (batch_size is not None and len(batch) == batch_size): | |
| batches.append(batch) | |
| batch = [] | |
| accum_num_frames = 0 | |
| batch.append(utt_id) | |
| accum_num_frames += mel_length | |
| if len(batch) > min_batch_size and not drop_last: | |
| batches.append(batch) | |
| return batches | |
| def __getitem__(self, idx): | |
| """Get specified idx items. | |
| Args: | |
| idx (int): Index of the item. | |
| Returns: | |
| str: Utterance id (only in return_utt_id = True). | |
| ndarray or tuple: Audio signal (T,) or (w/ sampling rate if return_sampling_rate = True). | |
| ndarray: Feature (T', C). | |
| """ | |
| batch = self.batches[idx] | |
| batch_items = [] | |
| for utt_id in batch: | |
| if self.allow_cache and self.caches.get(utt_id) is not None: | |
| items = self.caches[utt_id] | |
| else: | |
| fs, audio = self.audio_loader[utt_id] | |
| prompt = self.prompt_loader[utt_id] | |
| # prompt = torch.load(prompt_path).numpy() | |
| target = self.target_loader[utt_id] | |
| ssl_idx = self.idx_loader[utt_id] | |
| aux = self.aux_loader[utt_id] | |
| min_len = min(len(prompt), len(aux), len(target), len(ssl_idx)) | |
| assert ((abs(len(prompt) - min_len) <= self.length_tolerance) and | |
| (abs(len(target) - min_len) <= self.length_tolerance) and | |
| (abs(len(ssl_idx) - min_len) <= self.length_tolerance) and | |
| (abs(len(aux) - min_len) <= self.length_tolerance)), \ | |
| f"Audio feature lengths difference exceeds length tolerance for {utt_id}" | |
| prompt, aux, target, ssl_idx = prompt[:min_len], aux[:min_len], target[:min_len], ssl_idx[:min_len] | |
| if self.utt2spk is not None: | |
| spk=self.utt2spk[utt_id] | |
| gender=self.spk2gender[spk] | |
| else: | |
| gender = "M" # Default to male, so that not use pitch shift | |
| # normalize audio signal to be [-1, 1] | |
| audio = audio.astype(np.float32) | |
| audio /= 1 << (16 - 1) # assume that wav is PCM 16 bit | |
| if fs != self.extractor_fs: | |
| audio_for_extractor = librosa.resample(audio, orig_sr=fs, target_sr=self.extractor_fs) | |
| else: | |
| audio_for_extractor = audio | |
| # # perturb | |
| # if self.alpha_range[0] == self.alpha_range[1] == 1.0: | |
| # pass | |
| # else: | |
| # speed_alpha = random.random() * (self.alpha_range[1] - self.alpha_range[0]) + self.alpha_range[0] | |
| # self.transformer = sox.Transformer() | |
| # self.transformer.set_globals(verbosity=1) # disable warnings. | |
| # self.transformer.speed(factor=speed_alpha) | |
| # self.transformer.tempo(factor=1/speed_alpha) # recover the original length. Only affects timbre and pitch. | |
| # audio_for_extractor = self.transformer.build_array(input_array=audio_for_extractor, sample_rate_in=self.extractor_fs) | |
| audio_for_extractor = self.perturber(audio_for_extractor, self.extractor_fs, gender) | |
| if self.return_utt_id: | |
| items = utt_id, audio_for_extractor, prompt, aux, target, ssl_idx | |
| else: | |
| items = audio_for_extractor, prompt, aux, target, ssl_idx | |
| if self.allow_cache: | |
| self.caches[utt_id] = items | |
| batch_items.append(items) | |
| return batch_items | |
| def __len__(self): | |
| """Return dataset length. | |
| Returns: | |
| int: The length of dataset. | |
| """ | |
| return len(self.batches) | |
| class MelSCPDataset(Dataset): | |
| """PyTorch compatible feat dataset based on kaldi-stype scp files.""" | |
| def __init__( | |
| self, | |
| vqidx_scp, | |
| prompt_scp, | |
| utt2num_frames=None, | |
| min_num_frames=None, | |
| max_num_frames=None, | |
| return_utt_id=False, | |
| allow_cache=False, | |
| ): | |
| """Initialize dataset. | |
| Args: | |
| vqidx_scp (str): Kaldi-style fests.scp file. | |
| prompt_scp (str): Kaldi-style scp file. In this file, every utt is associated with its prompt's mel-spectrogram. | |
| min_num_frames (int): Threshold to remove short feature files. | |
| max_num_frames (int): Threshold to remove long feature files. | |
| return_utt_id (bool): Whether to return utterance id. | |
| allow_cache (bool): Whether to allow cache of the loaded files. | |
| """ | |
| # load scp as lazy dict | |
| vqidx_loader = _get_feats_scp_loader(vqidx_scp) | |
| prompt_loader = _get_feats_scp_loader(prompt_scp) | |
| vqidx_keys = list(set(prompt_loader.keys()) & set(vqidx_loader.keys())) | |
| utt2num_frames_loader = None | |
| if utt2num_frames is not None: | |
| with open(utt2num_frames, 'r') as f: | |
| utt2num_frames_loader = dict([(x.split()[0], int(x.split()[1])) for x in f.readlines()]) | |
| else: | |
| utt2num_frames_loader = dict([(k, mel.shape[0]) for k, mel in vqidx_loader.items()]) | |
| # filter by threshold | |
| if (min_num_frames or max_num_frames) is not None: | |
| mel_lengths = [utt2num_frames_loader[key] for key in vqidx_keys] | |
| idxs = [ | |
| idx | |
| for idx in range(len(vqidx_keys)) | |
| if (min_num_frames and mel_lengths[idx] > min_num_frames) and (max_num_frames and mel_lengths[idx] < max_num_frames) | |
| ] | |
| if len(vqidx_keys) != len(idxs): | |
| logging.warning( | |
| f"Some files are filtered by mel length threshold " | |
| f"({len(vqidx_keys)} -> {len(idxs)})." | |
| ) | |
| vqidx_keys = [vqidx_keys[idx] for idx in idxs] | |
| self.vqidx_loader = vqidx_loader | |
| self.prompt_loader = prompt_loader | |
| self.utt_ids = vqidx_keys | |
| self.return_utt_id = return_utt_id | |
| self.allow_cache = allow_cache | |
| if allow_cache: | |
| # NOTE(kan-bayashi): Manager is need to share memory in dataloader with num_workers > 0 | |
| self.manager = Manager() | |
| self.caches = self.manager.list() | |
| self.caches += [() for _ in range(len(self.utt_ids))] | |
| def __getitem__(self, idx): | |
| """Get specified idx items. | |
| Args: | |
| idx (int): Index of the item. | |
| Returns: | |
| str: Utterance id (only in return_utt_id = True). | |
| ndarray: Feature (T', C). | |
| """ | |
| if self.allow_cache and len(self.caches[idx]) != 0: | |
| return self.caches[idx] | |
| utt_id = self.utt_ids[idx] | |
| vqidx = self.vqidx_loader[utt_id].astype(int) | |
| prompt = self.prompt_loader[utt_id].copy() | |
| if self.return_utt_id: | |
| items = utt_id, vqidx, prompt | |
| else: | |
| items = vqidx, prompt | |
| if self.allow_cache: | |
| self.caches[idx] = items | |
| return items | |
| def __len__(self): | |
| """Return dataset length. | |
| Returns: | |
| int: The length of dataset. | |
| """ | |
| return len(self.utt_ids) | |