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
2401 lines
89 KiB
Python
2401 lines
89 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 itertools
|
||
import multiprocessing
|
||
import os
|
||
import random
|
||
from typing import Dict, Iterable, List, Optional, Tuple, Union
|
||
|
||
import librosa
|
||
import numpy as np
|
||
import soundfile as sf
|
||
from numpy.random import default_rng
|
||
from omegaconf import DictConfig, OmegaConf
|
||
from scipy.signal import convolve
|
||
from scipy.spatial.transform import Rotation
|
||
from tqdm import tqdm
|
||
|
||
from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
|
||
from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
|
||
from nemo.collections.audio.parts.utils.audio import db2mag, generate_approximate_noise_field, mag2db, pow2db, rms
|
||
from nemo.utils import logging
|
||
|
||
try:
|
||
import pyroomacoustics as pra
|
||
|
||
PRA = True
|
||
except ImportError:
|
||
PRA = False
|
||
|
||
try:
|
||
import h5py
|
||
|
||
HAS_H5PY = True
|
||
except ImportError:
|
||
HAS_H5PY = False
|
||
|
||
|
||
def check_angle(key: str, val: Union[float, Iterable[float]]) -> bool:
|
||
"""Check if the angle value is within the expected range. Input
|
||
values are in degrees.
|
||
|
||
Note:
|
||
azimuth: angle between a projection on the horizontal (xy) plane and
|
||
positive x axis. Increases counter-clockwise. Range: [-180, 180].
|
||
elevation: angle between a vector an its projection on the horizontal (xy) plane.
|
||
Positive above, negative below, i.e., north=+90, south=-90. Range: [-90, 90]
|
||
yaw: rotation around the z axis. Defined accoding to right-hand rule.
|
||
Range: [-180, 180]
|
||
pitch: rotation around the yʹ axis. Defined accoding to right-hand rule.
|
||
Range: [-90, 90]
|
||
roll: rotation around the xʺ axis. Defined accoding to right-hand rule.
|
||
Range: [-180, 180]
|
||
|
||
Args:
|
||
key: angle type
|
||
val: values in degrees
|
||
|
||
Returns:
|
||
True if all values are within the expected range.
|
||
"""
|
||
if np.isscalar(val):
|
||
min_val = max_val = val
|
||
else:
|
||
min_val = min(val)
|
||
max_val = max(val)
|
||
|
||
if key == 'azimuth' and -180 <= min_val <= max_val <= 180:
|
||
return True
|
||
if key == 'elevation' and -90 <= min_val <= max_val <= 90:
|
||
return True
|
||
if key == 'yaw' and -180 <= min_val <= max_val <= 180:
|
||
return True
|
||
if key == 'pitch' and -90 <= min_val <= max_val <= 90:
|
||
return True
|
||
if key == 'roll' and -180 <= min_val <= max_val <= 180:
|
||
return True
|
||
|
||
raise ValueError(f'Invalid value for angle {key} = {val}')
|
||
|
||
|
||
def wrap_to_180(angle: float) -> float:
|
||
"""Wrap an angle to range ±180 degrees.
|
||
|
||
Args:
|
||
angle: angle in degrees
|
||
|
||
Returns:
|
||
Angle in degrees wrapped to ±180 degrees.
|
||
"""
|
||
return angle - np.floor(angle / 360 + 1 / 2) * 360
|
||
|
||
|
||
class ArrayGeometry(object):
|
||
"""A class to simplify handling of array geometry.
|
||
|
||
Supports translation and rotation of the array and calculation of
|
||
spherical coordinates of a given point relative to the internal
|
||
coordinate system of the array.
|
||
|
||
Args:
|
||
mic_positions: 3D coordinates, with shape (num_mics, 3)
|
||
center: optional position of the center of the array. Defaults to the average of the coordinates.
|
||
internal_cs: internal coordinate system for the array relative to the global coordinate system.
|
||
Defaults to (x, y, z), and is rotated with the array.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
mic_positions: Union[np.ndarray, List],
|
||
center: Optional[np.ndarray] = None,
|
||
internal_cs: Optional[np.ndarray] = None,
|
||
):
|
||
if isinstance(mic_positions, Iterable):
|
||
mic_positions = np.array(mic_positions)
|
||
|
||
if not mic_positions.ndim == 2:
|
||
raise ValueError(
|
||
f'Expecting a 2D array specifying mic positions, but received {mic_positions.ndim}-dim array'
|
||
)
|
||
|
||
if not mic_positions.shape[1] == 3:
|
||
raise ValueError(f'Expecting 3D positions, but received {mic_positions.shape[1]}-dim positions')
|
||
|
||
mic_positions_center = np.mean(mic_positions, axis=0)
|
||
self.centered_positions = mic_positions - mic_positions_center
|
||
self.center = mic_positions_center if center is None else center
|
||
|
||
# Internal coordinate system
|
||
if internal_cs is None:
|
||
# Initially aligned with the global
|
||
self.internal_cs = np.eye(3)
|
||
else:
|
||
self.internal_cs = internal_cs
|
||
|
||
@property
|
||
def num_mics(self):
|
||
"""Return the number of microphones for the current array."""
|
||
return self.centered_positions.shape[0]
|
||
|
||
@property
|
||
def positions(self):
|
||
"""Absolute positions of the microphones."""
|
||
return self.centered_positions + self.center
|
||
|
||
@property
|
||
def internal_positions(self):
|
||
"""Positions in the internal coordinate system."""
|
||
return np.matmul(self.centered_positions, self.internal_cs.T)
|
||
|
||
@property
|
||
def radius(self):
|
||
"""Radius of the array, relative to the center."""
|
||
return max(np.linalg.norm(self.centered_positions, axis=1))
|
||
|
||
@staticmethod
|
||
def get_rotation(yaw: float = 0, pitch: float = 0, roll: float = 0) -> Rotation:
|
||
"""Get a Rotation object for given angles.
|
||
|
||
All angles are defined according to the right-hand rule.
|
||
|
||
Args:
|
||
yaw: rotation around the z axis
|
||
pitch: rotation around the yʹ axis
|
||
roll: rotation around the xʺ axis
|
||
|
||
Returns:
|
||
A rotation object constructed using the provided angles.
|
||
"""
|
||
check_angle('yaw', yaw)
|
||
check_angle('pitch', pitch)
|
||
check_angle('roll', roll)
|
||
|
||
return Rotation.from_euler('ZYX', [yaw, pitch, roll], degrees=True)
|
||
|
||
def translate(self, to: np.ndarray):
|
||
"""Translate the array center to a new point.
|
||
|
||
Translation does not change the centered positions or the internal coordinate system.
|
||
|
||
Args:
|
||
to: 3D point, shape (3,)
|
||
"""
|
||
self.center = to
|
||
|
||
def rotate(self, yaw: float = 0, pitch: float = 0, roll: float = 0):
|
||
"""Apply rotation on the mic array.
|
||
|
||
This rotates the centered microphone positions and the internal
|
||
coordinate system, it doesn't change the center of the array.
|
||
|
||
All angles are defined according to the right-hand rule.
|
||
For example, this means that a positive pitch will result in a rotation from z
|
||
to x axis, which will result in a reduced elevation with respect to the global
|
||
horizontal plane.
|
||
|
||
Args:
|
||
yaw: rotation around the z axis
|
||
pitch: rotation around the yʹ axis
|
||
roll: rotation around the xʺ axis
|
||
"""
|
||
# construct rotation using TB angles
|
||
rotation = self.get_rotation(yaw=yaw, pitch=pitch, roll=roll)
|
||
|
||
# rotate centered positions
|
||
self.centered_positions = rotation.apply(self.centered_positions)
|
||
|
||
# apply the same transformation on the internal coordinate system
|
||
self.internal_cs = rotation.apply(self.internal_cs)
|
||
|
||
def new_rotated_array(self, yaw: float = 0, pitch: float = 0, roll: float = 0):
|
||
"""Create a new array by rotating this array.
|
||
|
||
Args:
|
||
yaw: rotation around the z axis
|
||
pitch: rotation around the yʹ axis
|
||
roll: rotation around the xʺ axis
|
||
|
||
Returns:
|
||
A new ArrayGeometry object constructed using the provided angles.
|
||
"""
|
||
new_array = ArrayGeometry(mic_positions=self.positions, center=self.center, internal_cs=self.internal_cs)
|
||
new_array.rotate(yaw=yaw, pitch=pitch, roll=roll)
|
||
return new_array
|
||
|
||
def spherical_relative_to_array(
|
||
self, point: np.ndarray, use_internal_cs: bool = True
|
||
) -> Tuple[float, float, float]:
|
||
"""Return spherical coordinates of a point relative to the internal coordinate system.
|
||
|
||
Args:
|
||
point: 3D coordinate, shape (3,)
|
||
use_internal_cs: Calculate position relative to the internal coordinate system.
|
||
If `False`, the positions will be calculated relative to the
|
||
external coordinate system centered at `self.center`.
|
||
|
||
Returns:
|
||
A tuple (distance, azimuth, elevation) relative to the mic array.
|
||
"""
|
||
rel_position = point - self.center
|
||
distance = np.linalg.norm(rel_position)
|
||
|
||
if use_internal_cs:
|
||
# transform from the absolute coordinate system to the internal coordinate system
|
||
rel_position = np.matmul(self.internal_cs, rel_position)
|
||
|
||
# get azimuth
|
||
azimuth = np.arctan2(rel_position[1], rel_position[0]) / np.pi * 180
|
||
# get elevation
|
||
elevation = np.arcsin(rel_position[2] / distance) / np.pi * 180
|
||
|
||
return distance, azimuth, elevation
|
||
|
||
def __str__(self):
|
||
with np.printoptions(precision=3, suppress=True):
|
||
desc = f"{type(self)}:\ncenter =\n{self.center}\ncentered positions =\n{self.centered_positions}\nradius = \n{self.radius:.3}\nabsolute positions =\n{self.positions}\ninternal coordinate system =\n{self.internal_cs}\n\n"
|
||
return desc
|
||
|
||
def plot(self, elev=30, azim=-55, mic_size=25):
|
||
"""Plot microphone positions.
|
||
|
||
Args:
|
||
elev: elevation for the view of the plot
|
||
azim: azimuth for the view of the plot
|
||
mic_size: size of the microphone marker in the plot
|
||
"""
|
||
import matplotlib.pyplot as plt
|
||
|
||
fig = plt.figure()
|
||
ax = fig.add_subplot(projection='3d')
|
||
|
||
# show mic positions
|
||
for m in range(self.num_mics):
|
||
# show mic
|
||
ax.scatter(
|
||
self.positions[m, 0],
|
||
self.positions[m, 1],
|
||
self.positions[m, 2],
|
||
marker='o',
|
||
c='black',
|
||
s=mic_size,
|
||
depthshade=False,
|
||
)
|
||
# add label
|
||
ax.text(self.positions[m, 0], self.positions[m, 1], self.positions[m, 2], str(m), c='red', zorder=10)
|
||
|
||
# show the internal coordinate system
|
||
ax.quiver(
|
||
self.center[0],
|
||
self.center[1],
|
||
self.center[2],
|
||
self.internal_cs[:, 0],
|
||
self.internal_cs[:, 1],
|
||
self.internal_cs[:, 2],
|
||
length=self.radius,
|
||
label='internal cs',
|
||
normalize=False,
|
||
linestyle=':',
|
||
linewidth=1.0,
|
||
)
|
||
for dim, label in enumerate(['x′', 'y′', 'z′']):
|
||
label_pos = self.center + self.radius * self.internal_cs[dim]
|
||
ax.text(label_pos[0], label_pos[1], label_pos[2], label, tuple(self.internal_cs[dim]), c='blue')
|
||
try:
|
||
# Unfortunately, equal aspect ratio has been added very recently to Axes3D
|
||
ax.set_aspect('equal')
|
||
except NotImplementedError:
|
||
logging.warning('Equal aspect ratio not supported by Axes3D')
|
||
# Set view
|
||
ax.view_init(elev=elev, azim=azim)
|
||
# Set reasonable limits for all axes, even for the case of an unequal aspect ratio
|
||
ax.set_xlim([self.center[0] - self.radius, self.center[0] + self.radius])
|
||
ax.set_ylim([self.center[1] - self.radius, self.center[1] + self.radius])
|
||
ax.set_zlim([self.center[2] - self.radius, self.center[2] + self.radius])
|
||
|
||
ax.set_xlabel('x/m')
|
||
ax.set_ylabel('y/m')
|
||
ax.set_zlabel('z/m')
|
||
ax.set_title('Microphone positions')
|
||
ax.legend()
|
||
plt.show()
|
||
|
||
|
||
def convert_placement_to_range(
|
||
placement: dict, room_dim: Iterable[float], object_radius: float = 0
|
||
) -> List[List[float]]:
|
||
"""Given a placement dictionary, return ranges for each dimension.
|
||
|
||
Args:
|
||
placement: dictionary containing x, y, height, and min_to_wall
|
||
room_dim: dimensions of the room, shape (3,)
|
||
object_radius: radius of the object to be placed
|
||
|
||
Returns
|
||
List with a range of values for each dimensions.
|
||
"""
|
||
if not np.all(np.array(room_dim) > 0):
|
||
raise ValueError(f'Room dimensions must be positive: {room_dim}')
|
||
|
||
if object_radius < 0:
|
||
raise ValueError(f'Object radius must be non-negative: {object_radius}')
|
||
|
||
placement_range = [None] * 3
|
||
min_to_wall = placement.get('min_to_wall', 0)
|
||
|
||
if min_to_wall < 0:
|
||
raise ValueError(f'Min distance to wall must be positive: {min_to_wall}')
|
||
|
||
for idx, key in enumerate(['x', 'y', 'height']):
|
||
# Room dimension
|
||
dim = room_dim[idx]
|
||
# Construct the range
|
||
val = placement.get(key)
|
||
if val is None:
|
||
# No constrained specified on the coordinate of the mic center
|
||
min_val, max_val = 0, dim
|
||
elif np.isscalar(val):
|
||
min_val = max_val = val
|
||
else:
|
||
if len(val) != 2:
|
||
raise ValueError(f'Invalid value for placement for dim {idx}/{key}: {str(placement)}')
|
||
min_val, max_val = val
|
||
|
||
# Make sure the array is not too close to a wall
|
||
min_val = max(min_val, min_to_wall + object_radius)
|
||
max_val = min(max_val, dim - min_to_wall - object_radius)
|
||
|
||
if min_val > max_val or min(min_val, max_val) < 0:
|
||
raise ValueError(f'Invalid range dim {idx}/{key}: min={min_val}, max={max_val}')
|
||
|
||
placement_range[idx] = [min_val, max_val]
|
||
|
||
return placement_range
|
||
|
||
|
||
class RIRCorpusGenerator(object):
|
||
"""Creates a corpus of RIRs based on a defined configuration of rooms and microphone array.
|
||
|
||
RIRs are generated using `generate` method.
|
||
"""
|
||
|
||
def __init__(self, cfg: DictConfig):
|
||
"""
|
||
Args:
|
||
cfg: dictionary with parameters of the simulation
|
||
"""
|
||
logging.info("Initialize RIRCorpusGenerator")
|
||
self._cfg = cfg
|
||
self.check_cfg()
|
||
|
||
@property
|
||
def cfg(self):
|
||
"""Property holding the internal config of the object.
|
||
|
||
Note:
|
||
Changes to this config are not reflected in the state of the object.
|
||
Please create a new model with the updated config.
|
||
"""
|
||
return self._cfg
|
||
|
||
@property
|
||
def sample_rate(self):
|
||
return self._cfg.sample_rate
|
||
|
||
@cfg.setter
|
||
def cfg(self, cfg):
|
||
"""Property holding the internal config of the object.
|
||
|
||
Note:
|
||
Changes to this config are not reflected in the state of the object.
|
||
Please create a new model with the updated config.
|
||
"""
|
||
self._cfg = cfg
|
||
|
||
def check_cfg(self):
|
||
"""
|
||
Checks provided configuration to ensure it has the minimal required
|
||
configuration the values are in a reasonable range.
|
||
"""
|
||
# sample rate
|
||
sample_rate = self.cfg.get('sample_rate')
|
||
if sample_rate is None:
|
||
raise ValueError('Sample rate not provided.')
|
||
elif sample_rate < 0:
|
||
raise ValueError(f'Sample rate must to be positive: {sample_rate}')
|
||
|
||
# room configuration
|
||
room_cfg = self.cfg.get('room')
|
||
if room_cfg is None:
|
||
raise ValueError('Room configuration not provided')
|
||
|
||
if room_cfg.get('num') is None:
|
||
raise ValueError('Number of rooms per subset not provided')
|
||
|
||
if room_cfg.get('dim') is None:
|
||
raise ValueError('Room dimensions not provided')
|
||
|
||
for idx, key in enumerate(['width', 'length', 'height']):
|
||
dim = room_cfg.dim.get(key)
|
||
|
||
if dim is None:
|
||
# not provided
|
||
raise ValueError(f'Room {key} needs to be a scalar or a range, currently it is None')
|
||
elif np.isscalar(dim) and dim <= 0:
|
||
# fixed dimension
|
||
raise ValueError(f'A fixed dimension must be positive for {key}: {dim}')
|
||
elif len(dim) != 2 or not 0 < dim[0] < dim[1]:
|
||
# not a valid range
|
||
raise ValueError(f'Range must be specified with two positive increasing elements for {key}: {dim}')
|
||
|
||
rt60 = room_cfg.get('rt60')
|
||
if rt60 is None:
|
||
# not provided
|
||
raise ValueError('RT60 needs to be a scalar or a range, currently it is None')
|
||
elif np.isscalar(rt60) and rt60 <= 0:
|
||
# fixed dimension
|
||
raise ValueError(f'RT60 must be positive: {rt60}')
|
||
elif len(rt60) != 2 or not 0 < rt60[0] < rt60[1]:
|
||
# not a valid range
|
||
raise ValueError(f'RT60 range must be specified with two positive increasing elements: {rt60}')
|
||
|
||
# mic array
|
||
mic_cfg = self.cfg.get('mic_array')
|
||
if mic_cfg is None:
|
||
raise ValueError('Mic configuration not provided')
|
||
|
||
if mic_cfg.get('positions') == 'random':
|
||
# Only num_mics and placement are required
|
||
mic_cfg_keys = ['num_mics', 'placement']
|
||
else:
|
||
mic_cfg_keys = ['positions', 'placement', 'orientation']
|
||
|
||
for key in mic_cfg_keys:
|
||
if key not in mic_cfg:
|
||
raise ValueError(f'Mic array {key} not provided')
|
||
|
||
# source
|
||
source_cfg = self.cfg.get('source')
|
||
if source_cfg is None:
|
||
raise ValueError('Source configuration not provided')
|
||
|
||
if source_cfg.get('num') is None:
|
||
raise ValueError('Number of sources per room not provided')
|
||
elif source_cfg.num <= 0:
|
||
raise ValueError(f'Number of sources must be positive: {source_cfg.num}')
|
||
|
||
if 'placement' not in source_cfg:
|
||
raise ValueError('Source placement dictionary not provided')
|
||
|
||
# anechoic
|
||
if self.cfg.get('anechoic') is None:
|
||
raise ValueError('Anechoic configuratio not provided.')
|
||
|
||
def generate_room_params(self) -> dict:
|
||
"""Generate randomized room parameters based on the provided
|
||
configuration.
|
||
"""
|
||
# Prepare room sim parameters
|
||
if not PRA:
|
||
raise ImportError('pyroomacoustics is required for room simulation')
|
||
|
||
room_cfg = self.cfg.room
|
||
|
||
# Prepare rt60
|
||
if room_cfg.rt60 is None:
|
||
raise ValueError('Room RT60 needs to be a scalar or a range, currently it is None')
|
||
|
||
if np.isscalar(room_cfg.rt60):
|
||
assert room_cfg.rt60 > 0, f'RT60 should be positive: {room_cfg.rt60}'
|
||
rt60 = room_cfg.rt60
|
||
elif len(room_cfg.rt60) == 2:
|
||
assert (
|
||
0 < room_cfg.rt60[0] <= room_cfg.rt60[1]
|
||
), f'Expecting two non-decreasing values for RT60, received {room_cfg.rt60}'
|
||
rt60 = self.random.uniform(low=room_cfg.rt60[0], high=room_cfg.rt60[1])
|
||
else:
|
||
raise ValueError(f'Unexpected value for RT60: {room_cfg.rt60}')
|
||
|
||
# Generate a room with random dimensions
|
||
num_retries = self.cfg.get('num_retries', 20)
|
||
|
||
for n in range(num_retries):
|
||
|
||
# width, length, height
|
||
room_dim = np.zeros(3)
|
||
|
||
# prepare dimensions
|
||
for idx, key in enumerate(['width', 'length', 'height']):
|
||
# get configured dimension
|
||
dim = room_cfg.dim[key]
|
||
|
||
# set a value
|
||
if dim is None:
|
||
raise ValueError(f'Room {key} needs to be a scalar or a range, currently it is None')
|
||
elif np.isscalar(dim):
|
||
assert dim > 0, f'Dimension should be positive for {key}: {dim}'
|
||
room_dim[idx] = dim
|
||
elif len(dim) == 2:
|
||
assert 0 < dim[0] <= dim[1], f'Expecting two non-decreasing values for {key}, received {dim}'
|
||
# Reduce dimension if the previous attempt failed
|
||
room_dim[idx] = self.random.uniform(low=dim[0], high=dim[1] - n * (dim[1] - dim[0]) / num_retries)
|
||
else:
|
||
raise ValueError(f'Unexpected value for {key}: {dim}')
|
||
|
||
try:
|
||
# Get parameters from size and RT60
|
||
room_absorption, room_max_order = pra.inverse_sabine(rt60, room_dim)
|
||
break
|
||
except Exception as e:
|
||
logging.debug('Inverse sabine failed: %s', str(e))
|
||
# Inverse sabine may fail if the room is too large for the selected RT60.
|
||
# Try again by generate a smaller room.
|
||
room_absorption = room_max_order = None
|
||
continue
|
||
|
||
if room_absorption is None or room_max_order is None:
|
||
raise RuntimeError(f'Evaluation of parameters failed for RT60 {rt60}s and room size {room_dim}.')
|
||
|
||
# Return the required values
|
||
room_params = {
|
||
'dim': room_dim,
|
||
'absorption': room_absorption,
|
||
'max_order': room_max_order,
|
||
'rt60_theoretical': rt60,
|
||
'anechoic_absorption': self.cfg.anechoic.absorption,
|
||
'anechoic_max_order': self.cfg.anechoic.max_order,
|
||
'sample_rate': self.cfg.sample_rate,
|
||
}
|
||
return room_params
|
||
|
||
def generate_array(self, room_dim: Iterable[float]) -> ArrayGeometry:
|
||
"""Generate array placement for the current room and config.
|
||
|
||
Args:
|
||
room_dim: dimensions of the room, [width, length, height]
|
||
|
||
Returns:
|
||
Randomly placed microphone array.
|
||
"""
|
||
mic_cfg = self.cfg.mic_array
|
||
|
||
if mic_cfg.positions == 'random':
|
||
# Create a radom set of microphones
|
||
num_mics = mic_cfg.num_mics
|
||
mic_positions = []
|
||
|
||
# Each microphone is placed individually
|
||
placement_range = convert_placement_to_range(
|
||
placement=mic_cfg.placement, room_dim=room_dim, object_radius=0
|
||
)
|
||
|
||
# Randomize mic placement
|
||
for m in range(num_mics):
|
||
position_m = [None] * 3
|
||
for idx in range(3):
|
||
position_m[idx] = self.random.uniform(low=placement_range[idx][0], high=placement_range[idx][1])
|
||
mic_positions.append(position_m)
|
||
|
||
mic_array = ArrayGeometry(mic_positions)
|
||
|
||
else:
|
||
mic_array = ArrayGeometry(mic_cfg.positions)
|
||
|
||
# Randomize center placement
|
||
center = np.zeros(3)
|
||
placement_range = convert_placement_to_range(
|
||
placement=mic_cfg.placement, room_dim=room_dim, object_radius=mic_array.radius
|
||
)
|
||
|
||
for idx in range(len(center)):
|
||
center[idx] = self.random.uniform(low=placement_range[idx][0], high=placement_range[idx][1])
|
||
|
||
# Place the array at the configured center point
|
||
mic_array.translate(to=center)
|
||
|
||
# Randomize orientation
|
||
orientation = dict()
|
||
for key in ['yaw', 'roll', 'pitch']:
|
||
# angle for current orientation
|
||
angle = mic_cfg.orientation[key]
|
||
|
||
if angle is None:
|
||
raise ValueError(f'Mic array {key} should be a scalar or a range, currently it is set to None.')
|
||
|
||
# check it's within the expected range
|
||
check_angle(key, angle)
|
||
|
||
if np.isscalar(angle):
|
||
orientation[key] = angle
|
||
elif len(angle) == 2:
|
||
assert angle[0] <= angle[1], f"Expecting two non-decreasing values for {key}, received {angle}"
|
||
# generate integer values, for easier bucketing, if necessary
|
||
orientation[key] = self.random.uniform(low=angle[0], high=angle[1])
|
||
else:
|
||
raise ValueError(f'Unexpected value for orientation {key}: {angle}')
|
||
|
||
# Rotate the array to match the selected orientation
|
||
mic_array.rotate(**orientation)
|
||
|
||
return mic_array
|
||
|
||
def generate_source_position(self, room_dim: Iterable[float]) -> List[List[float]]:
|
||
"""Generate position for all sources in a room.
|
||
|
||
Args:
|
||
room_dim: dimensions of a 3D shoebox room
|
||
|
||
Returns:
|
||
List of source positions, with each position characterized with a 3D coordinate
|
||
"""
|
||
source_cfg = self.cfg.source
|
||
placement_range = convert_placement_to_range(placement=source_cfg.placement, room_dim=room_dim)
|
||
source_position = []
|
||
|
||
for n in range(source_cfg.num):
|
||
# generate a random point withing the range
|
||
s_pos = [None] * 3
|
||
for idx in range(len(s_pos)):
|
||
s_pos[idx] = self.random.uniform(low=placement_range[idx][0], high=placement_range[idx][1])
|
||
source_position.append(s_pos)
|
||
|
||
return source_position
|
||
|
||
def generate(self):
|
||
"""Generate RIR corpus.
|
||
|
||
This method will prepare randomized examples based on the current configuration,
|
||
run room simulations and save results to output_dir.
|
||
"""
|
||
logging.info("Generate RIR corpus")
|
||
|
||
# Initialize
|
||
self.random = default_rng(seed=self.cfg.random_seed)
|
||
|
||
# Prepare output dir
|
||
output_dir = self.cfg.output_dir
|
||
if output_dir.endswith('.yaml'):
|
||
output_dir = output_dir[:-5]
|
||
|
||
# Create absolute path
|
||
logging.info('Output dir set to: %s', output_dir)
|
||
|
||
# Generate all cases
|
||
for subset, num_rooms in self.cfg.room.num.items():
|
||
|
||
output_dir_subset = os.path.join(output_dir, subset)
|
||
examples = []
|
||
|
||
if not os.path.exists(output_dir_subset):
|
||
logging.info('Creating output directory: %s', output_dir_subset)
|
||
os.makedirs(output_dir_subset)
|
||
elif os.path.isdir(output_dir_subset) and len(os.listdir(output_dir_subset)) > 0:
|
||
raise RuntimeError(f'Output directory {output_dir_subset} is not empty.')
|
||
|
||
# Generate examples
|
||
for n_room in range(num_rooms):
|
||
|
||
# room info
|
||
room_params = self.generate_room_params()
|
||
|
||
# array placement
|
||
mic_array = self.generate_array(room_params['dim'])
|
||
|
||
# source placement
|
||
source_position = self.generate_source_position(room_params['dim'])
|
||
|
||
# file name for the file
|
||
room_filepath = os.path.join(output_dir_subset, f'{subset}_room_{n_room:06d}.h5')
|
||
|
||
# prepare example
|
||
example = {
|
||
'room_params': room_params,
|
||
'mic_array': mic_array,
|
||
'source_position': source_position,
|
||
'room_filepath': room_filepath,
|
||
}
|
||
examples.append(example)
|
||
|
||
# Simulation
|
||
if (num_workers := self.cfg.get('num_workers')) is None:
|
||
num_workers = os.cpu_count() - 1
|
||
|
||
if num_workers > 1:
|
||
logging.info(f'Simulate using {num_workers} workers')
|
||
with multiprocessing.Pool(processes=num_workers) as pool:
|
||
metadata = list(tqdm(pool.imap(simulate_room_kwargs, examples), total=len(examples)))
|
||
|
||
else:
|
||
logging.info('Simulate using a single worker')
|
||
metadata = []
|
||
for example in tqdm(examples, total=len(examples)):
|
||
metadata.append(simulate_room(**example))
|
||
|
||
# Save manifest
|
||
manifest_filepath = os.path.join(output_dir, f'{subset}_manifest.json')
|
||
|
||
if os.path.exists(manifest_filepath) and os.path.isfile(manifest_filepath):
|
||
raise RuntimeError(f'Manifest config file exists: {manifest_filepath}')
|
||
|
||
# Make all paths in the manifest relative to the output dir
|
||
for data in metadata:
|
||
data['room_filepath'] = os.path.relpath(data['room_filepath'], start=output_dir)
|
||
|
||
write_manifest(manifest_filepath, metadata)
|
||
|
||
# Generate plots with information about generated data
|
||
plot_filepath = os.path.join(output_dir, f'{subset}_info.png')
|
||
|
||
if os.path.exists(plot_filepath) and os.path.isfile(plot_filepath):
|
||
raise RuntimeError(f'Plot file exists: {plot_filepath}')
|
||
|
||
plot_rir_manifest_info(manifest_filepath, plot_filepath=plot_filepath)
|
||
|
||
# Save used configuration for reference
|
||
config_filepath = os.path.join(output_dir, 'config.yaml')
|
||
if os.path.exists(config_filepath) and os.path.isfile(config_filepath):
|
||
raise RuntimeError(f'Output config file exists: {config_filepath}')
|
||
|
||
OmegaConf.save(self.cfg, config_filepath, resolve=True)
|
||
|
||
|
||
def simulate_room_kwargs(kwargs: dict) -> dict:
|
||
"""Wrapper around `simulate_room` to handle kwargs.
|
||
|
||
`pool.map(simulate_room_kwargs, examples)` would be
|
||
equivalent to `pool.starstarmap(simulate_room, examples)`
|
||
if `starstarmap` would exist.
|
||
|
||
Args:
|
||
kwargs: kwargs that are forwarded to `simulate_room`
|
||
|
||
Returns:
|
||
Dictionary with metadata, see `simulate_room`
|
||
"""
|
||
return simulate_room(**kwargs)
|
||
|
||
|
||
def simulate_room(
|
||
room_params: dict,
|
||
mic_array: ArrayGeometry,
|
||
source_position: Iterable[Iterable[float]],
|
||
room_filepath: str,
|
||
) -> dict:
|
||
"""Simulate room
|
||
|
||
Args:
|
||
room_params: parameters of the room to be simulated
|
||
mic_array: defines positions of the microphones
|
||
source_positions: positions for all sources to be simulated
|
||
room_filepath: results are saved to this path
|
||
|
||
Returns:
|
||
Dictionary with metadata based on simulation setup
|
||
and simulation results. Used to create the corresponding
|
||
manifest file.
|
||
"""
|
||
# room with the selected parameters
|
||
room_sim = pra.ShoeBox(
|
||
room_params['dim'],
|
||
fs=room_params['sample_rate'],
|
||
materials=pra.Material(room_params['absorption']),
|
||
max_order=room_params['max_order'],
|
||
)
|
||
|
||
# same geometry for generating anechoic responses
|
||
room_anechoic = pra.ShoeBox(
|
||
room_params['dim'],
|
||
fs=room_params['sample_rate'],
|
||
materials=pra.Material(room_params['anechoic_absorption']),
|
||
max_order=room_params['anechoic_max_order'],
|
||
)
|
||
|
||
# Compute RIRs
|
||
for room in [room_sim, room_anechoic]:
|
||
# place the array
|
||
room.add_microphone_array(mic_array.positions.T)
|
||
|
||
# place the sources
|
||
for s_pos in source_position:
|
||
room.add_source(s_pos)
|
||
|
||
# generate RIRs
|
||
room.compute_rir()
|
||
|
||
# Get metadata for sources
|
||
source_distance = []
|
||
source_azimuth = []
|
||
source_elevation = []
|
||
for s_pos in source_position:
|
||
distance, azimuth, elevation = mic_array.spherical_relative_to_array(s_pos)
|
||
source_distance.append(distance)
|
||
source_azimuth.append(azimuth)
|
||
source_elevation.append(elevation)
|
||
|
||
# RIRs
|
||
rir_dataset = {
|
||
'rir': convert_rir_to_multichannel(room_sim.rir),
|
||
'anechoic': convert_rir_to_multichannel(room_anechoic.rir),
|
||
}
|
||
|
||
# Prepare metadata dict and return
|
||
metadata = {
|
||
'room_filepath': room_filepath,
|
||
'sample_rate': room_params['sample_rate'],
|
||
'dim': room_params['dim'],
|
||
'rir_absorption': room_params['absorption'],
|
||
'rir_max_order': room_params['max_order'],
|
||
'rir_rt60_theory': room_sim.rt60_theory(),
|
||
'rir_rt60_measured': room_sim.measure_rt60().mean(axis=0), # average across mics for each source
|
||
'anechoic_rt60_theory': room_anechoic.rt60_theory(),
|
||
'anechoic_rt60_measured': room_anechoic.measure_rt60().mean(axis=0), # average across mics for each source
|
||
'anechoic_absorption': room_params['anechoic_absorption'],
|
||
'anechoic_max_order': room_params['anechoic_max_order'],
|
||
'mic_positions': mic_array.positions,
|
||
'mic_center': mic_array.center,
|
||
'source_position': source_position,
|
||
'source_distance': source_distance,
|
||
'source_azimuth': source_azimuth,
|
||
'source_elevation': source_elevation,
|
||
'num_sources': len(source_position),
|
||
}
|
||
|
||
# Save simulated RIR
|
||
save_rir_simulation(room_filepath, rir_dataset, metadata)
|
||
|
||
return convert_numpy_to_serializable(metadata)
|
||
|
||
|
||
def save_rir_simulation(filepath: str, rir_dataset: Dict[str, List[np.array]], metadata: dict):
|
||
"""Save simulated RIRs and metadata.
|
||
|
||
Args:
|
||
filepath: Path to the file where the data will be saved.
|
||
rir_dataset: Dictionary with RIR data. Each item is a set of multi-channel RIRs.
|
||
metadata: Dictionary with related metadata.
|
||
"""
|
||
if not HAS_H5PY:
|
||
raise ImportError("Install h5py to use save_rir_simulation")
|
||
if os.path.exists(filepath):
|
||
raise RuntimeError(f'Output file exists: {filepath}')
|
||
|
||
num_sources = metadata['num_sources']
|
||
|
||
with h5py.File(filepath, 'w') as h5f:
|
||
# Save RIRs, each RIR set in a separate group
|
||
for rir_key, rir_value in rir_dataset.items():
|
||
if len(rir_value) != num_sources:
|
||
raise ValueError(
|
||
f'Each RIR dataset should have exactly {num_sources} elements. Current RIR {rir_key} has {len(rir_value)} elements'
|
||
)
|
||
|
||
rir_group = h5f.create_group(rir_key)
|
||
|
||
# RIRs for different sources are saved under [group]['idx']
|
||
for idx, rir in enumerate(rir_value):
|
||
rir_group.create_dataset(f'{idx}', data=rir_value[idx])
|
||
|
||
# Save metadata
|
||
metadata_group = h5f.create_group('metadata')
|
||
for key, value in metadata.items():
|
||
metadata_group.create_dataset(key, data=value)
|
||
|
||
|
||
def load_rir_simulation(filepath: str, source: int = 0, rir_key: str = 'rir') -> Tuple[np.ndarray, float]:
|
||
"""Load simulated RIRs and metadata.
|
||
|
||
Args:
|
||
filepath: Path to simulated RIR data
|
||
source: Index of a source.
|
||
rir_key: String to denote which RIR to load, if there are multiple available.
|
||
|
||
Returns:
|
||
Multichannel RIR as ndarray with shape (num_samples, num_channels) and scalar sample rate.
|
||
"""
|
||
if not HAS_H5PY:
|
||
raise ImportError("Install h5py to use load_rir_simulation")
|
||
with h5py.File(filepath, 'r') as h5f:
|
||
# Load RIR
|
||
rir = h5f[rir_key][f'{source}'][:]
|
||
|
||
# Load metadata
|
||
sample_rate = h5f['metadata']['sample_rate'][()]
|
||
|
||
return rir, sample_rate
|
||
|
||
|
||
def convert_numpy_to_serializable(data: Union[dict, float, np.ndarray]) -> Union[dict, float, np.ndarray]:
|
||
"""Convert all numpy estries to list.
|
||
Can be used to preprocess data before writing to a JSON file.
|
||
|
||
Args:
|
||
data: Dictionary, array or scalar.
|
||
|
||
Returns:
|
||
The same structure, but converted to list if
|
||
the input is np.ndarray, so `data` can be seralized.
|
||
"""
|
||
if isinstance(data, dict):
|
||
for key, val in data.items():
|
||
data[key] = convert_numpy_to_serializable(val)
|
||
elif isinstance(data, list):
|
||
data = [convert_numpy_to_serializable(d) for d in data]
|
||
elif isinstance(data, np.ndarray):
|
||
data = data.tolist()
|
||
elif isinstance(data, np.integer):
|
||
data = int(data)
|
||
elif isinstance(data, np.floating):
|
||
data = float(data)
|
||
elif isinstance(data, np.generic):
|
||
data = data.item()
|
||
|
||
return data
|
||
|
||
|
||
def convert_rir_to_multichannel(rir: List[List[np.ndarray]]) -> List[np.ndarray]:
|
||
"""Convert RIR to a list of arrays.
|
||
|
||
Args:
|
||
rir: list of lists, each element is a single-channel RIR
|
||
|
||
Returns:
|
||
List of multichannel RIRs
|
||
"""
|
||
num_mics = len(rir)
|
||
num_sources = len(rir[0])
|
||
|
||
mc_rir = [None] * num_sources
|
||
|
||
for n_source in range(num_sources):
|
||
rir_len = [len(rir[m][n_source]) for m in range(num_mics)]
|
||
max_len = max(rir_len)
|
||
mc_rir[n_source] = np.zeros((max_len, num_mics))
|
||
for n_mic, len_mic in enumerate(rir_len):
|
||
mc_rir[n_source][:len_mic, n_mic] = rir[n_mic][n_source]
|
||
|
||
return mc_rir
|
||
|
||
|
||
def plot_rir_manifest_info(filepath: str, plot_filepath: str = None):
|
||
"""Plot distribution of parameters from manifest file.
|
||
|
||
Args:
|
||
filepath: path to a RIR corpus manifest file
|
||
plot_filepath: path to save the plot at
|
||
"""
|
||
import matplotlib.pyplot as plt
|
||
|
||
metadata = read_manifest(filepath)
|
||
|
||
# source placement
|
||
source_distance = []
|
||
source_azimuth = []
|
||
source_elevation = []
|
||
source_height = []
|
||
|
||
# room config
|
||
rir_rt60_theory = []
|
||
rir_rt60_measured = []
|
||
anechoic_rt60_theory = []
|
||
anechoic_rt60_measured = []
|
||
|
||
# get the required data
|
||
for data in metadata:
|
||
# source config
|
||
source_distance += data['source_distance']
|
||
source_azimuth += data['source_azimuth']
|
||
source_elevation += data['source_elevation']
|
||
source_height += [s_pos[2] for s_pos in data['source_position']]
|
||
|
||
# room config
|
||
rir_rt60_theory.append(data['rir_rt60_theory'])
|
||
rir_rt60_measured += data['rir_rt60_measured']
|
||
anechoic_rt60_theory.append(data['anechoic_rt60_theory'])
|
||
anechoic_rt60_measured += data['anechoic_rt60_measured']
|
||
|
||
# plot
|
||
plt.figure(figsize=(12, 6))
|
||
|
||
plt.subplot(2, 4, 1)
|
||
plt.hist(source_distance, label='distance')
|
||
plt.xlabel('distance / m')
|
||
plt.ylabel('# examples')
|
||
plt.title('Source-to-array center distance')
|
||
|
||
plt.subplot(2, 4, 2)
|
||
plt.hist(source_azimuth, label='azimuth')
|
||
plt.xlabel('azimuth / deg')
|
||
plt.ylabel('# examples')
|
||
plt.title('Source-to-array center azimuth')
|
||
|
||
plt.subplot(2, 4, 3)
|
||
plt.hist(source_elevation, label='elevation')
|
||
plt.xlabel('elevation / deg')
|
||
plt.ylabel('# examples')
|
||
plt.title('Source-to-array center elevation')
|
||
|
||
plt.subplot(2, 4, 4)
|
||
plt.hist(source_height, label='source height')
|
||
plt.xlabel('height / m')
|
||
plt.ylabel('# examples')
|
||
plt.title('Source height')
|
||
|
||
plt.subplot(2, 4, 5)
|
||
plt.hist(rir_rt60_theory, label='theory')
|
||
plt.xlabel('RT60 / s')
|
||
plt.ylabel('# examples')
|
||
plt.title('RT60 theory')
|
||
|
||
plt.subplot(2, 4, 6)
|
||
plt.hist(rir_rt60_measured, label='measured')
|
||
plt.xlabel('RT60 / s')
|
||
plt.ylabel('# examples')
|
||
plt.title('RT60 measured')
|
||
|
||
plt.subplot(2, 4, 7)
|
||
plt.hist(anechoic_rt60_theory, label='theory')
|
||
plt.xlabel('RT60 / s')
|
||
plt.ylabel('# examples')
|
||
plt.title('RT60 theory (anechoic)')
|
||
|
||
plt.subplot(2, 4, 8)
|
||
plt.hist(anechoic_rt60_measured, label='measured')
|
||
plt.xlabel('RT60 / s')
|
||
plt.ylabel('# examples')
|
||
plt.title('RT60 measured (anechoic)')
|
||
|
||
for n in range(8):
|
||
plt.subplot(2, 4, n + 1)
|
||
plt.grid()
|
||
plt.legend(loc='lower left')
|
||
|
||
plt.tight_layout()
|
||
|
||
if plot_filepath is not None:
|
||
plt.savefig(plot_filepath)
|
||
plt.close()
|
||
logging.info('Plot saved at %s', plot_filepath)
|
||
|
||
|
||
class RIRMixGenerator(object):
|
||
"""Creates a dataset of mixed signals at the microphone
|
||
by combining target speech, background noise and interference.
|
||
|
||
Correspnding signals are are generated and saved
|
||
using the `generate` method.
|
||
|
||
Input configuration is expexted to have the following structure
|
||
```
|
||
sample_rate: sample rate used for simulation
|
||
room:
|
||
subset: manifest for RIR data
|
||
target:
|
||
subset: manifest for target source data
|
||
noise:
|
||
subset: manifest for noise data
|
||
interference:
|
||
subset: manifest for interference data
|
||
interference_probability: probability that interference is present
|
||
max_num_interferers: max number of interferers, randomly selected between 0 and max
|
||
mix:
|
||
subset:
|
||
num: number of examples to generate
|
||
rsnr: range of RSNR
|
||
rsir: range of RSIR
|
||
ref_mic: reference microphone
|
||
ref_mic_rms: desired RMS at ref_mic
|
||
```
|
||
"""
|
||
|
||
def __init__(self, cfg: DictConfig):
|
||
"""
|
||
Instantiate a RIRMixGenerator object.
|
||
|
||
Args:
|
||
cfg: generator configuration defining data for room,
|
||
target signal, noise, interference and mixture
|
||
"""
|
||
logging.info("Initialize RIRMixGenerator")
|
||
self._cfg = cfg
|
||
self.check_cfg()
|
||
|
||
self.subsets = self.cfg.room.keys()
|
||
logging.info('Initialized with %d subsets: %s', len(self.subsets), str(self.subsets))
|
||
|
||
# load manifests
|
||
self.metadata = dict()
|
||
for subset in self.subsets:
|
||
subset_data = dict()
|
||
|
||
logging.info('Loading data for %s', subset)
|
||
for key in ['room', 'target', 'noise', 'interference']:
|
||
try:
|
||
subset_data[key] = read_manifest(self.cfg[key][subset])
|
||
logging.info('\t%-*s: \t%d files', 15, key, len(subset_data[key]))
|
||
except Exception as e:
|
||
subset_data[key] = None
|
||
logging.info('\t%-*s: \t0 files', 15, key)
|
||
logging.warning('\t\tManifest data not loaded. Exception: %s', str(e))
|
||
|
||
self.metadata[subset] = subset_data
|
||
|
||
logging.info('Loaded all manifests')
|
||
|
||
self.num_retries = self.cfg.get('num_retries', 5)
|
||
|
||
@property
|
||
def cfg(self):
|
||
"""Property holding the internal config of the object.
|
||
|
||
Note:
|
||
Changes to this config are not reflected in the state of the object.
|
||
Please create a new model with the updated config.
|
||
"""
|
||
return self._cfg
|
||
|
||
@property
|
||
def sample_rate(self):
|
||
return self._cfg.sample_rate
|
||
|
||
@cfg.setter
|
||
def cfg(self, cfg):
|
||
"""Property holding the internal config of the object.
|
||
|
||
Note:
|
||
Changes to this config are not reflected in the state of the object.
|
||
Please create a new model with the updated config.
|
||
"""
|
||
self._cfg = cfg
|
||
|
||
def check_cfg(self):
|
||
"""
|
||
Checks provided configuration to ensure it has the minimal required
|
||
configuration the values are in a reasonable range.
|
||
"""
|
||
# sample rate
|
||
sample_rate = self.cfg.get('sample_rate')
|
||
if sample_rate is None:
|
||
raise ValueError('Sample rate not provided.')
|
||
elif sample_rate < 0:
|
||
raise ValueError(f'Sample rate must be positive: {sample_rate}')
|
||
|
||
# room configuration
|
||
room_cfg = self.cfg.get('room')
|
||
if not room_cfg:
|
||
raise ValueError(
|
||
'Room configuration not provided. Expecting RIR manifests in format {subset: path_to_manifest}'
|
||
)
|
||
|
||
# target configuration
|
||
target_cfg = self.cfg.get('target')
|
||
if not target_cfg:
|
||
raise ValueError(
|
||
'Target configuration not provided. Expecting audio manifests in format {subset: path_to_manifest}'
|
||
)
|
||
|
||
for key in ['azimuth', 'elevation', 'distance']:
|
||
value = target_cfg.get(key)
|
||
|
||
if value is None or np.isscalar(value):
|
||
# no constraint or a fixed dimension is ok
|
||
pass
|
||
elif len(value) != 2 or not value[0] < value[1]:
|
||
# not a valid range
|
||
raise ValueError(f'Range must be specified with two positive increasing elements for {key}: {value}')
|
||
|
||
# noise configuration
|
||
noise_cfg = self.cfg.get('noise')
|
||
if not noise_cfg:
|
||
raise ValueError(
|
||
'Noise configuration not provided. Expecting audio manifests in format {subset: path_to_manifest}'
|
||
)
|
||
|
||
# interference configuration
|
||
interference_cfg = self.cfg.get('interference')
|
||
if not interference_cfg:
|
||
logging.info('Interference configuration not provided.')
|
||
else:
|
||
interference_probability = interference_cfg.get('interference_probability', 0)
|
||
max_num_interferers = interference_cfg.get('max_num_interferers', 0)
|
||
min_azimuth_to_target = interference_cfg.get('min_azimuth_to_target', 0)
|
||
if interference_probability is not None:
|
||
if interference_probability < 0:
|
||
raise ValueError(
|
||
f'Interference probability must be non-negative. Current value: {interference_probability}'
|
||
)
|
||
elif interference_probability > 0:
|
||
assert (
|
||
max_num_interferers is not None and max_num_interferers > 0
|
||
), f'Max number of interferers must be positive. Current value: {max_num_interferers}'
|
||
assert (
|
||
min_azimuth_to_target is not None and min_azimuth_to_target >= 0
|
||
), 'Min azimuth to target must be non-negative'
|
||
|
||
# mix configuration
|
||
mix_cfg = self.cfg.get('mix')
|
||
if not mix_cfg:
|
||
raise ValueError('Mix configuration not provided. Expecting configuration for each subset.')
|
||
if 'ref_mic' not in mix_cfg:
|
||
raise ValueError('Reference microphone not defined.')
|
||
if 'ref_mic_rms' not in mix_cfg:
|
||
raise ValueError('Reference microphone RMS not defined.')
|
||
|
||
def generate_target(self, subset: str) -> dict:
|
||
"""
|
||
Prepare a dictionary with target configuration.
|
||
|
||
The output dictionary contains the following information
|
||
```
|
||
room_index: index of the selected room from the RIR corpus
|
||
room_filepath: path to the room simulation file
|
||
source: index of the selected source for the target
|
||
rt60: reverberation time of the selected room
|
||
num_mics: number of microphones
|
||
azimuth: azimuth of the target source, relative to the microphone array
|
||
elevation: elevation of the target source, relative to the microphone array
|
||
distance: distance of the target source, relative to the microphone array
|
||
audio_filepath: path to the audio file for the target source
|
||
text: text for the target source audio signal, if available
|
||
duration: duration of the target source audio signal
|
||
```
|
||
|
||
Args:
|
||
subset: string denoting a subset which will be used to selected target
|
||
audio and room parameters.
|
||
|
||
Returns:
|
||
Dictionary with target configuration, including room, source index, and audio information.
|
||
"""
|
||
|
||
# Utility function
|
||
def select_target_source(room_metadata, room_indices):
|
||
"""Find a room and a source that satisfies the constraints."""
|
||
for room_index in room_indices:
|
||
# Select room
|
||
room_data = room_metadata[room_index]
|
||
|
||
# Candidate sources
|
||
sources = self.random.choice(room_data['num_sources'], size=self.num_retries, replace=False)
|
||
|
||
# Select target source in this room
|
||
for source in sources:
|
||
# Check constraints
|
||
constraints_met = []
|
||
for constraint in ['azimuth', 'elevation', 'distance']:
|
||
if self.cfg.target.get(constraint) is not None:
|
||
# Check that the selected source is in the range
|
||
source_value = room_data[f'source_{constraint}'][source]
|
||
if self.cfg.target[constraint][0] <= source_value <= self.cfg.target[constraint][1]:
|
||
constraints_met.append(True)
|
||
else:
|
||
constraints_met.append(False)
|
||
# No need to check the remaining constraints
|
||
break
|
||
|
||
# Check if a feasible source is found
|
||
if all(constraints_met):
|
||
# A feasible source has been found
|
||
return source, room_index
|
||
|
||
return None, None
|
||
|
||
# Prepare room & source position
|
||
room_metadata = self.metadata[subset]['room']
|
||
room_indices = self.random.choice(len(room_metadata), size=self.num_retries, replace=False)
|
||
source, room_index = select_target_source(room_metadata, room_indices)
|
||
|
||
if source is None:
|
||
raise RuntimeError(f'Could not find a feasible source given target constraints {self.cfg.target}')
|
||
|
||
room_data = room_metadata[room_index]
|
||
|
||
# Optional: select subset of channels
|
||
num_available_mics = len(room_data['mic_positions'])
|
||
if 'mic_array' in self.cfg:
|
||
num_mics = self.cfg.mic_array['num_mics']
|
||
mic_selection = self.cfg.mic_array['selection']
|
||
|
||
if mic_selection == 'random':
|
||
logging.debug('Randomly selecting %d mics', num_mics)
|
||
selected_mics = self.random.choice(num_available_mics, size=num_mics, replace=False)
|
||
elif isinstance(mic_selection, Iterable):
|
||
logging.debug('Using explicitly selected mics: %s', str(mic_selection))
|
||
assert (
|
||
0 <= min(mic_selection) < num_available_mics
|
||
), f'Expecting mic_selection in range [0,{num_available_mics}), current value: {mic_selection}'
|
||
selected_mics = np.array(mic_selection)
|
||
else:
|
||
raise ValueError(f'Unexpected value for mic_selection: {mic_selection}')
|
||
else:
|
||
logging.debug('Using all %d available mics', num_available_mics)
|
||
num_mics = num_available_mics
|
||
selected_mics = np.arange(num_mics)
|
||
|
||
# Double-check the number of mics is as expected
|
||
assert (
|
||
len(selected_mics) == num_mics
|
||
), f'Expecting {num_mics} mics, but received {len(selected_mics)} mics: {selected_mics}'
|
||
logging.debug('Selected mics: %s', str(selected_mics))
|
||
|
||
# Calculate distance from the source to each microphone
|
||
mic_positions = np.array(room_data['mic_positions'])[selected_mics]
|
||
source_position = np.array(room_data['source_position'][source])
|
||
distance_source_to_mic = np.linalg.norm(mic_positions - source_position, axis=1)
|
||
|
||
# Handle relative paths
|
||
room_filepath = room_data['room_filepath']
|
||
if not os.path.isabs(room_filepath):
|
||
manifest_dir = os.path.dirname(self.cfg.room[subset])
|
||
room_filepath = os.path.join(manifest_dir, room_filepath)
|
||
|
||
target_cfg = {
|
||
'room_index': int(room_index),
|
||
'room_filepath': room_filepath,
|
||
'source': source,
|
||
'rt60': room_data['rir_rt60_measured'][source],
|
||
'selected_mics': selected_mics.tolist(),
|
||
# Positions
|
||
'source_position': source_position.tolist(),
|
||
'mic_positions': mic_positions.tolist(),
|
||
# Relative to center of the array
|
||
'azimuth': room_data['source_azimuth'][source],
|
||
'elevation': room_data['source_elevation'][source],
|
||
'distance': room_data['source_distance'][source],
|
||
# Relative to mics
|
||
'distance_source_to_mic': distance_source_to_mic,
|
||
}
|
||
|
||
return target_cfg
|
||
|
||
def generate_interference(self, subset: str, target_cfg: dict) -> List[dict]:
|
||
"""
|
||
Prepare a list of dictionaries with interference configuration.
|
||
|
||
Args:
|
||
subset: string denoting a subset which will be used to select interference audio.
|
||
target_cfg: dictionary with target configuration. This is used to determine
|
||
the minimal required duration for the noise signal.
|
||
|
||
Returns:
|
||
List of dictionary with interference configuration, including source index and audio information
|
||
for one or more interference sources.
|
||
"""
|
||
if self.metadata[subset]['interference'] is None:
|
||
# No interference to be configured
|
||
return None
|
||
|
||
# Configure interfering sources
|
||
max_num_sources = self.cfg.interference.get('max_num_interferers', 0)
|
||
interference_probability = self.cfg.interference.get('interference_probability', 0)
|
||
|
||
if (
|
||
max_num_sources >= 1
|
||
and interference_probability > 0
|
||
and self.random.uniform(low=0.0, high=1.0) < interference_probability
|
||
):
|
||
# interference present
|
||
num_interferers = self.random.integers(low=1, high=max_num_sources + 1)
|
||
else:
|
||
# interference not present
|
||
return None
|
||
|
||
# Room setup: same room as target
|
||
room_index = target_cfg['room_index']
|
||
room_data = self.metadata[subset]['room'][room_index]
|
||
feasible_sources = list(range(room_data['num_sources']))
|
||
# target source is not eligible
|
||
feasible_sources.remove(target_cfg['source'])
|
||
|
||
# Constraints for interfering sources
|
||
min_azimuth_to_target = self.cfg.interference.get('min_azimuth_to_target', 0)
|
||
|
||
# Prepare interference configuration
|
||
interference_cfg = []
|
||
for n in range(num_interferers):
|
||
|
||
# Select a source
|
||
source = None
|
||
while len(feasible_sources) > 0 and source is None:
|
||
|
||
# Select a potential source for the target
|
||
source = self.random.choice(feasible_sources)
|
||
feasible_sources.remove(source)
|
||
|
||
# Check azimuth separation
|
||
if min_azimuth_to_target > 0:
|
||
source_azimuth = room_data['source_azimuth'][source]
|
||
azimuth_diff = wrap_to_180(source_azimuth - target_cfg['azimuth'])
|
||
if abs(azimuth_diff) < min_azimuth_to_target:
|
||
# Try again
|
||
source = None
|
||
continue
|
||
|
||
if source is None:
|
||
logging.warning('Could not select a feasible interference source %d of %s', n, num_interferers)
|
||
|
||
# Return what we have for now or None
|
||
return interference_cfg if interference_cfg else None
|
||
|
||
# Current source setup
|
||
interfering_source = {
|
||
'source': source,
|
||
'selected_mics': target_cfg['selected_mics'],
|
||
'position': room_data['source_position'][source],
|
||
'azimuth': room_data['source_azimuth'][source],
|
||
'elevation': room_data['source_elevation'][source],
|
||
'distance': room_data['source_distance'][source],
|
||
}
|
||
|
||
# Done with interference for this source
|
||
interference_cfg.append(interfering_source)
|
||
|
||
return interference_cfg
|
||
|
||
def generate_mix(self, subset: str, target_cfg: dict) -> dict:
|
||
"""Generate scaling parameters for mixing
|
||
the target speech at the microphone, background noise
|
||
and interference signal at the microphone.
|
||
|
||
The output dictionary contains the following information
|
||
```
|
||
rsnr: reverberant signal-to-noise ratio
|
||
rsir: reverberant signal-to-interference ratio
|
||
ref_mic: reference microphone for calculating the metrics
|
||
ref_mic_rms: RMS of the signal at the reference microphone
|
||
```
|
||
|
||
Args:
|
||
subset: string denoting the subset of configuration
|
||
target_cfg: dictionary with target configuration
|
||
|
||
Returns:
|
||
Dictionary containing configured RSNR, RSIR, ref_mic
|
||
and RMS on ref_mic.
|
||
"""
|
||
mix_cfg = dict()
|
||
|
||
for key in ['rsnr', 'rsir', 'ref_mic', 'ref_mic_rms', 'min_duration']:
|
||
if key in self.cfg.mix[subset]:
|
||
# Take the value from subset config
|
||
value = self.cfg.mix[subset].get(key)
|
||
else:
|
||
# Take the global value
|
||
value = self.cfg.mix.get(key)
|
||
|
||
if value is None:
|
||
mix_cfg[key] = None
|
||
elif np.isscalar(value):
|
||
mix_cfg[key] = value
|
||
elif len(value) == 2:
|
||
# Select from the given range, including the upper bound
|
||
mix_cfg[key] = self.random.integers(low=value[0], high=value[1] + 1)
|
||
else:
|
||
# Select one of the multiple values
|
||
mix_cfg[key] = self.random.choice(value)
|
||
|
||
if mix_cfg['ref_mic'] == 'closest':
|
||
# Select the closest mic as the reference
|
||
mix_cfg['ref_mic'] = np.argmin(target_cfg['distance_source_to_mic'])
|
||
|
||
# Configuration for saving individual components
|
||
mix_cfg['save'] = OmegaConf.to_object(self.cfg.mix['save']) if 'save' in self.cfg.mix else {}
|
||
|
||
return mix_cfg
|
||
|
||
def generate(self):
|
||
"""Generate a corpus of microphone signals by mixing target, background noise
|
||
and interference signals.
|
||
|
||
This method will prepare randomized examples based on the current configuration,
|
||
run simulations and save results to output_dir.
|
||
"""
|
||
logging.info('Generate mixed signals')
|
||
|
||
# Initialize
|
||
self.random = default_rng(seed=self.cfg.random_seed)
|
||
|
||
# Prepare output dir
|
||
output_dir = self.cfg.output_dir
|
||
if output_dir.endswith('.yaml'):
|
||
output_dir = output_dir[:-5]
|
||
|
||
# Create absolute path
|
||
logging.info('Output dir set to: %s', output_dir)
|
||
|
||
# Generate all cases
|
||
for subset in self.subsets:
|
||
|
||
output_dir_subset = os.path.join(output_dir, subset)
|
||
examples = []
|
||
|
||
if not os.path.exists(output_dir_subset):
|
||
logging.info('Creating output directory: %s', output_dir_subset)
|
||
os.makedirs(output_dir_subset)
|
||
elif os.path.isdir(output_dir_subset) and len(os.listdir(output_dir_subset)) > 0:
|
||
raise RuntimeError(f'Output directory {output_dir_subset} is not empty.')
|
||
|
||
num_examples = self.cfg.mix[subset].num
|
||
logging.info('Preparing %d examples for subset %s', num_examples, subset)
|
||
|
||
# Generate examples
|
||
for n_example in tqdm(range(num_examples), total=num_examples, desc=f'Preparing {subset}'):
|
||
# prepare configuration
|
||
target_cfg = self.generate_target(subset)
|
||
interference_cfg = self.generate_interference(subset, target_cfg)
|
||
mix_cfg = self.generate_mix(subset, target_cfg)
|
||
|
||
# base file name
|
||
base_output_filepath = os.path.join(output_dir_subset, f'{subset}_example_{n_example:09d}')
|
||
|
||
# prepare example
|
||
example = {
|
||
'sample_rate': self.sample_rate,
|
||
'target_cfg': target_cfg,
|
||
'interference_cfg': interference_cfg,
|
||
'mix_cfg': mix_cfg,
|
||
'base_output_filepath': base_output_filepath,
|
||
}
|
||
|
||
examples.append(example)
|
||
|
||
# Audio data
|
||
audio_metadata = {
|
||
'target': self.metadata[subset]['target'],
|
||
'target_dir': os.path.dirname(self.cfg.target[subset]), # manifest_dir
|
||
'noise': self.metadata[subset]['noise'],
|
||
'noise_dir': os.path.dirname(self.cfg.noise[subset]), # manifest_dir
|
||
}
|
||
|
||
if interference_cfg is not None:
|
||
audio_metadata.update(
|
||
{
|
||
'interference': self.metadata[subset]['interference'],
|
||
'interference_dir': os.path.dirname(self.cfg.interference[subset]), # manifest_dir
|
||
}
|
||
)
|
||
|
||
# Simulation
|
||
if (num_workers := self.cfg.get('num_workers')) is None:
|
||
num_workers = os.cpu_count() - 1
|
||
|
||
if num_workers is not None and num_workers > 1:
|
||
logging.info(f'Simulate using {num_workers} workers')
|
||
examples_and_audio_metadata = zip(examples, itertools.repeat(audio_metadata, len(examples)))
|
||
with multiprocessing.Pool(processes=num_workers) as pool:
|
||
metadata = list(
|
||
tqdm(
|
||
pool.imap(simulate_room_mix_helper, examples_and_audio_metadata),
|
||
total=len(examples),
|
||
desc=f'Simulating {subset}',
|
||
)
|
||
)
|
||
else:
|
||
logging.info('Simulate using a single worker')
|
||
metadata = []
|
||
for example in tqdm(examples, total=len(examples), desc=f'Simulating {subset}'):
|
||
metadata.append(simulate_room_mix(**example, audio_metadata=audio_metadata))
|
||
|
||
# Save manifest
|
||
manifest_filepath = os.path.join(output_dir, f'{os.path.basename(output_dir)}_{subset}.json')
|
||
|
||
if os.path.exists(manifest_filepath) and os.path.isfile(manifest_filepath):
|
||
raise RuntimeError(f'Manifest config file exists: {manifest_filepath}')
|
||
|
||
# Make all paths in the manifest relative to the output dir
|
||
for data in tqdm(metadata, total=len(metadata), desc=f'Making filepaths relative {subset}'):
|
||
for key, val in data.items():
|
||
if key.endswith('_filepath') and val is not None:
|
||
data[key] = os.path.relpath(val, start=output_dir)
|
||
|
||
write_manifest(manifest_filepath, metadata)
|
||
|
||
# Generate plots with information about generated data
|
||
plot_filepath = os.path.join(output_dir, f'{os.path.basename(output_dir)}_{subset}_info.png')
|
||
|
||
if os.path.exists(plot_filepath) and os.path.isfile(plot_filepath):
|
||
raise RuntimeError(f'Plot file exists: {plot_filepath}')
|
||
|
||
plot_mix_manifest_info(manifest_filepath, plot_filepath=plot_filepath)
|
||
|
||
# Save used configuration for reference
|
||
config_filepath = os.path.join(output_dir, 'config.yaml')
|
||
if os.path.exists(config_filepath) and os.path.isfile(config_filepath):
|
||
raise RuntimeError(f'Output config file exists: {config_filepath}')
|
||
|
||
OmegaConf.save(self.cfg, config_filepath, resolve=True)
|
||
|
||
|
||
def convolve_rir(signal: np.ndarray, rir: np.ndarray) -> np.ndarray:
|
||
"""Convolve signal with a possibly multichannel IR in rir, i.e.,
|
||
calculate the following for each channel m:
|
||
|
||
signal_m = rir_m \ast signal
|
||
|
||
Args:
|
||
signal: single-channel signal (samples,)
|
||
rir: single- or multi-channel IR, (samples,) or (samples, channels)
|
||
|
||
Returns:
|
||
out: same length as signal, same number of channels as rir, shape (samples, channels)
|
||
"""
|
||
num_samples = len(signal)
|
||
if rir.ndim == 1:
|
||
# convolve and trim to length
|
||
out = convolve(signal, rir)[:num_samples]
|
||
elif rir.ndim == 2:
|
||
num_channels = rir.shape[1]
|
||
out = np.zeros((num_samples, num_channels))
|
||
for m in range(num_channels):
|
||
out[:, m] = convolve(signal, rir[:, m])[:num_samples]
|
||
|
||
else:
|
||
raise RuntimeError(f'RIR with {rir.ndim} not supported')
|
||
|
||
return out
|
||
|
||
|
||
def calculate_drr(rir: np.ndarray, sample_rate: float, n_direct: List[int], n_0_ms=2.5) -> List[float]:
|
||
"""Calculate direct-to-reverberant ratio (DRR) from the measured RIR.
|
||
|
||
Calculation is done as in eq. (3) from [1].
|
||
|
||
Args:
|
||
rir: room impulse response, shape (num_samples, num_channels)
|
||
sample_rate: sample rate for the impulse response
|
||
n_direct: direct path delay
|
||
n_0_ms: window around n_direct for calculating the direct path energy
|
||
|
||
Returns:
|
||
Calculated DRR for each channel of the input RIR.
|
||
|
||
References:
|
||
[1] Eaton et al, The ACE challenge: Corpus description and performance evaluation, WASPAA 2015
|
||
"""
|
||
# Define a window around the direct path delay
|
||
n_0 = int(n_0_ms * sample_rate / 1000)
|
||
|
||
len_rir, num_channels = rir.shape
|
||
drr = [None] * num_channels
|
||
for m in range(num_channels):
|
||
|
||
# Window around the direct path
|
||
dir_start = max(n_direct[m] - n_0, 0)
|
||
dir_end = n_direct[m] + n_0
|
||
|
||
# Power of the direct component
|
||
pow_dir = np.sum(np.abs(rir[dir_start:dir_end, m]) ** 2) / len_rir
|
||
|
||
# Power of the reverberant component
|
||
pow_reverberant = (np.sum(np.abs(rir[0:dir_start, m]) ** 2) + np.sum(np.abs(rir[dir_end:, m]) ** 2)) / len_rir
|
||
|
||
# DRR in dB
|
||
drr[m] = pow2db(pow_dir / pow_reverberant)
|
||
|
||
return drr
|
||
|
||
|
||
def normalize_max(x: np.ndarray, max_db: float = 0, eps: float = 1e-16) -> np.ndarray:
|
||
"""Normalize max input value to max_db full scale (±1).
|
||
|
||
Args:
|
||
x: input signal
|
||
max_db: desired max magnitude compared to full scale
|
||
eps: small regularization constant
|
||
|
||
Returns:
|
||
Normalized signal with max absolute value max_db.
|
||
"""
|
||
max_val = db2mag(max_db)
|
||
return max_val * x / (np.max(np.abs(x)) + eps)
|
||
|
||
|
||
def simultaneously_active_rms(
|
||
x: np.ndarray,
|
||
y: np.ndarray,
|
||
sample_rate: float,
|
||
rms_threshold_db: float = -60,
|
||
window_len_ms: float = 200,
|
||
min_active_duration: float = 0.5,
|
||
) -> Tuple[float, float]:
|
||
"""Calculate RMS over segments where both input signals are active.
|
||
|
||
Args:
|
||
x: first input signal
|
||
y: second input signal
|
||
sample_rate: sample rate for input signals in Hz
|
||
rms_threshold_db: threshold for determining activity of the signal, relative
|
||
to max absolute value
|
||
window_len_ms: window length in milliseconds, used for calculating segmental RMS
|
||
min_active_duration: minimal duration of the active segments
|
||
|
||
Returns:
|
||
RMS value over active segments for x and y.
|
||
"""
|
||
if len(x) != len(y):
|
||
raise RuntimeError(f'Expecting signals of same length: len(x)={len(x)}, len(y)={len(y)}')
|
||
window_len = int(window_len_ms * sample_rate / 1000)
|
||
rms_threshold = db2mag(rms_threshold_db) # linear scale
|
||
|
||
x_normalized = normalize_max(x)
|
||
y_normalized = normalize_max(y)
|
||
|
||
x_active_power = y_active_power = active_len = 0
|
||
for start in range(0, len(x) - window_len, window_len):
|
||
window = slice(start, start + window_len)
|
||
|
||
# check activity on the scaled signal
|
||
x_window_rms = rms(x_normalized[window])
|
||
y_window_rms = rms(y_normalized[window])
|
||
|
||
if x_window_rms > rms_threshold and y_window_rms > rms_threshold:
|
||
# sum the power of the original non-scaled signal
|
||
x_active_power += np.sum(np.abs(x[window]) ** 2)
|
||
y_active_power += np.sum(np.abs(y[window]) ** 2)
|
||
active_len += window_len
|
||
|
||
if active_len < int(min_active_duration * sample_rate):
|
||
raise RuntimeError(
|
||
f'Signals are simultaneously active less than {min_active_duration} s: only {active_len/sample_rate} s'
|
||
)
|
||
|
||
# normalize
|
||
x_active_power /= active_len
|
||
y_active_power /= active_len
|
||
|
||
return np.sqrt(x_active_power), np.sqrt(y_active_power)
|
||
|
||
|
||
def scaled_disturbance(
|
||
signal: np.ndarray,
|
||
disturbance: np.ndarray,
|
||
sdr: float,
|
||
sample_rate: float = None,
|
||
ref_channel: int = 0,
|
||
eps: float = 1e-16,
|
||
) -> np.ndarray:
|
||
"""
|
||
Args:
|
||
signal: numpy array, shape (num_samples, num_channels)
|
||
disturbance: numpy array, same shape as signal
|
||
sdr: desired signal-to-disturbance ration
|
||
sample_rate: sample rate of the input signals
|
||
ref_channel: ref mic used to calculate RMS
|
||
eps: regularization constant
|
||
|
||
Returns:
|
||
Scaled disturbance, so that signal-to-disturbance ratio at ref_channel
|
||
is approximately equal to input SDR during simultaneously active
|
||
segment of signal and disturbance.
|
||
"""
|
||
if signal.shape != disturbance.shape:
|
||
raise ValueError(f'Signal and disturbance shapes do not match: {signal.shape} != {disturbance.shape}')
|
||
|
||
# set scaling based on RMS at ref_mic
|
||
signal_rms, disturbance_rms = simultaneously_active_rms(
|
||
signal[:, ref_channel], disturbance[:, ref_channel], sample_rate=sample_rate
|
||
)
|
||
disturbance_gain = db2mag(-sdr) * signal_rms / (disturbance_rms + eps)
|
||
# scale disturbance
|
||
scaled_disturbance = disturbance_gain * disturbance
|
||
return scaled_disturbance
|
||
|
||
|
||
def prepare_source_signal(
|
||
signal_type: str,
|
||
sample_rate: int,
|
||
audio_data: List[dict],
|
||
audio_dir: Optional[str] = None,
|
||
min_duration: Optional[int] = None,
|
||
ref_signal: Optional[np.ndarray] = None,
|
||
mic_positions: Optional[np.ndarray] = None,
|
||
num_retries: int = 10,
|
||
) -> tuple:
|
||
"""Prepare an audio signal for a source.
|
||
|
||
Args:
|
||
signal_type: 'point' or 'diffuse'
|
||
sample_rate: Sampling rate for the signal
|
||
audio_data: List of audio items, each is a dictionary with audio_filepath, duration, offset and optionally text
|
||
audio_dir: Base directory for resolving paths, e.g., manifest basedir
|
||
min_duration: Minimal duration to be loaded if ref_signal is not provided, in seconds
|
||
ref_signal: Optional, used to determine the length of the signal
|
||
mic_positions: Optional, used to prepare approximately diffuse signal
|
||
num_retries: Number of retries when selecting the source files
|
||
|
||
Returns:
|
||
(audio_signal, metadata), where audio_signal is an ndarray and metadata is a dictionary
|
||
with audio filepaths, durations and offsets
|
||
"""
|
||
if signal_type not in ['point', 'diffuse']:
|
||
raise ValueError(f'Unexpected signal type {signal_type}.')
|
||
|
||
if audio_data is None:
|
||
# No data to load
|
||
return None
|
||
|
||
metadata = {}
|
||
|
||
if ref_signal is None:
|
||
audio_signal = None
|
||
# load at least one sample if min_duration is not provided
|
||
samples_to_load = int(min_duration * sample_rate) if min_duration is not None else 1
|
||
source_signals_metadata = {'audio_filepath': [], 'duration': [], 'offset': [], 'text': []}
|
||
|
||
while samples_to_load > 0:
|
||
# Select a random item and load the audio
|
||
item = random.choice(audio_data)
|
||
|
||
audio_filepath = item['audio_filepath']
|
||
if not os.path.isabs(audio_filepath) and audio_dir is not None:
|
||
audio_filepath = os.path.join(audio_dir, audio_filepath)
|
||
|
||
# Load audio
|
||
check_min_sample_rate(audio_filepath, sample_rate)
|
||
audio_segment = AudioSegment.from_file(
|
||
audio_file=audio_filepath,
|
||
target_sr=sample_rate,
|
||
duration=item['duration'],
|
||
offset=item.get('offset', 0),
|
||
)
|
||
|
||
if signal_type == 'point':
|
||
if audio_segment.num_channels > 1:
|
||
raise RuntimeError(
|
||
f'Expecting single-channel source signal, but received {audio_segment.num_channels}. File: {audio_filepath}'
|
||
)
|
||
else:
|
||
raise ValueError(f'Unexpected signal type {signal_type}.')
|
||
|
||
source_signals_metadata['audio_filepath'].append(audio_filepath)
|
||
source_signals_metadata['duration'].append(item['duration'])
|
||
source_signals_metadata['duration'].append(item.get('offset', 0))
|
||
source_signals_metadata['text'].append(item.get('text'))
|
||
|
||
# not perfect, since different files may have different distributions
|
||
segment_samples = normalize_max(audio_segment.samples)
|
||
# concatenate
|
||
audio_signal = (
|
||
np.concatenate((audio_signal, segment_samples)) if audio_signal is not None else segment_samples
|
||
)
|
||
# remaining samples
|
||
samples_to_load -= len(segment_samples)
|
||
|
||
# Finally, we need only the metadata for the complete signal
|
||
metadata = {
|
||
'duration': sum(source_signals_metadata['duration']),
|
||
'offset': 0,
|
||
}
|
||
|
||
# Add text only if all source signals have text
|
||
if all([isinstance(tt, str) for tt in source_signals_metadata['text']]):
|
||
metadata['text'] = ' '.join(source_signals_metadata['text'])
|
||
else:
|
||
# Load a signal with total_len samples and ensure it has enough simultaneous activity/overlap with ref_signal
|
||
# Concatenate multiple files if necessary
|
||
total_len = len(ref_signal)
|
||
|
||
for n in range(num_retries):
|
||
|
||
audio_signal = None
|
||
source_signals_metadata = {'audio_filepath': [], 'duration': [], 'offset': []}
|
||
|
||
if signal_type == 'point':
|
||
samples_to_load = total_len
|
||
elif signal_type == 'diffuse':
|
||
# Load longer signal so it can be reshaped into (samples, mics) and
|
||
# used to generate approximately diffuse noise field
|
||
num_mics = len(mic_positions)
|
||
samples_to_load = num_mics * total_len
|
||
|
||
while samples_to_load > 0:
|
||
# Select an audio file
|
||
item = random.choice(audio_data)
|
||
|
||
audio_filepath = item['audio_filepath']
|
||
if not os.path.isabs(audio_filepath) and audio_dir is not None:
|
||
audio_filepath = os.path.join(audio_dir, audio_filepath)
|
||
|
||
# Load audio signal
|
||
check_min_sample_rate(audio_filepath, sample_rate)
|
||
|
||
if (max_offset := item['duration'] - np.ceil(samples_to_load / sample_rate)) > 0:
|
||
# Load with a random offset if the example is longer than samples_to_load
|
||
offset = random.uniform(0, max_offset)
|
||
duration = -1
|
||
else:
|
||
# Load the whole file
|
||
offset, duration = 0, item['duration']
|
||
audio_segment = AudioSegment.from_file(
|
||
audio_file=audio_filepath, target_sr=sample_rate, duration=duration, offset=offset
|
||
)
|
||
|
||
# Prepare a single-channel signal
|
||
if audio_segment.num_channels == 1:
|
||
# Take all samples
|
||
segment_samples = audio_segment.samples
|
||
else:
|
||
# Take a random channel
|
||
selected_channel = random.choice(range(audio_segment.num_channels))
|
||
segment_samples = audio_segment.samples[:, selected_channel]
|
||
|
||
source_signals_metadata['audio_filepath'].append(audio_filepath)
|
||
source_signals_metadata['duration'].append(len(segment_samples) / sample_rate)
|
||
source_signals_metadata['offset'].append(offset)
|
||
|
||
# not perfect, since different files may have different distributions
|
||
segment_samples = normalize_max(segment_samples)
|
||
# concatenate
|
||
audio_signal = (
|
||
np.concatenate((audio_signal, segment_samples)) if audio_signal is not None else segment_samples
|
||
)
|
||
# remaining samples
|
||
samples_to_load -= len(segment_samples)
|
||
|
||
if signal_type == 'diffuse' and num_mics > 1:
|
||
try:
|
||
# Trim and reshape to num_mics to prepare num_mics source signals
|
||
audio_signal = audio_signal[: num_mics * total_len].reshape(num_mics, -1).T
|
||
|
||
# Make spherically diffuse noise
|
||
audio_signal = generate_approximate_noise_field(
|
||
mic_positions=np.array(mic_positions), noise_signal=audio_signal, sample_rate=sample_rate
|
||
)
|
||
except Exception as e:
|
||
logging.info('Failed to generate approximate noise field: %s', str(e))
|
||
logging.info('Try again.')
|
||
# Try again
|
||
audio_signal, source_signals_metadata = None, {}
|
||
continue
|
||
|
||
# Trim to length
|
||
audio_signal = audio_signal[:total_len, ...]
|
||
|
||
# Include the channel dimension if the reference includes it
|
||
if ref_signal.ndim == 2 and audio_signal.ndim == 1:
|
||
audio_signal = audio_signal[:, None]
|
||
|
||
try:
|
||
# Signal and ref_signal should be simultaneously active
|
||
simultaneously_active_rms(ref_signal, audio_signal, sample_rate=sample_rate)
|
||
# We have enough overlap
|
||
break
|
||
except Exception as e:
|
||
# Signal and ref_signal are not overlapping, try again
|
||
logging.info('Exception: %s', str(e))
|
||
logging.info('Signals are not overlapping, try again.')
|
||
audio_signal, source_signals_metadata = None, {}
|
||
continue
|
||
|
||
if audio_signal is None:
|
||
logging.warning('Audio signal not set: %s.', signal_type)
|
||
|
||
metadata['source_signals'] = source_signals_metadata
|
||
|
||
return audio_signal, metadata
|
||
|
||
|
||
def check_min_sample_rate(filepath: str, sample_rate: float):
|
||
"""Make sure the file's sample rate is at least sample_rate.
|
||
This will make sure that we have only downsampling if loading
|
||
this file, while upsampling is not permitted.
|
||
|
||
Args:
|
||
filepath: path to a file
|
||
sample_rate: desired sample rate
|
||
"""
|
||
file_sample_rate = librosa.get_samplerate(path=filepath)
|
||
if file_sample_rate < sample_rate:
|
||
raise RuntimeError(
|
||
f'Sample rate ({file_sample_rate}) is lower than the desired sample rate ({sample_rate}). File: {filepath}.'
|
||
)
|
||
|
||
|
||
def simulate_room_mix(
|
||
sample_rate: int,
|
||
target_cfg: dict,
|
||
interference_cfg: dict,
|
||
mix_cfg: dict,
|
||
audio_metadata: dict,
|
||
base_output_filepath: str,
|
||
max_amplitude: float = 0.999,
|
||
eps: float = 1e-16,
|
||
) -> dict:
|
||
"""Simulate mixture signal at the microphone, including target, noise and
|
||
interference signals and mixed at specific RSNR and RSIR.
|
||
|
||
Args:
|
||
sample_rate: Sample rate for all signals
|
||
target_cfg: Dictionary with configuration of the target. Includes
|
||
room_filepath, source index, audio_filepath, duration
|
||
noise_cfg: List of dictionaries, where each item includes audio_filepath,
|
||
offset and duration.
|
||
interference_cfg: List of dictionaries, where each item contains source
|
||
index
|
||
mix_cfg: Dictionary with the mixture configuration. Includes RSNR, RSIR,
|
||
ref_mic and ref_mic_rms.
|
||
audio_metadata: Dictionary with a list of files for target, noise and interference
|
||
base_output_filepath: All output audio files will be saved with this prefix by
|
||
adding a diffierent suffix for each component, e.g., _mic.wav.
|
||
max_amplitude: Maximum amplitude of the mic signal, used to prevent clipping.
|
||
eps: Small regularization constant.
|
||
|
||
Returns:
|
||
Dictionary with metadata based on the mixture setup and
|
||
simulation results. This corresponds to a line of the
|
||
output manifest file.
|
||
"""
|
||
|
||
# Local utilities
|
||
def load_rir(
|
||
room_filepath: str, source: int, selected_mics: list, sample_rate: float, rir_key: str = 'rir'
|
||
) -> np.ndarray:
|
||
"""Load a RIR and check that the sample rate is matching the desired sample rate
|
||
|
||
Args:
|
||
room_filepath: Path to a room simulation in an h5 file
|
||
source: Index of the desired source
|
||
sample_rate: Sample rate of the simulation
|
||
rir_key: Key of the RIR to load from the simulation.
|
||
|
||
Returns:
|
||
Numpy array with shape (num_samples, num_channels)
|
||
"""
|
||
rir, rir_sample_rate = load_rir_simulation(room_filepath, source=source, rir_key=rir_key)
|
||
if rir_sample_rate != sample_rate:
|
||
raise RuntimeError(
|
||
f'RIR sample rate ({sample_rate}) is not matching the expected sample rate ({sample_rate}). File: {room_filepath}'
|
||
)
|
||
return rir[:, selected_mics]
|
||
|
||
def get_early_rir(
|
||
rir: np.ndarray, rir_anechoic: np.ndarray, sample_rate: int, early_duration: float = 0.050
|
||
) -> np.ndarray:
|
||
"""Return only the early part of the RIR."""
|
||
early_len = int(early_duration * sample_rate)
|
||
direct_path_delay = np.min(np.argmax(rir_anechoic, axis=0))
|
||
rir_early = rir.copy()
|
||
rir_early[direct_path_delay + early_len :, :] = 0
|
||
return rir_early
|
||
|
||
def save_audio(
|
||
base_path: str,
|
||
tag: str,
|
||
audio_signal: Optional[np.ndarray],
|
||
sample_rate: int,
|
||
save: str = 'all',
|
||
ref_mic: Optional[int] = None,
|
||
format: str = 'wav',
|
||
subtype: str = 'float',
|
||
):
|
||
"""Save audio signal and return filepath."""
|
||
if (audio_signal is None) or (not save):
|
||
return None
|
||
|
||
if save == 'ref_mic':
|
||
# save only ref_mic
|
||
audio_signal = audio_signal[:, ref_mic]
|
||
|
||
audio_filepath = base_path + f'_{tag}.{format}'
|
||
sf.write(audio_filepath, audio_signal, sample_rate, subtype)
|
||
|
||
return audio_filepath
|
||
|
||
# Target RIRs
|
||
target_rir = load_rir(
|
||
target_cfg['room_filepath'],
|
||
source=target_cfg['source'],
|
||
selected_mics=target_cfg['selected_mics'],
|
||
sample_rate=sample_rate,
|
||
)
|
||
target_rir_anechoic = load_rir(
|
||
target_cfg['room_filepath'],
|
||
source=target_cfg['source'],
|
||
sample_rate=sample_rate,
|
||
selected_mics=target_cfg['selected_mics'],
|
||
rir_key='anechoic',
|
||
)
|
||
target_rir_early = get_early_rir(rir=target_rir, rir_anechoic=target_rir_anechoic, sample_rate=sample_rate)
|
||
|
||
# Target signals
|
||
target_signal, target_metadata = prepare_source_signal(
|
||
signal_type='point',
|
||
sample_rate=sample_rate,
|
||
audio_data=audio_metadata['target'],
|
||
audio_dir=audio_metadata['target_dir'],
|
||
min_duration=mix_cfg['min_duration'],
|
||
)
|
||
source_signals_metadata = {'target': target_metadata['source_signals']}
|
||
|
||
# Convolve target
|
||
target_reverberant = convolve_rir(target_signal, target_rir)
|
||
target_anechoic = convolve_rir(target_signal, target_rir_anechoic)
|
||
target_early = convolve_rir(target_signal, target_rir_early)
|
||
|
||
# Prepare noise signal
|
||
noise, noise_metadata = prepare_source_signal(
|
||
signal_type='diffuse',
|
||
sample_rate=sample_rate,
|
||
mic_positions=target_cfg['mic_positions'],
|
||
audio_data=audio_metadata['noise'],
|
||
audio_dir=audio_metadata['noise_dir'],
|
||
ref_signal=target_reverberant,
|
||
)
|
||
source_signals_metadata['noise'] = noise_metadata['source_signals']
|
||
|
||
# Prepare interference signal
|
||
if interference_cfg is None:
|
||
interference = None
|
||
else:
|
||
# Load interference signals
|
||
interference = 0
|
||
source_signals_metadata['interference'] = []
|
||
for i_cfg in interference_cfg:
|
||
# Load single-channel signal for directional interference
|
||
i_signal, i_metadata = prepare_source_signal(
|
||
signal_type='point',
|
||
sample_rate=sample_rate,
|
||
audio_data=audio_metadata['interference'],
|
||
audio_dir=audio_metadata['interference_dir'],
|
||
ref_signal=target_signal,
|
||
)
|
||
source_signals_metadata['interference'].append(i_metadata['source_signals'])
|
||
# Load RIR from the same room as the target, but a difference source
|
||
i_rir = load_rir(
|
||
target_cfg['room_filepath'],
|
||
source=i_cfg['source'],
|
||
selected_mics=i_cfg['selected_mics'],
|
||
sample_rate=sample_rate,
|
||
)
|
||
# Convolve interference
|
||
i_reverberant = convolve_rir(i_signal, i_rir)
|
||
# Sum
|
||
interference += i_reverberant
|
||
|
||
# Scale and add components of the signal
|
||
mic = target_reverberant.copy()
|
||
|
||
if noise is not None:
|
||
noise = scaled_disturbance(
|
||
signal=target_reverberant,
|
||
disturbance=noise,
|
||
sdr=mix_cfg['rsnr'],
|
||
sample_rate=sample_rate,
|
||
ref_channel=mix_cfg['ref_mic'],
|
||
)
|
||
# Update mic signal
|
||
mic += noise
|
||
|
||
if interference is not None:
|
||
interference = scaled_disturbance(
|
||
signal=target_reverberant,
|
||
disturbance=interference,
|
||
sdr=mix_cfg['rsir'],
|
||
sample_rate=sample_rate,
|
||
ref_channel=mix_cfg['ref_mic'],
|
||
)
|
||
# Update mic signal
|
||
mic += interference
|
||
|
||
# Set the final mic signal level
|
||
mic_rms = rms(mic[:, mix_cfg['ref_mic']])
|
||
global_gain = db2mag(mix_cfg['ref_mic_rms']) / (mic_rms + eps)
|
||
mic_max = np.max(np.abs(mic))
|
||
if (clipped_max := mic_max * global_gain) > max_amplitude:
|
||
# Downscale the global gain to prevent clipping + adjust ref_mic_rms accordingly
|
||
clipping_prevention_gain = max_amplitude / clipped_max
|
||
global_gain *= clipping_prevention_gain
|
||
mix_cfg['ref_mic_rms'] += mag2db(clipping_prevention_gain)
|
||
|
||
logging.debug(
|
||
'Clipping prevented for example %s (protection gain: %.2f dB)',
|
||
base_output_filepath,
|
||
mag2db(clipping_prevention_gain),
|
||
)
|
||
|
||
# save signals
|
||
signals = {
|
||
'mic': mic,
|
||
'target_reverberant': target_reverberant,
|
||
'target_anechoic': target_anechoic,
|
||
'target_early': target_early,
|
||
'noise': noise,
|
||
'interference': interference,
|
||
}
|
||
|
||
metadata = {}
|
||
|
||
for tag, signal in signals.items():
|
||
|
||
if signal is not None:
|
||
# scale all signal components with the global gain
|
||
signal = global_gain * signal
|
||
|
||
audio_filepath = save_audio(
|
||
base_path=base_output_filepath,
|
||
tag=tag,
|
||
audio_signal=signal,
|
||
sample_rate=sample_rate,
|
||
save=mix_cfg['save'].get(tag, 'all'),
|
||
ref_mic=mix_cfg['ref_mic'],
|
||
format=mix_cfg['save'].get('format', 'wav'),
|
||
subtype=mix_cfg['save'].get('subtype', 'float'),
|
||
)
|
||
|
||
if tag == 'mic':
|
||
metadata['audio_filepath'] = audio_filepath
|
||
else:
|
||
metadata[tag + '_filepath'] = audio_filepath
|
||
|
||
# Add metadata
|
||
metadata.update(
|
||
{
|
||
'text': target_metadata.get('text'),
|
||
'duration': target_metadata['duration'],
|
||
'target_cfg': target_cfg,
|
||
'interference_cfg': interference_cfg,
|
||
'mix_cfg': mix_cfg,
|
||
'ref_channel': mix_cfg.get('ref_mic'),
|
||
'rt60': target_cfg.get('rt60'),
|
||
'drr': calculate_drr(target_rir, sample_rate, n_direct=np.argmax(target_rir_anechoic, axis=0)),
|
||
'rsnr': None if noise is None else mix_cfg['rsnr'],
|
||
'rsir': None if interference is None else mix_cfg['rsir'],
|
||
'source_signals': source_signals_metadata,
|
||
}
|
||
)
|
||
|
||
return convert_numpy_to_serializable(metadata)
|
||
|
||
|
||
def simulate_room_mix_helper(example_and_audio_metadata: tuple) -> dict:
|
||
"""Wrapper around `simulate_room_mix` for pool.imap.
|
||
|
||
Args:
|
||
args: example and audio_metadata that are forwarded to `simulate_room_mix`
|
||
|
||
Returns:
|
||
Dictionary with metadata, see `simulate_room_mix`
|
||
"""
|
||
example, audio_metadata = example_and_audio_metadata
|
||
return simulate_room_mix(**example, audio_metadata=audio_metadata)
|
||
|
||
|
||
def plot_mix_manifest_info(filepath: str, plot_filepath: str = None):
|
||
"""Plot distribution of parameters from the manifest file.
|
||
|
||
Args:
|
||
filepath: path to a RIR corpus manifest file
|
||
plot_filepath: path to save the plot at
|
||
"""
|
||
import matplotlib.pyplot as plt
|
||
|
||
metadata = read_manifest(filepath)
|
||
|
||
# target info
|
||
target_distance = []
|
||
target_azimuth = []
|
||
target_elevation = []
|
||
target_duration = []
|
||
|
||
# room config
|
||
rt60 = []
|
||
drr = []
|
||
|
||
# noise
|
||
rsnr = []
|
||
rsir = []
|
||
|
||
# get the required data
|
||
for data in metadata:
|
||
# target info
|
||
target_distance.append(data['target_cfg']['distance'])
|
||
target_azimuth.append(data['target_cfg']['azimuth'])
|
||
target_elevation.append(data['target_cfg']['elevation'])
|
||
target_duration.append(data['duration'])
|
||
|
||
# room config
|
||
rt60.append(data['rt60'])
|
||
drr += data['drr'] # average DRR across all mics
|
||
|
||
# noise
|
||
if data['rsnr'] is not None:
|
||
rsnr.append(data['rsnr'])
|
||
|
||
if data['rsir'] is not None:
|
||
rsir.append(data['rsir'])
|
||
|
||
# plot
|
||
plt.figure(figsize=(12, 6))
|
||
|
||
plt.subplot(2, 4, 1)
|
||
plt.hist(target_distance, label='distance')
|
||
plt.xlabel('distance / m')
|
||
plt.ylabel('# examples')
|
||
plt.title('Target-to-array distance')
|
||
|
||
plt.subplot(2, 4, 2)
|
||
plt.hist(target_azimuth, label='azimuth')
|
||
plt.xlabel('azimuth / deg')
|
||
plt.ylabel('# examples')
|
||
plt.title('Target-to-array azimuth')
|
||
|
||
plt.subplot(2, 4, 3)
|
||
plt.hist(target_elevation, label='elevation')
|
||
plt.xlabel('elevation / deg')
|
||
plt.ylabel('# examples')
|
||
plt.title('Target-to-array elevation')
|
||
|
||
plt.subplot(2, 4, 4)
|
||
plt.hist(target_duration, label='duration')
|
||
plt.xlabel('time / s')
|
||
plt.ylabel('# examples')
|
||
plt.title('Target duration')
|
||
|
||
plt.subplot(2, 4, 5)
|
||
plt.hist(rt60, label='RT60')
|
||
plt.xlabel('RT60 / s')
|
||
plt.ylabel('# examples')
|
||
plt.title('RT60')
|
||
|
||
plt.subplot(2, 4, 6)
|
||
plt.hist(drr, label='DRR')
|
||
plt.xlabel('DRR / dB')
|
||
plt.ylabel('# examples')
|
||
plt.title('DRR [avg over mics]')
|
||
|
||
if len(rsnr) > 0:
|
||
plt.subplot(2, 4, 7)
|
||
plt.hist(rsnr, label='RSNR')
|
||
plt.xlabel('RSNR / dB')
|
||
plt.ylabel('# examples')
|
||
plt.title(f'RSNR [{100 * len(rsnr) / len(rt60):.0f}% ex]')
|
||
|
||
if len(rsir):
|
||
plt.subplot(2, 4, 8)
|
||
plt.hist(rsir, label='RSIR')
|
||
plt.xlabel('RSIR / dB')
|
||
plt.ylabel('# examples')
|
||
plt.title(f'RSIR [{100 * len(rsir) / len(rt60):.0f}% ex]')
|
||
|
||
for n in range(8):
|
||
plt.subplot(2, 4, n + 1)
|
||
plt.grid()
|
||
plt.legend(loc='lower left')
|
||
|
||
plt.tight_layout()
|
||
|
||
if plot_filepath is not None:
|
||
plt.savefig(plot_filepath)
|
||
plt.close()
|
||
logging.info('Plot saved at %s', plot_filepath)
|