wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
1171 lines
47 KiB
Python
1171 lines
47 KiB
Python
# Copyright (c) 2022, NVIDIA CORPORATION. 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.
|
|
|
|
import abc
|
|
import math
|
|
import random
|
|
from collections import OrderedDict, namedtuple
|
|
from dataclasses import dataclass
|
|
from typing import Callable, Dict, List, Optional, Tuple, Type, Union
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import torch
|
|
|
|
from nemo.collections.asr.parts.preprocessing.segment import AudioSegment, ChannelSelectorType
|
|
from nemo.collections.common.parts.preprocessing import collections
|
|
from nemo.collections.common.parts.utils import flatten
|
|
from nemo.core.classes import Dataset
|
|
from nemo.core.neural_types import AudioSignal, EncodedRepresentation, LengthsType, NeuralType
|
|
from nemo.utils import logging
|
|
|
|
__all__ = [
|
|
'AudioToTargetDataset',
|
|
'AudioToTargetWithReferenceDataset',
|
|
'AudioToTargetWithEmbeddingDataset',
|
|
]
|
|
|
|
|
|
def _audio_collate_fn(batch: List[dict]) -> Tuple[torch.Tensor]:
|
|
"""Collate a batch of items returned by __getitem__.
|
|
Examples for each signal are zero padded to the same length
|
|
(batch_length), which is determined by the longest example.
|
|
Lengths of the original signals are returned in the output.
|
|
|
|
Args:
|
|
batch: List of dictionaries. Each element of the list
|
|
has the following format
|
|
```
|
|
{
|
|
'signal_0': 1D or 2D tensor,
|
|
'signal_1': 1D or 2D tensor,
|
|
...
|
|
'signal_N': 1D or 2D tensor,
|
|
}
|
|
```
|
|
1D tensors have shape (num_samples,) and 2D tensors
|
|
have shape (num_channels, num_samples)
|
|
|
|
Returns:
|
|
A tuple containing signal tensor and signal length tensor (in samples)
|
|
for each signal.
|
|
The output has the following format:
|
|
```
|
|
(signal_0, signal_0_length, signal_1, signal_1_length, ..., signal_N, signal_N_length)
|
|
```
|
|
Note that the output format is obtained by interleaving signals and their length.
|
|
"""
|
|
signals = batch[0].keys()
|
|
|
|
batched = tuple()
|
|
|
|
for signal in signals:
|
|
signal_length = [b[signal].shape[-1] for b in batch]
|
|
# Batch length is determined by the longest signal in the batch
|
|
batch_length = max(signal_length)
|
|
b_signal = []
|
|
for s_len, b in zip(signal_length, batch):
|
|
# check if padding is necessary
|
|
if s_len < batch_length:
|
|
if b[signal].ndim == 1:
|
|
# single-channel signal
|
|
pad = (0, batch_length - s_len)
|
|
elif b[signal].ndim == 2:
|
|
# multi-channel signal
|
|
pad = (0, batch_length - s_len, 0, 0)
|
|
else:
|
|
raise RuntimeError(
|
|
f'Signal {signal} has unsuported dimensions {signal.shape}. Currently, only 1D and 2D arrays are supported.'
|
|
)
|
|
b[signal] = torch.nn.functional.pad(b[signal], pad)
|
|
# append the current padded signal
|
|
b_signal.append(b[signal])
|
|
# (signal_batched, signal_length)
|
|
batched += (torch.stack(b_signal), torch.tensor(signal_length, dtype=torch.int32))
|
|
|
|
# Currently, outputs are expected to be in a tuple, where each element must correspond
|
|
# to the output type in the OrderedDict returned by output_types.
|
|
#
|
|
# Therefore, we return batched signals by interleaving signals and their length:
|
|
# (signal_0, signal_0_length, signal_1, signal_1_length, ...)
|
|
return batched
|
|
|
|
|
|
@dataclass
|
|
class SignalSetup:
|
|
signals: List[str] # signal names
|
|
duration: Optional[Union[float, list]] = None # duration for each signal
|
|
channel_selectors: Optional[List[ChannelSelectorType]] = None # channel selector for loading each signal
|
|
|
|
|
|
class ASRAudioProcessor:
|
|
"""Class that processes an example from Audio collection and returns
|
|
a dictionary with prepared signals.
|
|
|
|
For example, the output dictionary may be the following
|
|
```
|
|
{
|
|
'input_signal': input_signal_tensor,
|
|
'target_signal': target_signal_tensor,
|
|
'reference_signal': reference_signal_tensor,
|
|
'embedding_vector': embedding_vector
|
|
}
|
|
```
|
|
Keys in the output dictionary are ordered with synchronous signals given first,
|
|
followed by asynchronous signals and embedding.
|
|
|
|
Args:
|
|
sample_rate: sample rate used for all audio signals
|
|
random_offset: If `True`, offset will be randomized when loading a subsegment
|
|
from a file.
|
|
normalization_signal: Normalize all audio with a factor that ensures the signal
|
|
`example[normalization_signal]` in `process` is in range [-1, 1].
|
|
All other audio signals are scaled by the same factor. Default is
|
|
`None`, corresponding to no normalization.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sample_rate: float,
|
|
random_offset: bool,
|
|
normalization_signal: Optional[str] = None,
|
|
eps: float = 1e-8,
|
|
):
|
|
self.sample_rate = sample_rate
|
|
self.random_offset = random_offset
|
|
self.normalization_signal = normalization_signal
|
|
self.eps = eps
|
|
|
|
self.sync_setup = None
|
|
self.async_setup = None
|
|
self.embedding_setup = None
|
|
|
|
@property
|
|
def sample_rate(self) -> float:
|
|
return self._sample_rate
|
|
|
|
@sample_rate.setter
|
|
def sample_rate(self, value: float):
|
|
if value <= 0:
|
|
raise ValueError(f'Sample rate must be positive, received {value}')
|
|
|
|
self._sample_rate = value
|
|
|
|
@property
|
|
def random_offset(self) -> bool:
|
|
return self._random_offset
|
|
|
|
@random_offset.setter
|
|
def random_offset(self, value: bool):
|
|
self._random_offset = value
|
|
|
|
@property
|
|
def sync_setup(self) -> SignalSetup:
|
|
"""Return the current setup for synchronous signals.
|
|
|
|
Returns:
|
|
A dataclass containing the list of signals, their
|
|
duration and channel selectors.
|
|
"""
|
|
return self._sync_setup
|
|
|
|
@sync_setup.setter
|
|
def sync_setup(self, value: Optional[SignalSetup]):
|
|
"""Setup signals to be loaded synchronously.
|
|
|
|
Args:
|
|
value: An instance of SignalSetup with the following fields
|
|
- signals: list of signals (keys of example.audio_signals) which will be loaded
|
|
synchronously with the same start time and duration.
|
|
- duration: Duration for each signal to be loaded.
|
|
If duration is set to None, the whole file will be loaded.
|
|
- channel_selectors: A list of channel selector for each signal. If channel selector
|
|
is None, all channels in the audio file will be loaded.
|
|
"""
|
|
if value is None or isinstance(value, SignalSetup):
|
|
self._sync_setup = value
|
|
else:
|
|
raise ValueError(f'Unexpected type {type(value)} for value {value}.')
|
|
|
|
@property
|
|
def async_setup(self) -> SignalSetup:
|
|
"""Return the current setup for asynchronous signals.
|
|
|
|
Returns:
|
|
A dataclass containing the list of signals, their
|
|
duration and channel selectors.
|
|
"""
|
|
return self._async_setup
|
|
|
|
@async_setup.setter
|
|
def async_setup(self, value: Optional[SignalSetup]):
|
|
"""Setup signals to be loaded asynchronously.
|
|
|
|
Args:
|
|
Args:
|
|
value: An instance of SignalSetup with the following fields
|
|
- signals: list of signals (keys of example.audio_signals) which will be loaded
|
|
asynchronously with signals possibly having different start and duration
|
|
- duration: Duration for each signal to be loaded.
|
|
If duration is set to None, the whole file will be loaded.
|
|
- channel_selectors: A list of channel selector for each signal. If channel selector
|
|
is None, all channels in the audio file will be loaded.
|
|
"""
|
|
if value is None or isinstance(value, SignalSetup):
|
|
self._async_setup = value
|
|
else:
|
|
raise ValueError(f'Unexpected type {type(value)} for value {value}.')
|
|
|
|
@property
|
|
def embedding_setup(self) -> SignalSetup:
|
|
"""Setup signals corresponding to an embedding vector."""
|
|
return self._embedding_setup
|
|
|
|
@embedding_setup.setter
|
|
def embedding_setup(self, value: SignalSetup):
|
|
"""Setup signals corresponding to an embedding vector.
|
|
|
|
Args:
|
|
value: An instance of SignalSetup with the following fields
|
|
- signals: list of signals (keys of example.audio_signals) which will be loaded
|
|
as embedding vectors.
|
|
"""
|
|
if value is None or isinstance(value, SignalSetup):
|
|
self._embedding_setup = value
|
|
else:
|
|
raise ValueError(f'Unexpected type {type(value)} for value {value}.')
|
|
|
|
def process(self, example: collections.Audio.OUTPUT_TYPE) -> Dict[str, torch.Tensor]:
|
|
"""Process an example from a collection of audio examples.
|
|
|
|
Args:
|
|
example: an example from Audio collection.
|
|
|
|
Returns:
|
|
An ordered dictionary of signals and their tensors.
|
|
For example, the output dictionary may be the following
|
|
```
|
|
{
|
|
'input_signal': input_signal_tensor,
|
|
'target_signal': target_signal_tensor,
|
|
'reference_signal': reference_signal_tensor,
|
|
'embedding_vector': embedding_vector
|
|
}
|
|
```
|
|
Keys in the output dictionary are ordered with synchronous signals given first,
|
|
followed by asynchronous signals and embedding.
|
|
"""
|
|
audio = self.load_audio(example=example)
|
|
audio = self.process_audio(audio=audio)
|
|
return audio
|
|
|
|
def load_audio(self, example: collections.Audio.OUTPUT_TYPE) -> Dict[str, torch.Tensor]:
|
|
"""Given an example, load audio from `example.audio_files` and prepare
|
|
the output dictionary.
|
|
|
|
Args:
|
|
example: An example from an audio collection
|
|
|
|
Returns:
|
|
An ordered dictionary of signals and their tensors.
|
|
For example, the output dictionary may be the following
|
|
```
|
|
{
|
|
'input_signal': input_signal_tensor,
|
|
'target_signal': target_signal_tensor,
|
|
'reference_signal': reference_signal_tensor,
|
|
'embedding_vector': embedding_vector
|
|
}
|
|
```
|
|
Keys in the output dictionary are ordered with synchronous signals given first,
|
|
followed by asynchronous signals and embedding.
|
|
"""
|
|
output = OrderedDict()
|
|
|
|
if self.sync_setup is not None:
|
|
# Load all signals with the same start and duration
|
|
sync_signals = self.load_sync_signals(example)
|
|
output.update(sync_signals)
|
|
|
|
if self.async_setup is not None:
|
|
# Load each signal independently
|
|
async_signals = self.load_async_signals(example)
|
|
output.update(async_signals)
|
|
|
|
# Load embedding vector
|
|
if self.embedding_setup is not None:
|
|
embedding = self.load_embedding(example)
|
|
output.update(embedding)
|
|
|
|
if not output:
|
|
raise RuntimeError('Output dictionary is empty. Please use `_setup` methods to setup signals to be loaded')
|
|
|
|
return output
|
|
|
|
def process_audio(self, audio: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
"""Process audio signals available in the input dictionary.
|
|
|
|
Args:
|
|
audio: A dictionary containing loaded signals `signal: tensor`
|
|
|
|
Returns:
|
|
An ordered dictionary of signals and their tensors.
|
|
"""
|
|
if self.normalization_signal:
|
|
# Normalize all audio with a factor that ensures the normalization signal is in range [-1, 1].
|
|
norm_scale = audio[self.normalization_signal].abs().max()
|
|
|
|
# Do not normalize embeddings
|
|
skip_signals = self.embedding_setup.signals if self.embedding_setup is not None else []
|
|
|
|
# Normalize audio signals
|
|
for signal in audio:
|
|
if signal not in skip_signals:
|
|
# All audio signals are scaled by the same factor.
|
|
# This ensures that the relative level between signals is preserved.
|
|
audio[signal] = audio[signal] / (norm_scale + self.eps)
|
|
|
|
return audio
|
|
|
|
def load_sync_signals(self, example: collections.Audio.OUTPUT_TYPE) -> Dict[str, torch.Tensor]:
|
|
"""Load signals with the same start and duration.
|
|
|
|
Args:
|
|
example: an example from audio collection
|
|
|
|
Returns:
|
|
An ordered dictionary of signals and their tensors.
|
|
"""
|
|
output = OrderedDict()
|
|
sync_audio_files = [example.audio_files[s] for s in self.sync_setup.signals]
|
|
|
|
sync_samples = self.get_samples_synchronized(
|
|
audio_files=sync_audio_files,
|
|
channel_selectors=self.sync_setup.channel_selectors,
|
|
sample_rate=self.sample_rate,
|
|
duration=self.sync_setup.duration,
|
|
fixed_offset=example.offset,
|
|
random_offset=self.random_offset,
|
|
)
|
|
|
|
for signal, samples in zip(self.sync_setup.signals, sync_samples):
|
|
output[signal] = torch.tensor(samples)
|
|
|
|
return output
|
|
|
|
def load_async_signals(self, example: collections.Audio.OUTPUT_TYPE) -> Dict[str, torch.Tensor]:
|
|
"""Load each async signal independently, no constraints on starting
|
|
from the same time.
|
|
|
|
Args:
|
|
example: an example from audio collection
|
|
|
|
Returns:
|
|
An ordered dictionary of signals and their tensors.
|
|
"""
|
|
output = OrderedDict()
|
|
for idx, signal in enumerate(self.async_setup.signals):
|
|
samples = self.get_samples(
|
|
audio_file=example.audio_files[signal],
|
|
sample_rate=self.sample_rate,
|
|
duration=self.async_setup.duration[idx],
|
|
channel_selector=self.async_setup.channel_selectors[idx],
|
|
fixed_offset=example.offset,
|
|
random_offset=self.random_offset,
|
|
)
|
|
output[signal] = torch.tensor(samples)
|
|
return output
|
|
|
|
@classmethod
|
|
def get_samples(
|
|
cls,
|
|
audio_file: str,
|
|
sample_rate: int,
|
|
duration: Optional[float] = None,
|
|
channel_selector: ChannelSelectorType = None,
|
|
fixed_offset: float = 0,
|
|
random_offset: bool = False,
|
|
) -> np.ndarray:
|
|
"""Get samples from an audio file.
|
|
For a single-channel signal, the output is shape (num_samples,).
|
|
For a multi-channel signal, the output is shape (num_samples, num_channels).
|
|
|
|
Args:
|
|
audio_file: path to an audio file
|
|
sample_rate: desired sample rate for output samples
|
|
duration: Optional desired duration of output samples.
|
|
If `None`, the complete file will be loaded.
|
|
If set, a segment of `duration` seconds will be loaded.
|
|
channel_selector: Optional channel selector, for selecting a subset of channels.
|
|
fixed_offset: Optional fixed offset when loading samples.
|
|
random_offset: If `True`, offset will be randomized when loading a short segment
|
|
from a file. The value is randomized between fixed_offset and
|
|
max_offset (set depending on the duration and fixed_offset).
|
|
|
|
Returns:
|
|
Numpy array with samples from audio file.
|
|
The array has shape (num_samples,) for a single-channel signal
|
|
or (num_channels, num_samples) for a multi-channel signal.
|
|
"""
|
|
output = cls.get_samples_synchronized(
|
|
audio_files=[audio_file],
|
|
sample_rate=sample_rate,
|
|
duration=duration,
|
|
channel_selectors=[channel_selector],
|
|
fixed_offset=fixed_offset,
|
|
random_offset=random_offset,
|
|
)
|
|
|
|
return output[0]
|
|
|
|
@classmethod
|
|
def get_samples_synchronized(
|
|
cls,
|
|
audio_files: List[str],
|
|
sample_rate: int,
|
|
duration: Optional[float] = None,
|
|
channel_selectors: Optional[List[ChannelSelectorType]] = None,
|
|
fixed_offset: float = 0,
|
|
random_offset: bool = False,
|
|
) -> List[np.ndarray]:
|
|
"""Get samples from multiple files with the same start and end point.
|
|
|
|
Args:
|
|
audio_files: list of paths to audio files
|
|
sample_rate: desired sample rate for output samples
|
|
duration: Optional desired duration of output samples.
|
|
If `None`, the complete files will be loaded.
|
|
If set, a segment of `duration` seconds will be loaded from
|
|
all files. Segment is synchronized across files, so that
|
|
start and end points are the same.
|
|
channel_selectors: Optional channel selector for each signal, for selecting
|
|
a subset of channels.
|
|
fixed_offset: Optional fixed offset when loading samples.
|
|
random_offset: If `True`, offset will be randomized when loading a short segment
|
|
from a file. The value is randomized between fixed_offset and
|
|
max_offset (set depending on the duration and fixed_offset).
|
|
|
|
Returns:
|
|
List with the same size as `audio_files` but containing numpy arrays
|
|
with samples from each audio file.
|
|
Each array has shape (num_samples,) or (num_channels, num_samples), for single-
|
|
or multi-channel signal, respectively.
|
|
For example, if `audio_files = [path/to/file_1.wav, path/to/file_2.wav]`,
|
|
the output will be a list `output = [samples_file_1, samples_file_2]`.
|
|
"""
|
|
if channel_selectors is None:
|
|
channel_selectors = [None] * len(audio_files)
|
|
|
|
if duration is None:
|
|
# Load complete files starting from a fixed offset
|
|
offset = fixed_offset # fixed offset
|
|
num_samples = None # no constrain on the number of samples
|
|
|
|
else:
|
|
# Fixed duration of the output
|
|
audio_durations = cls.get_duration(audio_files)
|
|
min_audio_duration = min(audio_durations)
|
|
available_duration = min_audio_duration - fixed_offset
|
|
|
|
if available_duration <= 0:
|
|
raise ValueError(f'Fixed offset {fixed_offset}s is larger than shortest file {min_audio_duration}s.')
|
|
|
|
if duration + fixed_offset > min_audio_duration:
|
|
# The shortest file is shorter than the requested duration
|
|
logging.debug(
|
|
f'Shortest file ({min_audio_duration}s) is less than the desired duration {duration}s + fixed offset {fixed_offset}s. Returned signals will be shortened to {available_duration} seconds.'
|
|
)
|
|
offset = fixed_offset
|
|
duration = available_duration
|
|
elif random_offset:
|
|
# Randomize offset based on the shortest file
|
|
max_offset = min_audio_duration - duration
|
|
offset = random.uniform(fixed_offset, max_offset)
|
|
else:
|
|
# Fixed offset
|
|
offset = fixed_offset
|
|
|
|
# Fixed number of samples
|
|
num_samples = math.floor(duration * sample_rate)
|
|
|
|
output = []
|
|
|
|
# Prepare segments
|
|
for idx, audio_file in enumerate(audio_files):
|
|
segment_samples = cls.get_samples_from_file(
|
|
audio_file=audio_file,
|
|
sample_rate=sample_rate,
|
|
offset=offset,
|
|
num_samples=num_samples,
|
|
channel_selector=channel_selectors[idx],
|
|
)
|
|
output.append(segment_samples)
|
|
|
|
return output
|
|
|
|
@classmethod
|
|
def get_samples_from_file(
|
|
cls,
|
|
audio_file: Union[str, List[str]],
|
|
sample_rate: int,
|
|
offset: float,
|
|
num_samples: Optional[int] = None,
|
|
channel_selector: Optional[ChannelSelectorType] = None,
|
|
) -> np.ndarray:
|
|
"""Get samples from a single or multiple files.
|
|
If loading samples from multiple files, they will
|
|
be concatenated along the channel dimension.
|
|
|
|
Args:
|
|
audio_file: path or a list of paths.
|
|
sample_rate: sample rate of the loaded samples
|
|
offset: fixed offset in seconds
|
|
num_samples: Optional, number of samples to load.
|
|
If `None`, all available samples will be loaded.
|
|
channel_selector: Select a subset of available channels.
|
|
|
|
Returns:
|
|
An array with shape (samples,) or (channels, samples)
|
|
"""
|
|
if isinstance(audio_file, str):
|
|
# Load samples from a single file
|
|
segment_samples = cls.get_segment_from_file(
|
|
audio_file=audio_file,
|
|
sample_rate=sample_rate,
|
|
offset=offset,
|
|
num_samples=num_samples,
|
|
channel_selector=channel_selector,
|
|
)
|
|
elif isinstance(audio_file, list):
|
|
# Load samples from multiple files and form a multi-channel signal
|
|
segment_samples = []
|
|
for a_file in audio_file:
|
|
a_file_samples = cls.get_segment_from_file(
|
|
audio_file=a_file,
|
|
sample_rate=sample_rate,
|
|
offset=offset,
|
|
num_samples=num_samples,
|
|
channel_selector=channel_selector,
|
|
)
|
|
segment_samples.append(a_file_samples)
|
|
segment_samples = cls.list_to_multichannel(segment_samples)
|
|
elif audio_file is None:
|
|
# Support for inference, when the target signal is `None`
|
|
segment_samples = []
|
|
else:
|
|
raise RuntimeError(f'Unexpected audio_file type {type(audio_file)}')
|
|
return segment_samples
|
|
|
|
@staticmethod
|
|
def get_segment_from_file(
|
|
audio_file: str,
|
|
sample_rate: int,
|
|
offset: float,
|
|
num_samples: Optional[int] = None,
|
|
channel_selector: Optional[ChannelSelectorType] = None,
|
|
) -> np.ndarray:
|
|
"""Get a segment of samples from a single audio file.
|
|
|
|
Args:
|
|
audio_file: path to an audio file
|
|
sample_rate: sample rate of the loaded samples
|
|
offset: fixed offset in seconds
|
|
num_samples: Optional, number of samples to load.
|
|
If `None`, all available samples will be loaded.
|
|
channel_selector: Select a subset of available channels.
|
|
|
|
Returns:
|
|
An array with shape (samples,) or (channels, samples)
|
|
"""
|
|
if num_samples is None:
|
|
segment = AudioSegment.from_file(
|
|
audio_file=audio_file,
|
|
target_sr=sample_rate,
|
|
offset=offset,
|
|
channel_selector=channel_selector,
|
|
)
|
|
|
|
else:
|
|
segment = AudioSegment.segment_from_file(
|
|
audio_file=audio_file,
|
|
target_sr=sample_rate,
|
|
n_segments=num_samples,
|
|
offset=offset,
|
|
channel_selector=channel_selector,
|
|
)
|
|
|
|
if segment.samples.ndim == 1:
|
|
# Single-channel signal
|
|
return segment.samples
|
|
elif segment.samples.ndim == 2:
|
|
# Use multi-channel format as (channels, samples)
|
|
return segment.samples.T
|
|
else:
|
|
raise RuntimeError(f'Unexpected samples shape: {segment.samples.shape}')
|
|
|
|
@staticmethod
|
|
def list_to_multichannel(signal: Union[np.ndarray, List[np.ndarray]]) -> np.ndarray:
|
|
"""Convert a list of signals into a multi-channel signal by concatenating
|
|
the elements of the list along the channel dimension.
|
|
|
|
If input is not a list, it is returned unmodified.
|
|
|
|
Args:
|
|
signal: list of arrays
|
|
|
|
Returns:
|
|
Numpy array obtained by concatenating the elements of the list
|
|
along the channel dimension (axis=0).
|
|
"""
|
|
if not isinstance(signal, list):
|
|
# Nothing to do there
|
|
return signal
|
|
elif len(signal) == 0:
|
|
# Nothing to do, return as is
|
|
return signal
|
|
elif len(signal) == 1:
|
|
# Nothing to concatenate, return the original format
|
|
return signal[0]
|
|
|
|
# If multiple signals are provided in a list, we concatenate them along the channel dimension
|
|
if signal[0].ndim == 1:
|
|
# Single-channel individual files
|
|
mc_signal = np.stack(signal, axis=0)
|
|
elif signal[0].ndim == 2:
|
|
# Multi-channel individual files
|
|
mc_signal = np.concatenate(signal, axis=0)
|
|
else:
|
|
raise RuntimeError(f'Unexpected target with {signal[0].ndim} dimensions.')
|
|
|
|
return mc_signal
|
|
|
|
@staticmethod
|
|
def get_duration(audio_files: List[str]) -> List[float]:
|
|
"""Get duration for each audio file in `audio_files`.
|
|
|
|
Args:
|
|
audio_files: list of paths to audio files
|
|
|
|
Returns:
|
|
List of durations in seconds.
|
|
"""
|
|
duration = [librosa.get_duration(path=f) for f in flatten(audio_files)]
|
|
return duration
|
|
|
|
def load_embedding(self, example: collections.Audio.OUTPUT_TYPE) -> Dict[str, torch.Tensor]:
|
|
"""Given an example, load embedding from `example.audio_files[embedding]`
|
|
and return it in a dictionary.
|
|
|
|
Args:
|
|
example: An example from audio collection
|
|
|
|
Returns:
|
|
An dictionary of embedding keys and their tensors.
|
|
"""
|
|
output = OrderedDict()
|
|
for idx, signal in enumerate(self.embedding_setup.signals):
|
|
embedding_file = example.audio_files[signal]
|
|
embedding = self.load_embedding_vector(embedding_file)
|
|
output[signal] = torch.tensor(embedding)
|
|
return output
|
|
|
|
@staticmethod
|
|
def load_embedding_vector(filepath: str) -> np.ndarray:
|
|
"""Load an embedding vector from a file.
|
|
|
|
Args:
|
|
filepath: path to a file storing a vector.
|
|
Currently, it is assumed the file is a npy file.
|
|
|
|
Returns:
|
|
Array loaded from filepath.
|
|
"""
|
|
if filepath.endswith('.npy'):
|
|
with open(filepath, 'rb') as f:
|
|
embedding = np.load(f)
|
|
else:
|
|
raise RuntimeError(f'Unknown embedding file format in file: {filepath}')
|
|
|
|
return embedding
|
|
|
|
|
|
class BaseAudioDataset(Dataset):
|
|
"""Base class of audio datasets, providing common functionality
|
|
for other audio datasets.
|
|
|
|
Args:
|
|
collection: Collection of audio examples prepared from manifest files.
|
|
audio_processor: Used to process every example from the collection.
|
|
A callable with `process` method. For reference,
|
|
please check ASRAudioProcessor.
|
|
"""
|
|
|
|
@property
|
|
@abc.abstractmethod
|
|
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
|
"""Returns definitions of module output ports."""
|
|
|
|
def __init__(self, collection: collections.Audio, audio_processor: Callable, output_type: Type[namedtuple]):
|
|
"""Instantiates an audio dataset."""
|
|
super().__init__()
|
|
|
|
self.collection = collection
|
|
self.audio_processor = audio_processor
|
|
self.output_type = output_type
|
|
|
|
def num_channels(self, signal_key) -> int:
|
|
"""Returns the number of channels for a particular signal in
|
|
items prepared by this dictionary.
|
|
|
|
More specifically, this will get the tensor from the first
|
|
item in the dataset, check if it's a one- or two-dimensional
|
|
tensor, and return the number of channels based on the size
|
|
of the first axis (shape[0]).
|
|
|
|
NOTE:
|
|
This assumes that all examples have the same number of channels.
|
|
|
|
Args:
|
|
signal_key: string, used to select a signal from the dictionary
|
|
output by __getitem__
|
|
|
|
Returns:
|
|
Number of channels for the selected signal.
|
|
"""
|
|
# Assumption: whole dataset has the same number of channels
|
|
item = self.__getitem__(0)
|
|
|
|
if item[signal_key].ndim == 1:
|
|
return 1
|
|
elif item[signal_key].ndim == 2:
|
|
return item[signal_key].shape[0]
|
|
else:
|
|
raise RuntimeError(
|
|
f'Unexpected number of dimension for signal {signal_key} with shape {item[signal_key].shape}'
|
|
)
|
|
|
|
def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
|
|
"""Return a single example from the dataset.
|
|
|
|
Args:
|
|
index: integer index of an example in the collection
|
|
|
|
Returns:
|
|
Dictionary providing mapping from signal to its tensor.
|
|
For example:
|
|
```
|
|
{
|
|
'input_signal': input_signal_tensor,
|
|
'target_signal': target_signal_tensor,
|
|
}
|
|
```
|
|
"""
|
|
example = self.collection[index]
|
|
output = self.audio_processor.process(example=example)
|
|
|
|
return output
|
|
|
|
def __len__(self) -> int:
|
|
"""Return the number of examples in the dataset."""
|
|
return len(self.collection)
|
|
|
|
def _collate_fn(self, batch) -> Tuple[torch.Tensor]:
|
|
"""Collate items in a batch."""
|
|
return self.output_type(*_audio_collate_fn(batch))
|
|
|
|
|
|
AudioToTargetExample = namedtuple(
|
|
typename='AudioToTargetExample', field_names='input_signal input_length target_signal target_length'
|
|
)
|
|
|
|
|
|
class AudioToTargetDataset(BaseAudioDataset):
|
|
"""A dataset for audio-to-audio tasks where the goal is to use
|
|
an input signal to recover the corresponding target signal.
|
|
|
|
Each line of the manifest file is expected to have the following format:
|
|
|
|
.. code-block:: json
|
|
|
|
{"input_key": "path/to/input.wav", "target_key": "path/to/target.wav", "duration": "duration_in_seconds"}
|
|
|
|
Additionally, multiple audio files may be provided for each key in the manifest, for example,
|
|
|
|
.. code-block:: json
|
|
|
|
{"input_key": "path/to/input.wav", "target_key": ["path/to/path_to_target_ch0.wav", "path/to/path_to_target_ch1.wav"], "duration": "duration_in_seconds"}
|
|
|
|
Keys for input and target signals can be configured in the constructor (`input_key` and `target_key`).
|
|
|
|
Args:
|
|
manifest_filepath: Path to manifest file in a format described above.
|
|
sample_rate: Sample rate for loaded audio signals.
|
|
input_key: Key pointing to input audio files in the manifest
|
|
target_key: Key pointing to target audio files in manifest
|
|
audio_duration: Optional duration of each item returned by __getitem__.
|
|
If `None`, complete audio will be loaded.
|
|
If set, a random subsegment will be loaded synchronously from
|
|
target and audio, i.e., with the same start and end point.
|
|
random_offset: If `True`, offset will be randomized when loading a subsegment
|
|
from a file.
|
|
max_duration: If audio exceeds this length, do not include in dataset.
|
|
min_duration: If audio is less than this length, do not include in dataset.
|
|
max_utts: Limit number of utterances.
|
|
input_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
target_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
normalization_signal: Normalize audio signals with a scale that ensures the normalization signal is in range [-1, 1].
|
|
All audio signals are scaled by the same factor. Supported values are `None` (no normalization),
|
|
'input_signal', 'target_signal'.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
manifest_filepath: str,
|
|
sample_rate: int,
|
|
input_key: str,
|
|
target_key: str,
|
|
audio_duration: Optional[float] = None,
|
|
random_offset: bool = False,
|
|
max_duration: Optional[float] = None,
|
|
min_duration: Optional[float] = None,
|
|
max_utts: Optional[int] = None,
|
|
input_channel_selector: Optional[int] = None,
|
|
target_channel_selector: Optional[int] = None,
|
|
normalization_signal: Optional[str] = None,
|
|
):
|
|
audio_to_manifest_key = {
|
|
'input_signal': input_key,
|
|
'target_signal': target_key,
|
|
}
|
|
|
|
collection = collections.AudioCollection(
|
|
manifest_files=manifest_filepath,
|
|
audio_to_manifest_key=audio_to_manifest_key,
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
max_number=max_utts,
|
|
)
|
|
|
|
audio_processor = ASRAudioProcessor(
|
|
sample_rate=sample_rate,
|
|
random_offset=random_offset,
|
|
normalization_signal=normalization_signal,
|
|
)
|
|
audio_processor.sync_setup = SignalSetup(
|
|
signals=['input_signal', 'target_signal'],
|
|
duration=audio_duration,
|
|
channel_selectors=[input_channel_selector, target_channel_selector],
|
|
)
|
|
|
|
super().__init__(collection=collection, audio_processor=audio_processor, output_type=AudioToTargetExample)
|
|
|
|
@property
|
|
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
|
"""Returns definitions of module output ports.
|
|
|
|
Returns:
|
|
OrderedDict: Dictionary containing the following items:
|
|
input_signal:
|
|
Batched single- or multi-channel input audio signal
|
|
input_length:
|
|
Batched original length of each input signal
|
|
target_signal:
|
|
Batched single- or multi-channel target audio signal
|
|
target_length:
|
|
Batched original length of each target signal
|
|
"""
|
|
sc_audio_type = NeuralType(('B', 'T'), AudioSignal())
|
|
mc_audio_type = NeuralType(('B', 'C', 'T'), AudioSignal())
|
|
|
|
return OrderedDict(
|
|
input_signal=sc_audio_type if self.num_channels('input_signal') == 1 else mc_audio_type,
|
|
input_length=NeuralType(('B',), LengthsType()),
|
|
target_signal=sc_audio_type if self.num_channels('target_signal') == 1 else mc_audio_type,
|
|
target_length=NeuralType(('B',), LengthsType()),
|
|
)
|
|
|
|
|
|
AudioToTargetWithReferenceExample = namedtuple(
|
|
typename='AudioToTargetWithReferenceExample',
|
|
field_names='input_signal input_length target_signal target_length reference_signal reference_length',
|
|
)
|
|
|
|
|
|
class AudioToTargetWithReferenceDataset(BaseAudioDataset):
|
|
"""A dataset for audio-to-audio tasks where the goal is to use
|
|
an input signal to recover the corresponding target signal and an
|
|
additional reference signal is available.
|
|
|
|
This can be used, for example, when a reference signal is
|
|
available from
|
|
- enrollment utterance for the target signal
|
|
- echo reference from playback
|
|
- reference from another sensor that correlates with the target signal
|
|
|
|
Each line of the manifest file is expected to have the following format
|
|
|
|
.. code-block:: json
|
|
|
|
{"input_key": "path/to/input.wav", "target_key": "path/to/path_to_target.wav", "reference_key": "path/to/path_to_reference.wav", "duration": "duration_in_seconds"}
|
|
|
|
Keys for input, target and reference signals can be configured in the constructor.
|
|
|
|
Args:
|
|
manifest_filepath: Path to manifest file in a format described above.
|
|
sample_rate: Sample rate for loaded audio signals.
|
|
input_key: Key pointing to input audio files in the manifest
|
|
target_key: Key pointing to target audio files in manifest
|
|
reference_key: Key pointing to reference audio files in manifest
|
|
audio_duration: Optional duration of each item returned by __getitem__.
|
|
If `None`, complete audio will be loaded.
|
|
If set, a random subsegment will be loaded synchronously from
|
|
target and audio, i.e., with the same start and end point.
|
|
random_offset: If `True`, offset will be randomized when loading a subsegment
|
|
from a file.
|
|
max_duration: If audio exceeds this length, do not include in dataset.
|
|
min_duration: If audio is less than this length, do not include in dataset.
|
|
max_utts: Limit number of utterances.
|
|
input_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
target_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
reference_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
reference_is_synchronized: If True, it is assumed that the reference signal is synchronized
|
|
with the input signal, so the same subsegment will be loaded as for
|
|
input and target. If False, reference signal will be loaded independently
|
|
from input and target.
|
|
reference_duration: Optional, can be used to set a fixed duration of the reference utterance. If `None`,
|
|
complete audio file will be loaded.
|
|
normalization_signal: Normalize audio signals with a scale that ensures the normalization signal is in range [-1, 1].
|
|
All audio signals are scaled by the same factor. Supported values are `None` (no normalization),
|
|
'input_signal', 'target_signal', 'reference_signal'.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
manifest_filepath: str,
|
|
sample_rate: int,
|
|
input_key: str,
|
|
target_key: str,
|
|
reference_key: str,
|
|
audio_duration: Optional[float] = None,
|
|
random_offset: bool = False,
|
|
max_duration: Optional[float] = None,
|
|
min_duration: Optional[float] = None,
|
|
max_utts: Optional[int] = None,
|
|
input_channel_selector: Optional[int] = None,
|
|
target_channel_selector: Optional[int] = None,
|
|
reference_channel_selector: Optional[int] = None,
|
|
reference_is_synchronized: bool = True,
|
|
reference_duration: Optional[float] = None,
|
|
normalization_signal: Optional[str] = None,
|
|
):
|
|
audio_to_manifest_key = {
|
|
'input_signal': input_key,
|
|
'target_signal': target_key,
|
|
'reference_signal': reference_key,
|
|
}
|
|
|
|
collection = collections.AudioCollection(
|
|
manifest_files=manifest_filepath,
|
|
audio_to_manifest_key=audio_to_manifest_key,
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
max_number=max_utts,
|
|
)
|
|
|
|
audio_processor = ASRAudioProcessor(
|
|
sample_rate=sample_rate,
|
|
random_offset=random_offset,
|
|
normalization_signal=normalization_signal,
|
|
)
|
|
|
|
if reference_is_synchronized:
|
|
audio_processor.sync_setup = SignalSetup(
|
|
signals=['input_signal', 'target_signal', 'reference_signal'],
|
|
duration=audio_duration,
|
|
channel_selectors=[input_channel_selector, target_channel_selector, reference_channel_selector],
|
|
)
|
|
else:
|
|
audio_processor.sync_setup = SignalSetup(
|
|
signals=['input_signal', 'target_signal'],
|
|
duration=audio_duration,
|
|
channel_selectors=[input_channel_selector, target_channel_selector],
|
|
)
|
|
audio_processor.async_setup = SignalSetup(
|
|
signals=['reference_signal'],
|
|
duration=[reference_duration],
|
|
channel_selectors=[reference_channel_selector],
|
|
)
|
|
|
|
super().__init__(
|
|
collection=collection, audio_processor=audio_processor, output_type=AudioToTargetWithReferenceExample
|
|
)
|
|
|
|
@property
|
|
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
|
"""Returns definitions of module output ports.
|
|
|
|
Returns:
|
|
OrderedDict: Dictionary containing the following items:
|
|
input_signal:
|
|
Batched single- or multi-channel input audio signal
|
|
input_length:
|
|
Batched original length of each input signal
|
|
target_signal:
|
|
Batched single- or multi-channel target audio signal
|
|
target_length:
|
|
Batched original length of each target signal
|
|
reference_signal:
|
|
Batched single- or multi-channel reference audio signal
|
|
reference_length:
|
|
Batched original length of each reference signal
|
|
"""
|
|
sc_audio_type = NeuralType(('B', 'T'), AudioSignal())
|
|
mc_audio_type = NeuralType(('B', 'C', 'T'), AudioSignal())
|
|
|
|
return OrderedDict(
|
|
input_signal=sc_audio_type if self.num_channels('input_signal') == 1 else mc_audio_type,
|
|
input_length=NeuralType(('B',), LengthsType()),
|
|
target_signal=sc_audio_type if self.num_channels('target_signal') == 1 else mc_audio_type,
|
|
target_length=NeuralType(('B',), LengthsType()),
|
|
reference_signal=sc_audio_type if self.num_channels('reference_signal') == 1 else mc_audio_type,
|
|
reference_length=NeuralType(('B',), LengthsType()),
|
|
)
|
|
|
|
|
|
AudioToTargetWithEmbeddingExample = namedtuple(
|
|
typename='AudioToTargetWithEmbeddingExample',
|
|
field_names='input_signal input_length target_signal target_length embedding_vector embedding_length',
|
|
)
|
|
|
|
|
|
class AudioToTargetWithEmbeddingDataset(BaseAudioDataset):
|
|
"""A dataset for audio-to-audio tasks where the goal is to use
|
|
an input signal to recover the corresponding target signal and an
|
|
additional embedding signal. It is assumed that the embedding
|
|
is in a form of a vector.
|
|
|
|
Each line of the manifest file is expected to have the following format
|
|
|
|
.. code-block:: json
|
|
|
|
{"input_key": "path/to/input.wav", "target_key": "path/to/path_to_target.wav", "embedding_key": "path/to/path_to_reference.npy", "duration": "duration_in_seconds"}
|
|
|
|
Keys for input, target and embedding signals can be configured in the constructor.
|
|
|
|
Args:
|
|
manifest_filepath: Path to manifest file in a format described above.
|
|
sample_rate: Sample rate for loaded audio signals.
|
|
input_key: Key pointing to input audio files in the manifest
|
|
target_key: Key pointing to target audio files in manifest
|
|
embedding_key: Key pointing to embedding files in manifest
|
|
audio_duration: Optional duration of each item returned by __getitem__.
|
|
If `None`, complete audio will be loaded.
|
|
If set, a random subsegment will be loaded synchronously from
|
|
target and audio, i.e., with the same start and end point.
|
|
random_offset: If `True`, offset will be randomized when loading a subsegment
|
|
from a file.
|
|
max_duration: If audio exceeds this length, do not include in dataset.
|
|
min_duration: If audio is less than this length, do not include in dataset.
|
|
max_utts: Limit number of utterances.
|
|
input_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
target_channel_selector: Optional, select subset of channels from each input audio file.
|
|
If `None`, all channels will be loaded.
|
|
normalization_signal: Normalize audio signals with a scale that ensures the normalization signal is in range [-1, 1].
|
|
All audio signals are scaled by the same factor. Supported values are `None` (no normalization),
|
|
'input_signal', 'target_signal'.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
manifest_filepath: str,
|
|
sample_rate: int,
|
|
input_key: str,
|
|
target_key: str,
|
|
embedding_key: str,
|
|
audio_duration: Optional[float] = None,
|
|
random_offset: bool = False,
|
|
max_duration: Optional[float] = None,
|
|
min_duration: Optional[float] = None,
|
|
max_utts: Optional[int] = None,
|
|
input_channel_selector: Optional[int] = None,
|
|
target_channel_selector: Optional[int] = None,
|
|
normalization_signal: Optional[str] = None,
|
|
):
|
|
audio_to_manifest_key = {
|
|
'input_signal': input_key,
|
|
'target_signal': target_key,
|
|
'embedding_vector': embedding_key,
|
|
}
|
|
|
|
collection = collections.AudioCollection(
|
|
manifest_files=manifest_filepath,
|
|
audio_to_manifest_key=audio_to_manifest_key,
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
max_number=max_utts,
|
|
)
|
|
|
|
audio_processor = ASRAudioProcessor(
|
|
sample_rate=sample_rate,
|
|
random_offset=random_offset,
|
|
normalization_signal=normalization_signal,
|
|
)
|
|
audio_processor.sync_setup = SignalSetup(
|
|
signals=['input_signal', 'target_signal'],
|
|
duration=audio_duration,
|
|
channel_selectors=[input_channel_selector, target_channel_selector],
|
|
)
|
|
audio_processor.embedding_setup = SignalSetup(signals=['embedding_vector'])
|
|
|
|
super().__init__(
|
|
collection=collection, audio_processor=audio_processor, output_type=AudioToTargetWithEmbeddingExample
|
|
)
|
|
|
|
@property
|
|
def output_types(self) -> Optional[Dict[str, NeuralType]]:
|
|
"""Returns definitions of module output ports.
|
|
|
|
Returns:
|
|
OrderedDict: Dictionary containing the following items:
|
|
input_signal:
|
|
Batched single- or multi-channel input audio signal
|
|
input_length:
|
|
Batched original length of each input signal
|
|
target_signal:
|
|
Batched single- or multi-channel target audio signal
|
|
target_length:
|
|
Batched original length of each target signal
|
|
embedding_vector:
|
|
Batched embedded vector format
|
|
embedding_length:
|
|
Batched original length of each embedding vector
|
|
"""
|
|
sc_audio_type = NeuralType(('B', 'T'), AudioSignal())
|
|
mc_audio_type = NeuralType(('B', 'C', 'T'), AudioSignal())
|
|
|
|
return OrderedDict(
|
|
input_signal=sc_audio_type if self.num_channels('input_signal') == 1 else mc_audio_type,
|
|
input_length=NeuralType(('B',), LengthsType()),
|
|
target_signal=sc_audio_type if self.num_channels('target_signal') == 1 else mc_audio_type,
|
|
target_length=NeuralType(('B',), LengthsType()),
|
|
embedding_vector=NeuralType(('B', 'D'), EncodedRepresentation()),
|
|
embedding_length=NeuralType(('B',), LengthsType()),
|
|
)
|