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
1157 lines
49 KiB
Python
1157 lines
49 KiB
Python
# Copyright (c) 2020, 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 os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import soundfile as sf
|
|
import torch.cuda
|
|
from omegaconf import OmegaConf
|
|
|
|
from nemo.collections.asr.parts.utils.manifest_utils import write_manifest
|
|
from nemo.collections.audio.data import audio_to_audio_dataset
|
|
from nemo.collections.audio.data.audio_to_audio import (
|
|
ASRAudioProcessor,
|
|
AudioToTargetDataset,
|
|
AudioToTargetWithEmbeddingDataset,
|
|
AudioToTargetWithReferenceDataset,
|
|
_audio_collate_fn,
|
|
)
|
|
from nemo.collections.audio.data.audio_to_audio_lhotse import (
|
|
LhotseAudioToTargetDataset,
|
|
convert_manifest_nemo_to_lhotse,
|
|
)
|
|
from nemo.collections.audio.parts.utils.audio import get_segment_start
|
|
from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
|
|
|
|
|
|
class TestAudioDatasets:
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize('num_channels', [1, 2])
|
|
@pytest.mark.parametrize('num_targets', [1, 3])
|
|
def test_list_to_multichannel(self, num_channels, num_targets):
|
|
"""Test conversion of a list of arrays into"""
|
|
random_seed = 42
|
|
num_samples = 1000
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
# Multi-channel signal
|
|
golden_target = _rng.normal(size=(num_channels * num_targets, num_samples))
|
|
|
|
# Create a list of num_targets signals with num_channels channels
|
|
target_list = [golden_target[n * num_channels : (n + 1) * num_channels, :] for n in range(num_targets)]
|
|
|
|
# Check the original signal is not modified
|
|
assert (ASRAudioProcessor.list_to_multichannel(golden_target) == golden_target).all()
|
|
# Check the list is converted back to the original signal
|
|
assert (ASRAudioProcessor.list_to_multichannel(target_list) == golden_target).all()
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize('num_channels', [1, 2])
|
|
def test_processor_process_audio(self, num_channels):
|
|
"""Test signal normalization in process_audio."""
|
|
num_samples = 1000
|
|
num_examples = 30
|
|
|
|
signals = ['input_signal', 'target_signal', 'reference_signal']
|
|
|
|
for normalization_signal in [None] + signals:
|
|
# Create processor
|
|
processor = ASRAudioProcessor(
|
|
sample_rate=16000, random_offset=False, normalization_signal=normalization_signal
|
|
)
|
|
|
|
# Generate random signals
|
|
for n in range(num_examples):
|
|
example = {signal: torch.randn(num_channels, num_samples) for signal in signals}
|
|
processed_example = processor.process_audio(example)
|
|
|
|
# Expected scale
|
|
if normalization_signal:
|
|
scale = 1.0 / (example[normalization_signal].abs().max() + processor.eps)
|
|
else:
|
|
scale = 1.0
|
|
|
|
# Make sure all signals are scaled as expected
|
|
for signal in signals:
|
|
assert torch.allclose(
|
|
processed_example[signal], example[signal] * scale
|
|
), f'Failed example {n} signal {signal}'
|
|
|
|
@pytest.mark.unit
|
|
def test_audio_collate_fn(self):
|
|
"""Test `_audio_collate_fn`"""
|
|
batch_size = 16
|
|
random_seed = 42
|
|
atol = 1e-5
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
signal_to_channels = {
|
|
'input_signal': 2,
|
|
'target_signal': 1,
|
|
'reference_signal': 1,
|
|
}
|
|
|
|
signal_to_length = {
|
|
'input_signal': _rng.integers(low=5, high=25, size=batch_size),
|
|
'target_signal': _rng.integers(low=5, high=25, size=batch_size),
|
|
'reference_signal': _rng.integers(low=5, high=25, size=batch_size),
|
|
}
|
|
|
|
# Generate batch
|
|
batch = []
|
|
for n in range(batch_size):
|
|
item = dict()
|
|
for signal, num_channels in signal_to_channels.items():
|
|
random_signal = _rng.normal(size=(num_channels, signal_to_length[signal][n]))
|
|
random_signal = np.squeeze(random_signal) # get rid of channel dimention for single-channel
|
|
item[signal] = torch.tensor(random_signal)
|
|
batch.append(item)
|
|
|
|
# Run UUT
|
|
batched = _audio_collate_fn(batch)
|
|
|
|
batched_signals = {
|
|
'input_signal': batched[0].cpu().detach().numpy(),
|
|
'target_signal': batched[2].cpu().detach().numpy(),
|
|
'reference_signal': batched[4].cpu().detach().numpy(),
|
|
}
|
|
|
|
batched_lengths = {
|
|
'input_signal': batched[1].cpu().detach().numpy(),
|
|
'target_signal': batched[3].cpu().detach().numpy(),
|
|
'reference_signal': batched[5].cpu().detach().numpy(),
|
|
}
|
|
|
|
# Check outputs
|
|
for signal, b_signal in batched_signals.items():
|
|
for n in range(batch_size):
|
|
# Check length
|
|
uut_length = batched_lengths[signal][n]
|
|
golden_length = signal_to_length[signal][n]
|
|
assert (
|
|
uut_length == golden_length
|
|
), f'Example {n} signal {signal} length mismatch: batched ({uut_length}) != golden ({golden_length})'
|
|
|
|
uut_signal = b_signal[n][:uut_length, ...]
|
|
golden_signal = batch[n][signal][:uut_length, ...].cpu().detach().numpy()
|
|
assert np.allclose(
|
|
uut_signal, golden_signal, atol=atol
|
|
), f'Example {n} signal {signal} value mismatch.'
|
|
|
|
@pytest.mark.unit
|
|
def test_audio_to_target_dataset(self):
|
|
"""Test AudioWithTargetDataset in different configurations.
|
|
|
|
Test below cover the following:
|
|
1) no constraints
|
|
2) filtering based on signal duration
|
|
3) use with channel selector
|
|
4) use with fixed audio duration and random subsegments
|
|
5) collate a batch of items
|
|
|
|
In this use case, each line of the manifest file has the following format:
|
|
```
|
|
{
|
|
'input_filepath': 'path/to/input.wav',
|
|
'target_filepath': 'path/to/path_to_target.wav',
|
|
'duration': duration_of_input,
|
|
}
|
|
```
|
|
"""
|
|
# Data setup
|
|
random_seed = 42
|
|
sample_rate = 16000
|
|
num_examples = 25
|
|
data_num_channels = {
|
|
'input_signal': 4,
|
|
'target_signal': 2,
|
|
}
|
|
data_min_duration = 2.0
|
|
data_max_duration = 8.0
|
|
data_key = {
|
|
'input_signal': 'input_filepath',
|
|
'target_signal': 'target_filepath',
|
|
}
|
|
|
|
# Tolerance
|
|
atol = 1e-6
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
# Input and target signals have the same duration
|
|
data_duration = np.round(_rng.uniform(low=data_min_duration, high=data_max_duration, size=num_examples), 3)
|
|
data_duration_samples = np.floor(data_duration * sample_rate).astype(int)
|
|
|
|
data = dict()
|
|
for signal, num_channels in data_num_channels.items():
|
|
data[signal] = []
|
|
for n in range(num_examples):
|
|
if num_channels == 1:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(data_duration_samples[n]))
|
|
else:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(num_channels, data_duration_samples[n]))
|
|
data[signal].append(random_signal)
|
|
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
|
|
# Build metadata for manifest
|
|
metadata = []
|
|
|
|
for n in range(num_examples):
|
|
|
|
meta = dict()
|
|
|
|
for signal in data:
|
|
# filenames
|
|
signal_filename = f'{signal}_{n:02d}.wav'
|
|
|
|
# write audio files
|
|
sf.write(os.path.join(test_dir, signal_filename), data[signal][n].T, sample_rate, 'float')
|
|
|
|
# update metadata
|
|
meta[data_key[signal]] = signal_filename
|
|
|
|
meta['duration'] = data_duration[n]
|
|
metadata.append(meta)
|
|
|
|
# Save manifest
|
|
manifest_filepath = os.path.join(test_dir, 'manifest.json')
|
|
write_manifest(manifest_filepath, metadata)
|
|
|
|
# Test 1
|
|
# - No constraints on channels or duration
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
# Also test the corresponding factory
|
|
config = {
|
|
'manifest_filepath': manifest_filepath,
|
|
'input_key': data_key['input_signal'],
|
|
'target_key': data_key['target_signal'],
|
|
'sample_rate': sample_rate,
|
|
}
|
|
dataset_factory = audio_to_audio_dataset.get_audio_to_target_dataset(config)
|
|
|
|
# Prepare lhotse manifest
|
|
cuts_path = manifest_filepath.replace('.json', '_cuts.jsonl')
|
|
convert_manifest_nemo_to_lhotse(
|
|
input_manifest=manifest_filepath,
|
|
output_manifest=cuts_path,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
)
|
|
|
|
# Prepare lhotse dataset
|
|
config_lhotse = {
|
|
'cuts_path': cuts_path,
|
|
'use_lhotse': True,
|
|
'sample_rate': sample_rate,
|
|
'batch_size': 1,
|
|
}
|
|
dl_lhotse = get_lhotse_dataloader_from_config(
|
|
OmegaConf.create(config_lhotse), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
|
)
|
|
dataset_lhotse = [item for item in dl_lhotse]
|
|
|
|
# Test number of channels
|
|
for signal in data:
|
|
assert data_num_channels[signal] == dataset.num_channels(
|
|
signal
|
|
), f'Num channels not correct for signal {signal}'
|
|
assert data_num_channels[signal] == dataset_factory.num_channels(
|
|
signal
|
|
), f'Num channels not correct for signal {signal}'
|
|
|
|
# Test returned examples
|
|
for n in range(num_examples):
|
|
for signal in data:
|
|
golden_signal = data[signal][n]
|
|
|
|
for use_lhotse in [False, True]:
|
|
item_signal = (
|
|
dataset_lhotse[n][signal].squeeze(0) if use_lhotse else dataset.__getitem__(n)[signal]
|
|
)
|
|
item_factory_signal = dataset_factory.__getitem__(n)[signal]
|
|
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Test 1, use_lhotse={use_lhotse}: Signal {signal} item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 1, use_lhotse={use_lhotse}: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
assert np.allclose(
|
|
item_factory_signal, golden_signal, atol=atol
|
|
), f'Test 1, use_lhotse={use_lhotse}: Failed for factory example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 2
|
|
# - Filtering based on signal duration
|
|
min_duration = 3.5
|
|
max_duration = 7.5
|
|
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
# Prepare lhotse dataset
|
|
config_lhotse = {
|
|
'cuts_path': cuts_path,
|
|
'use_lhotse': True,
|
|
'min_duration': min_duration,
|
|
'max_duration': max_duration,
|
|
'sample_rate': sample_rate,
|
|
'batch_size': 1,
|
|
}
|
|
dl_lhotse = get_lhotse_dataloader_from_config(
|
|
OmegaConf.create(config_lhotse), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
|
)
|
|
dataset_lhotse = [item for item in dl_lhotse]
|
|
|
|
filtered_examples = [n for n, val in enumerate(data_duration) if min_duration <= val <= max_duration]
|
|
|
|
for n in range(len(dataset)):
|
|
for use_lhotse in [False, True]:
|
|
for signal in data:
|
|
item_signal = (
|
|
dataset_lhotse[n][signal].squeeze(0) if use_lhotse else dataset.__getitem__(n)[signal]
|
|
)
|
|
golden_signal = data[signal][filtered_examples[n]]
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Test 2, use_lhotse={use_lhotse}: Signal {signal} item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 2, use_lhotse={use_lhotse}: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 3
|
|
# - Use channel selector
|
|
channel_selector = {
|
|
'input_signal': [0, 2],
|
|
'target_signal': 1,
|
|
}
|
|
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
input_channel_selector=channel_selector['input_signal'],
|
|
target_channel_selector=channel_selector['target_signal'],
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
for n in range(len(dataset)):
|
|
item = dataset.__getitem__(n)
|
|
|
|
for signal in data:
|
|
cs = channel_selector[signal]
|
|
item_signal = item[signal].cpu().detach().numpy()
|
|
golden_signal = data[signal][n][cs, ...]
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Signal {signal}: item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 3: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 4
|
|
# - Use fixed duration (random segment selection)
|
|
audio_duration = 4.0
|
|
audio_duration_samples = int(np.floor(audio_duration * sample_rate))
|
|
|
|
filtered_examples = [n for n, val in enumerate(data_duration) if val >= audio_duration]
|
|
|
|
for random_offset in [True, False]:
|
|
# Test subsegments with the default fixed offset and a random offset
|
|
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
sample_rate=sample_rate,
|
|
min_duration=audio_duration,
|
|
audio_duration=audio_duration,
|
|
random_offset=random_offset, # random offset when selecting subsegment
|
|
)
|
|
|
|
# Prepare lhotse dataset
|
|
config_lhotse = {
|
|
'cuts_path': cuts_path,
|
|
'use_lhotse': True,
|
|
'min_duration': audio_duration,
|
|
'truncate_duration': audio_duration,
|
|
'truncate_offset_type': 'random' if random_offset else 'start',
|
|
'sample_rate': sample_rate,
|
|
'batch_size': 1,
|
|
}
|
|
dl_lhotse = get_lhotse_dataloader_from_config(
|
|
OmegaConf.create(config_lhotse), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
|
)
|
|
dataset_lhotse = [item for item in dl_lhotse]
|
|
|
|
for n in range(len(dataset)):
|
|
for use_lhotse in [False, True]:
|
|
item = dataset_lhotse[n] if use_lhotse else dataset.__getitem__(n)
|
|
golden_start = golden_end = None
|
|
for signal in data:
|
|
item_signal = item[signal].squeeze(0) if use_lhotse else item[signal]
|
|
full_golden_signal = data[signal][filtered_examples[n]]
|
|
|
|
# Find random segment using correlation on the first channel
|
|
# of the first signal, and then use it fixed for other signals
|
|
if golden_start is None:
|
|
golden_start = get_segment_start(
|
|
signal=full_golden_signal[0, :], segment=item_signal[0, :]
|
|
)
|
|
if not random_offset:
|
|
assert (
|
|
golden_start == 0
|
|
), f'Test 4, use_lhotse={use_lhotse}: Expecting the signal to start at 0 when random_offset is False'
|
|
|
|
golden_end = golden_start + audio_duration_samples
|
|
golden_signal = full_golden_signal[..., golden_start:golden_end]
|
|
|
|
# Test length is correct
|
|
assert (
|
|
item_signal.shape[-1] == audio_duration_samples
|
|
), f'Test 4, use_lhotse={use_lhotse}: Signal length ({item_signal.shape[-1]}) not matching the expected length ({audio_duration_samples})'
|
|
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Test 4, use_lhotse={use_lhotse}: Signal {signal} item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
# Test signal values
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 4, use_lhotse={use_lhotse}: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 5:
|
|
# - Test collate_fn
|
|
batch_size = 16
|
|
|
|
for use_lhotse in [False, True]:
|
|
if use_lhotse:
|
|
# Get batch from lhotse dataloader
|
|
config_lhotse['batch_size'] = batch_size
|
|
dl_lhotse = get_lhotse_dataloader_from_config(
|
|
OmegaConf.create(config_lhotse),
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=LhotseAudioToTargetDataset(),
|
|
)
|
|
batched = next(iter(dl_lhotse))
|
|
else:
|
|
# Get examples from dataset and collate into a batch
|
|
batch = [dataset.__getitem__(n) for n in range(batch_size)]
|
|
batched = dataset.collate_fn(batch)
|
|
|
|
# Test all shapes and lengths
|
|
for n, signal in enumerate(data.keys()):
|
|
length = signal.replace('_signal', '_length')
|
|
|
|
if isinstance(batched, dict):
|
|
signal_shape = batched[signal].shape
|
|
signal_len = batched[length]
|
|
else:
|
|
signal_shape = batched[2 * n].shape
|
|
signal_len = batched[2 * n + 1]
|
|
|
|
assert signal_shape == (
|
|
batch_size,
|
|
data_num_channels[signal],
|
|
audio_duration_samples,
|
|
), f'Test 5, use_lhotse={use_lhotse}: Unexpected signal {signal} shape {signal_shape}'
|
|
assert (
|
|
len(signal_len) == batch_size
|
|
), f'Test 5, use_lhotse={use_lhotse}: Unexpected length of signal_len ({len(signal_len)})'
|
|
assert all(
|
|
signal_len == audio_duration_samples
|
|
), f'Test 5, use_lhotse={use_lhotse}: Unexpected signal_len {signal_len}'
|
|
|
|
@pytest.mark.unit
|
|
def test_audio_to_target_dataset_with_target_list(self):
|
|
"""Test AudioWithTargetDataset when the input manifest has a list
|
|
of audio files in the target key.
|
|
|
|
In this use case, each line of the manifest file has the following format:
|
|
```
|
|
{
|
|
'input_filepath': 'path/to/input.wav',
|
|
'target_filepath': ['path/to/path_to_target_ch0.wav', 'path/to/path_to_target_ch1.wav'],
|
|
'duration': duration_of_input,
|
|
}
|
|
```
|
|
"""
|
|
# Data setup
|
|
random_seed = 42
|
|
sample_rate = 16000
|
|
num_examples = 25
|
|
data_num_channels = {
|
|
'input_signal': 4,
|
|
'target_signal': 2,
|
|
}
|
|
data_min_duration = 2.0
|
|
data_max_duration = 8.0
|
|
data_key = {
|
|
'input_signal': 'input_filepath',
|
|
'target_signal': 'target_filepath',
|
|
}
|
|
|
|
# Tolerance
|
|
atol = 1e-6
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
# Input and target signals have the same duration
|
|
data_duration = np.round(_rng.uniform(low=data_min_duration, high=data_max_duration, size=num_examples), 3)
|
|
data_duration_samples = np.floor(data_duration * sample_rate).astype(int)
|
|
|
|
data = dict()
|
|
for signal, num_channels in data_num_channels.items():
|
|
data[signal] = []
|
|
for n in range(num_examples):
|
|
if num_channels == 1:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(data_duration_samples[n]))
|
|
else:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(num_channels, data_duration_samples[n]))
|
|
data[signal].append(random_signal)
|
|
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
|
|
# Build metadata for manifest
|
|
metadata = []
|
|
|
|
for n in range(num_examples):
|
|
|
|
meta = dict()
|
|
|
|
for signal in data:
|
|
if signal == 'target_signal':
|
|
# Save targets as individual files
|
|
signal_filename = []
|
|
for ch in range(data_num_channels[signal]):
|
|
# add current filename
|
|
signal_filename.append(f'{signal}_{n:02d}_ch_{ch}.wav')
|
|
# write audio file
|
|
sf.write(
|
|
os.path.join(test_dir, signal_filename[-1]),
|
|
data[signal][n][ch, :],
|
|
sample_rate,
|
|
'float',
|
|
)
|
|
else:
|
|
# single file
|
|
signal_filename = f'{signal}_{n:02d}.wav'
|
|
|
|
# write audio files
|
|
sf.write(os.path.join(test_dir, signal_filename), data[signal][n].T, sample_rate, 'float')
|
|
|
|
# update metadata
|
|
meta[data_key[signal]] = signal_filename
|
|
|
|
meta['duration'] = data_duration[n]
|
|
metadata.append(meta)
|
|
|
|
# Save manifest
|
|
manifest_filepath = os.path.join(test_dir, 'manifest.json')
|
|
write_manifest(manifest_filepath, metadata)
|
|
|
|
# Test 1
|
|
# - No constraints on channels or duration
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
config = {
|
|
'manifest_filepath': manifest_filepath,
|
|
'input_key': data_key['input_signal'],
|
|
'target_key': data_key['target_signal'],
|
|
'sample_rate': sample_rate,
|
|
}
|
|
dataset_factory = audio_to_audio_dataset.get_audio_to_target_dataset(config)
|
|
|
|
# Prepare lhotse manifest
|
|
cuts_path = manifest_filepath.replace('.json', '_cuts.jsonl')
|
|
convert_manifest_nemo_to_lhotse(
|
|
input_manifest=manifest_filepath,
|
|
output_manifest=cuts_path,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
)
|
|
|
|
# Prepare lhotse dataset
|
|
config_lhotse = {
|
|
'cuts_path': cuts_path,
|
|
'use_lhotse': True,
|
|
'sample_rate': sample_rate,
|
|
'batch_size': 1,
|
|
}
|
|
dl_lhotse = get_lhotse_dataloader_from_config(
|
|
OmegaConf.create(config_lhotse), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
|
)
|
|
dataset_lhotse = [item for item in dl_lhotse]
|
|
|
|
for n in range(num_examples):
|
|
for use_lhotse in [False, True]:
|
|
item = dataset_lhotse[n] if use_lhotse else dataset.__getitem__(n)
|
|
item_factory = dataset_factory.__getitem__(n)
|
|
for signal in data:
|
|
item_signal = item[signal].squeeze(0) if use_lhotse else item[signal]
|
|
golden_signal = data[signal][n]
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Test 1, use_lhotse={use_lhotse}: Signal {signal} item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 1, use_lhotse={use_lhotse}: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
assert np.allclose(
|
|
item_factory[signal], golden_signal, atol=atol
|
|
), f'Test 1, use_lhotse={use_lhotse}: Failed for factory example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 2
|
|
# Set target as the first channel of input_filepath and all files listed in target_filepath.
|
|
# In this case, the target will have 3 channels.
|
|
# Note: this is currently not supported by lhotse, so we only test the default dataset here.
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=[data_key['input_signal'], data_key['target_signal']],
|
|
target_channel_selector=0,
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
for n in range(num_examples):
|
|
item = dataset.__getitem__(n)
|
|
|
|
for signal in data:
|
|
item_signal = item[signal].cpu().detach().numpy()
|
|
golden_signal = data[signal][n]
|
|
if signal == 'target_signal':
|
|
# add the first channel of the input
|
|
golden_signal = np.concatenate([data['input_signal'][n][0:1, ...], golden_signal], axis=0)
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Signal {signal}: item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 2: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
@pytest.mark.unit
|
|
def test_audio_to_target_dataset_for_inference(self):
|
|
"""Test AudioWithTargetDataset when target_key is
|
|
not set, i.e., it is `None`. This is the case, e.g., when
|
|
running inference, and a target is not available.
|
|
|
|
In this use case, each line of the manifest file has the following format:
|
|
```
|
|
{
|
|
'input_filepath': 'path/to/input.wav',
|
|
'duration': duration_of_input,
|
|
}
|
|
```
|
|
"""
|
|
# Data setup
|
|
random_seed = 42
|
|
sample_rate = 16000
|
|
num_examples = 25
|
|
data_num_channels = {
|
|
'input_signal': 4,
|
|
}
|
|
data_min_duration = 2.0
|
|
data_max_duration = 8.0
|
|
data_key = {
|
|
'input_signal': 'input_filepath',
|
|
}
|
|
|
|
# Tolerance
|
|
atol = 1e-6
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
# Input and target signals have the same duration
|
|
data_duration = np.round(_rng.uniform(low=data_min_duration, high=data_max_duration, size=num_examples), 3)
|
|
data_duration_samples = np.floor(data_duration * sample_rate).astype(int)
|
|
|
|
data = dict()
|
|
for signal, num_channels in data_num_channels.items():
|
|
data[signal] = []
|
|
for n in range(num_examples):
|
|
if num_channels == 1:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(data_duration_samples[n]))
|
|
else:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(num_channels, data_duration_samples[n]))
|
|
data[signal].append(random_signal)
|
|
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
# Build metadata for manifest
|
|
metadata = []
|
|
for n in range(num_examples):
|
|
meta = dict()
|
|
for signal in data:
|
|
# filenames
|
|
signal_filename = f'{signal}_{n:02d}.wav'
|
|
# write audio files
|
|
sf.write(os.path.join(test_dir, signal_filename), data[signal][n].T, sample_rate, 'float')
|
|
# update metadata
|
|
meta[data_key[signal]] = signal_filename
|
|
meta['duration'] = data_duration[n]
|
|
metadata.append(meta)
|
|
|
|
# Save manifest
|
|
manifest_filepath = os.path.join(test_dir, 'manifest.json')
|
|
write_manifest(manifest_filepath, metadata)
|
|
|
|
# Test 1
|
|
# - No constraints on channels or duration
|
|
dataset = AudioToTargetDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=None, # target_signal will be empty
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
# Also test the corresponding factory
|
|
config = {
|
|
'manifest_filepath': manifest_filepath,
|
|
'input_key': data_key['input_signal'],
|
|
'target_key': None,
|
|
'sample_rate': sample_rate,
|
|
}
|
|
dataset_factory = audio_to_audio_dataset.get_audio_to_target_dataset(config)
|
|
|
|
# Prepare lhotse manifest
|
|
cuts_path = manifest_filepath.replace('.json', '_cuts.jsonl')
|
|
convert_manifest_nemo_to_lhotse(
|
|
input_manifest=manifest_filepath,
|
|
output_manifest=cuts_path,
|
|
input_key=data_key['input_signal'],
|
|
target_key=None,
|
|
)
|
|
|
|
# Prepare lhotse dataset
|
|
config_lhotse = {
|
|
'cuts_path': cuts_path,
|
|
'use_lhotse': True,
|
|
'sample_rate': sample_rate,
|
|
'batch_size': 1,
|
|
}
|
|
dl_lhotse = get_lhotse_dataloader_from_config(
|
|
OmegaConf.create(config_lhotse), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
|
)
|
|
dataset_lhotse = [item for item in dl_lhotse]
|
|
|
|
for n in range(num_examples):
|
|
|
|
for label in ['original', 'factory', 'lhotse']:
|
|
|
|
if label == 'original':
|
|
item = dataset.__getitem__(n)
|
|
elif label == 'factory':
|
|
item = dataset_factory.__getitem__(n)
|
|
elif label == 'lhotse':
|
|
item = dataset_lhotse[n]
|
|
else:
|
|
raise ValueError(f'Unknown label {label}')
|
|
|
|
# Check target is None
|
|
if 'target_signal' in item:
|
|
assert item['target_signal'].numel() == 0, f'{label}: target_signal is expected to be empty.'
|
|
|
|
# Check valid signals
|
|
for signal in data:
|
|
|
|
item_signal = item[signal].squeeze(0) if label == 'lhotse' else item[signal]
|
|
golden_signal = data[signal][n]
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'{label} -- Signal {signal}: item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'{label} -- Test 1: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
@pytest.mark.unit
|
|
def test_audio_to_target_with_reference_dataset(self):
|
|
"""Test AudioWithTargetWithReferenceDataset in different configurations.
|
|
|
|
1) reference synchronized with input and target
|
|
2) reference not synchronized
|
|
|
|
In this use case, each line of the manifest file has the following format:
|
|
```
|
|
{
|
|
'input_filepath': 'path/to/input.wav',
|
|
'target_filepath': 'path/to/path_to_target.wav',
|
|
'reference_filepath': 'path/to/path_to_reference.wav',
|
|
'duration': duration_of_input,
|
|
}
|
|
```
|
|
"""
|
|
# Data setup
|
|
random_seed = 42
|
|
sample_rate = 16000
|
|
num_examples = 25
|
|
data_num_channels = {
|
|
'input_signal': 4,
|
|
'target_signal': 2,
|
|
'reference_signal': 1,
|
|
}
|
|
data_min_duration = 2.0
|
|
data_max_duration = 8.0
|
|
data_key = {
|
|
'input_signal': 'input_filepath',
|
|
'target_signal': 'target_filepath',
|
|
'reference_signal': 'reference_filepath',
|
|
}
|
|
|
|
# Tolerance
|
|
atol = 1e-6
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
# Input and target signals have the same duration
|
|
data_duration = np.round(_rng.uniform(low=data_min_duration, high=data_max_duration, size=num_examples), 3)
|
|
data_duration_samples = np.floor(data_duration * sample_rate).astype(int)
|
|
|
|
data = dict()
|
|
for signal, num_channels in data_num_channels.items():
|
|
data[signal] = []
|
|
for n in range(num_examples):
|
|
if num_channels == 1:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(data_duration_samples[n]))
|
|
else:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(num_channels, data_duration_samples[n]))
|
|
data[signal].append(random_signal)
|
|
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
|
|
# Build metadata for manifest
|
|
metadata = []
|
|
|
|
for n in range(num_examples):
|
|
|
|
meta = dict()
|
|
|
|
for signal in data:
|
|
# filenames
|
|
signal_filename = f'{signal}_{n:02d}.wav'
|
|
|
|
# write audio files
|
|
sf.write(os.path.join(test_dir, signal_filename), data[signal][n].T, sample_rate, 'float')
|
|
|
|
# update metadata
|
|
meta[data_key[signal]] = signal_filename
|
|
|
|
meta['duration'] = data_duration[n]
|
|
metadata.append(meta)
|
|
|
|
# Save manifest
|
|
manifest_filepath = os.path.join(test_dir, 'manifest.json')
|
|
write_manifest(manifest_filepath, metadata)
|
|
|
|
# Test 1
|
|
# - No constraints on channels or duration
|
|
# - Reference is not synchronized with input and target, so whole reference signal will be loaded
|
|
dataset = AudioToTargetWithReferenceDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
reference_key=data_key['reference_signal'],
|
|
reference_is_synchronized=False,
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
# Also test the corresponding factory
|
|
config = {
|
|
'manifest_filepath': manifest_filepath,
|
|
'input_key': data_key['input_signal'],
|
|
'target_key': data_key['target_signal'],
|
|
'reference_key': data_key['reference_signal'],
|
|
'reference_is_synchronized': False,
|
|
'sample_rate': sample_rate,
|
|
}
|
|
dataset_factory = audio_to_audio_dataset.get_audio_to_target_with_reference_dataset(config)
|
|
|
|
for n in range(num_examples):
|
|
item = dataset.__getitem__(n)
|
|
item_factory = dataset_factory.__getitem__(n)
|
|
|
|
for signal in data:
|
|
item_signal = item[signal].cpu().detach().numpy()
|
|
golden_signal = data[signal][n]
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Signal {signal}: item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 1: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
item_factory_signal = item_factory[signal].cpu().detach().numpy()
|
|
assert np.allclose(
|
|
item_factory_signal, golden_signal, atol=atol
|
|
), f'Test 1: Failed for factory example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 2
|
|
# - Use fixed duration (random segment selection)
|
|
# - Reference is synchronized with input and target, so the same segment of reference signal will be loaded
|
|
audio_duration = 4.0
|
|
audio_duration_samples = int(np.floor(audio_duration * sample_rate))
|
|
dataset = AudioToTargetWithReferenceDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
reference_key=data_key['reference_signal'],
|
|
reference_is_synchronized=True,
|
|
sample_rate=sample_rate,
|
|
min_duration=audio_duration,
|
|
audio_duration=audio_duration,
|
|
random_offset=True,
|
|
)
|
|
|
|
filtered_examples = [n for n, val in enumerate(data_duration) if val >= audio_duration]
|
|
|
|
for n in range(len(dataset)):
|
|
item = dataset.__getitem__(n)
|
|
|
|
golden_start = golden_end = None
|
|
for signal in data:
|
|
item_signal = item[signal].cpu().detach().numpy()
|
|
full_golden_signal = data[signal][filtered_examples[n]]
|
|
|
|
# Find random segment using correlation on the first channel
|
|
# of the first signal, and then use it fixed for other signals
|
|
if golden_start is None:
|
|
golden_start = get_segment_start(signal=full_golden_signal[0, :], segment=item_signal[0, :])
|
|
golden_end = golden_start + audio_duration_samples
|
|
golden_signal = full_golden_signal[..., golden_start:golden_end]
|
|
|
|
# Test length is correct
|
|
assert (
|
|
item_signal.shape[-1] == audio_duration_samples
|
|
), f'Test 2: Signal {signal} length ({item_signal.shape[-1]}) not matching the expected length ({audio_duration_samples})'
|
|
|
|
# Test signal values
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 2: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 3
|
|
# - Use fixed duration (random segment selection)
|
|
# - Reference is not synchronized with input and target, so whole reference signal will be loaded
|
|
audio_duration = 4.0
|
|
audio_duration_samples = int(np.floor(audio_duration * sample_rate))
|
|
dataset = AudioToTargetWithReferenceDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
reference_key=data_key['reference_signal'],
|
|
reference_is_synchronized=False,
|
|
sample_rate=sample_rate,
|
|
min_duration=audio_duration,
|
|
audio_duration=audio_duration,
|
|
random_offset=True,
|
|
)
|
|
|
|
filtered_examples = [n for n, val in enumerate(data_duration) if val >= audio_duration]
|
|
|
|
for n in range(len(dataset)):
|
|
item = dataset.__getitem__(n)
|
|
|
|
golden_start = golden_end = None
|
|
for signal in data:
|
|
item_signal = item[signal].cpu().detach().numpy()
|
|
full_golden_signal = data[signal][filtered_examples[n]]
|
|
|
|
if signal == 'reference_signal':
|
|
# Complete signal is loaded for reference
|
|
golden_signal = full_golden_signal
|
|
else:
|
|
# Find random segment using correlation on the first channel
|
|
# of the first signal, and then use it fixed for other signals
|
|
if golden_start is None:
|
|
golden_start = get_segment_start(
|
|
signal=full_golden_signal[0, :], segment=item_signal[0, :]
|
|
)
|
|
golden_end = golden_start + audio_duration_samples
|
|
golden_signal = full_golden_signal[..., golden_start:golden_end]
|
|
|
|
# Test length is correct
|
|
assert (
|
|
item_signal.shape[-1] == audio_duration_samples
|
|
), f'Test 3: Signal {signal} length ({item_signal.shape[-1]}) not matching the expected length ({audio_duration_samples})'
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Signal {signal}: item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
# Test signal values
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 3: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 4:
|
|
# - Test collate_fn
|
|
batch_size = 16
|
|
batch = [dataset.__getitem__(n) for n in range(batch_size)]
|
|
_ = dataset.collate_fn(batch)
|
|
|
|
@pytest.mark.unit
|
|
def test_audio_to_target_with_embedding_dataset(self):
|
|
"""Test AudioWithTargetWithEmbeddingDataset.
|
|
|
|
In this use case, each line of the manifest file has the following format:
|
|
```
|
|
{
|
|
'input_filepath': 'path/to/input.wav',
|
|
'target_filepath': 'path/to/path_to_target.wav',
|
|
'embedding_filepath': 'path/to/path_to_embedding.npy',
|
|
'duration': duration_of_input,
|
|
}
|
|
```
|
|
"""
|
|
# Data setup
|
|
random_seed = 42
|
|
sample_rate = 16000
|
|
num_examples = 25
|
|
data_num_channels = {
|
|
'input_signal': 4,
|
|
'target_signal': 2,
|
|
'embedding_vector': 1,
|
|
}
|
|
data_min_duration = 2.0
|
|
data_max_duration = 8.0
|
|
embedding_length = 64 # 64-dimensional embedding vector
|
|
data_key = {
|
|
'input_signal': 'input_filepath',
|
|
'target_signal': 'target_filepath',
|
|
'embedding_vector': 'embedding_filepath',
|
|
}
|
|
|
|
# Tolerance
|
|
atol = 1e-6
|
|
|
|
# Generate random signals
|
|
_rng = np.random.default_rng(seed=random_seed)
|
|
|
|
# Input and target signals have the same duration
|
|
data_duration = np.round(_rng.uniform(low=data_min_duration, high=data_max_duration, size=num_examples), 3)
|
|
data_duration_samples = np.floor(data_duration * sample_rate).astype(int)
|
|
|
|
data = dict()
|
|
for signal, num_channels in data_num_channels.items():
|
|
data[signal] = []
|
|
for n in range(num_examples):
|
|
data_length = embedding_length if signal == 'embedding_vector' else data_duration_samples[n]
|
|
|
|
if num_channels == 1:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(data_length))
|
|
else:
|
|
random_signal = _rng.uniform(low=-0.5, high=0.5, size=(num_channels, data_length))
|
|
data[signal].append(random_signal)
|
|
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
|
|
# Build metadata for manifest
|
|
metadata = []
|
|
|
|
for n in range(num_examples):
|
|
|
|
meta = dict()
|
|
|
|
for signal in data:
|
|
if signal == 'embedding_vector':
|
|
signal_filename = f'{signal}_{n:02d}.npy'
|
|
np.save(os.path.join(test_dir, signal_filename), data[signal][n])
|
|
|
|
else:
|
|
# filenames
|
|
signal_filename = f'{signal}_{n:02d}.wav'
|
|
|
|
# write audio files
|
|
sf.write(os.path.join(test_dir, signal_filename), data[signal][n].T, sample_rate, 'float')
|
|
|
|
# update metadata
|
|
meta[data_key[signal]] = signal_filename
|
|
|
|
meta['duration'] = data_duration[n]
|
|
metadata.append(meta)
|
|
|
|
# Save manifest
|
|
manifest_filepath = os.path.join(test_dir, 'manifest.json')
|
|
write_manifest(manifest_filepath, metadata)
|
|
|
|
# Test 1
|
|
# - No constraints on channels or duration
|
|
dataset = AudioToTargetWithEmbeddingDataset(
|
|
manifest_filepath=manifest_filepath,
|
|
input_key=data_key['input_signal'],
|
|
target_key=data_key['target_signal'],
|
|
embedding_key=data_key['embedding_vector'],
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
# Also test the corresponding factory
|
|
config = {
|
|
'manifest_filepath': manifest_filepath,
|
|
'input_key': data_key['input_signal'],
|
|
'target_key': data_key['target_signal'],
|
|
'embedding_key': data_key['embedding_vector'],
|
|
'sample_rate': sample_rate,
|
|
}
|
|
dataset_factory = audio_to_audio_dataset.get_audio_to_target_with_embedding_dataset(config)
|
|
|
|
for n in range(num_examples):
|
|
item = dataset.__getitem__(n)
|
|
item_factory = dataset_factory.__getitem__(n)
|
|
|
|
for signal in data:
|
|
item_signal = item[signal].cpu().detach().numpy()
|
|
golden_signal = data[signal][n]
|
|
assert (
|
|
item_signal.shape == golden_signal.shape
|
|
), f'Signal {signal}: item shape {item_signal.shape} not matching reference shape {golden_signal.shape}'
|
|
assert np.allclose(
|
|
item_signal, golden_signal, atol=atol
|
|
), f'Test 1: Failed for example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
item_factory_signal = item_factory[signal].cpu().detach().numpy()
|
|
assert np.allclose(
|
|
item_factory_signal, golden_signal, atol=atol
|
|
), f'Test 1: Failed for factory example {n}, signal {signal} (random seed {random_seed})'
|
|
|
|
# Test 2:
|
|
# - Test collate_fn
|
|
batch_size = 16
|
|
batch = [dataset.__getitem__(n) for n in range(batch_size)]
|
|
_ = dataset.collate_fn(batch)
|