# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional, Tuple import torch.utils.data from nemo.collections.asr.parts.preprocessing.segment import available_formats as valid_sf_formats from nemo.collections.asr.parts.utils.manifest_utils import read_manifest from nemo.collections.tts.parts.preprocessing.feature_processors import FeatureProcessor from nemo.collections.tts.parts.utils.tts_dataset_utils import ( filter_dataset_by_duration, get_weighted_sampler, load_audio, resample_batch, sample_audio, stack_tensors, ) from nemo.core.classes import Dataset from nemo.utils import logging VALID_FILE_FORMATS = ';'.join(['wav', 'mp3', 'flac', 'opus'] + [fmt.lower() for fmt in valid_sf_formats.keys()]) @dataclass class DatasetMeta: manifest_path: Path audio_dir: Path sample_weight: float = 1.0 audio_tar_filepaths: Optional[List[str]] = None @dataclass class DatasetSample: dataset_name: str manifest_entry: dict audio_dir: Path def audio_collate_fn(batch: List[dict], resample_rates: Optional[Tuple[int]] = None): dataset_name_list = [] audio_filepath_list = [] audio_list = [] audio_len_list = [] for example in batch: dataset_name_list.append(example["dataset_name"]) audio_filepath_list.append(example["audio_filepath"]) audio_list.append(example["audio"]) audio_len_list.append(example["audio_len"]) batch_audio_len = torch.IntTensor(audio_len_list) audio_max_len = int(batch_audio_len.max().item()) batch_audio = stack_tensors(audio_list, max_lens=[audio_max_len]) if resample_rates: batch_audio, batch_audio_len = resample_batch( audio=batch_audio, audio_len=batch_audio_len, input_sample_rate=resample_rates[0], output_sample_rate=resample_rates[1], ) batch_dict = { "dataset_names": dataset_name_list, "audio_filepaths": audio_filepath_list, "audio": batch_audio, "audio_lens": batch_audio_len, } return batch_dict def preprocess_manifest( dataset_name: str, dataset: DatasetMeta, min_duration: float, max_duration: float, ): entries = read_manifest(dataset.manifest_path) filtered_entries, total_hours, filtered_hours = filter_dataset_by_duration( entries=entries, min_duration=min_duration, max_duration=max_duration ) logging.info(dataset_name) logging.info(f"Original # of files: {len(entries)}") logging.info(f"Filtered # of files: {len(filtered_entries)}") logging.info(f"Original duration: {total_hours:.2f} hours") logging.info(f"Filtered duration: {filtered_hours:.2f} hours") samples = [] sample_weights = [] for entry in filtered_entries: sample = DatasetSample(dataset_name=dataset_name, manifest_entry=entry, audio_dir=Path(dataset.audio_dir)) samples.append(sample) sample_weights.append(dataset.sample_weight) return samples, sample_weights class VocoderDataset(Dataset): """ Class for processing and loading Vocoder training examples. Args: dataset_meta: Dict of dataset names (string) to dataset metadata. sample_rate: Sample rate to load audio as. If the audio is stored at a different sample rate, then it will be resampled using librosa. resample_rate: Optional sample rate to resample to, using torch-based resampling. n_samples: Optional int, if provided then n_samples samples will be randomly sampled from the full audio file. weighted_sampling_steps_per_epoch: Optional int, If provided, then data will be sampled (with replacement) based on the sample weights provided in the dataset metadata. If None, then sample weights will be ignored. feature_processors: Optional, list of feature processors to run on training examples. min_duration: Optional float, if provided audio files in the training manifest shorter than 'min_duration' will be ignored. max_duration: Optional float, if provided audio files in the training manifest longer than 'max_duration' will be ignored. trunc_duration: Optional int, if provided audio will be truncated to at most 'trunc_duration' seconds. volume_norm: Whether to apply volume normalization to loaded audio. """ def __init__( self, dataset_meta: Dict, sample_rate: int, resample_rate: Optional[int] = None, n_samples: Optional[int] = None, weighted_sampling_steps_per_epoch: Optional[int] = None, feature_processors: Optional[Dict[str, FeatureProcessor]] = None, min_duration: Optional[float] = None, max_duration: Optional[float] = None, trunc_duration: Optional[float] = None, volume_norm: bool = False, ): super().__init__() self.sample_rate = sample_rate if resample_rate and self.sample_rate != resample_rate: self.resample_rates = [sample_rate, resample_rate] else: self.resample_rates = None self.n_samples = n_samples self.trunc_duration = trunc_duration self.volume_norm = volume_norm self.weighted_sampling_steps_per_epoch = weighted_sampling_steps_per_epoch self.load_precomputed_mel = False if feature_processors: logging.info(f"Found feature processors {feature_processors.keys()}") self.feature_processors = list(feature_processors.values()) else: self.feature_processors = [] self.data_samples = [] self.sample_weights = [] for dataset_name, dataset_info in dataset_meta.items(): dataset = DatasetMeta(**dataset_info) samples, weights = preprocess_manifest( dataset_name=dataset_name, dataset=dataset, min_duration=min_duration, max_duration=max_duration, ) self.data_samples += samples self.sample_weights += weights def get_sampler(self, batch_size: int, world_size: int) -> Optional[torch.utils.data.Sampler]: if not self.weighted_sampling_steps_per_epoch: return None sampler = get_weighted_sampler( sample_weights=self.sample_weights, batch_size=batch_size, world_size=world_size, num_steps=self.weighted_sampling_steps_per_epoch, ) return sampler def __len__(self): return len(self.data_samples) def __getitem__(self, index): data = self.data_samples[index] if self.n_samples: audio_array, _, audio_filepath_rel = sample_audio( manifest_entry=data.manifest_entry, audio_dir=data.audio_dir, sample_rate=self.sample_rate, n_samples=self.n_samples, volume_norm=self.volume_norm, ) else: audio_array, _, audio_filepath_rel = load_audio( manifest_entry=data.manifest_entry, audio_dir=data.audio_dir, sample_rate=self.sample_rate, max_duration=self.trunc_duration, volume_norm=self.volume_norm, ) audio = torch.tensor(audio_array, dtype=torch.float32) audio_len = audio.shape[0] example = { "dataset_name": data.dataset_name, "audio_filepath": audio_filepath_rel, "audio": audio, "audio_len": audio_len, } for processor in self.feature_processors: processor.process(example) return example def collate_fn(self, batch): return audio_collate_fn(batch, resample_rates=self.resample_rates)