sonusai 1.0.16__cp311-abi3-macosx_11_0_arm64.whl

sonusai/mixture/spectral_mask.py

@@ -0,0 +1,51 @@
+from ..datatypes import AudioF
+from ..datatypes import SpectralMask
+
+
+def apply_spectral_mask(audio_f: AudioF, spectral_mask: SpectralMask, seed: int | None = None) -> AudioF:
+    """Apply frequency and time masking
+
+    Implementation of SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition
+
+    Ref: https://arxiv.org/pdf/1904.08779.pdf
+
+    f_width consecutive bins [f_start, f_start + f_width) are masked, where f_width is chosen from a uniform
+    distribution from 0 to the f_max_width, and f_start is chosen from [0, bins - f_width).
+
+    t_width consecutive frames [t_start, t_start + t_width) are masked, where t_width is chosen from a uniform
+    distribution from 0 to the t_max_width, and t_start is chosen from [0, frames - t_width).
+
+    A time mask cannot be wider than t_max_percent times the number of frames.
+
+    :param audio_f: Numpy array of transform audio data [frames, bins]
+    :param spectral_mask: Spectral mask parameters
+    :param seed: Random number seed
+    :return: Augmented feature
+    """
+    import numpy as np
+
+    if audio_f.ndim != 2:
+        raise ValueError("feature input must have three dimensions [frames, bins]")
+
+    frames, bins = audio_f.shape
+
+    f_max_width = spectral_mask.f_max_width
+    if f_max_width not in range(0, bins + 1):
+        f_max_width = bins
+
+    rng = np.random.default_rng(seed)
+
+    # apply f_num frequency masks to the feature
+    for _ in range(spectral_mask.f_num):
+        f_width = int(rng.uniform(0, f_max_width))
+        f_start = rng.integers(0, bins - f_width, endpoint=True)
+        audio_f[:, f_start : f_start + f_width] = 0
+
+    # apply t_num time masks to the feature
+    t_upper_bound = int(spectral_mask.t_max_percent / 100 * frames)
+    for _ in range(spectral_mask.t_num):
+        t_width = min(int(rng.uniform(0, spectral_mask.t_max_width)), t_upper_bound)
+        t_start = rng.integers(0, frames - t_width, endpoint=True)
+        audio_f[t_start : t_start + t_width, :] = 0
+
+    return audio_f

sonusai/mixture/truth.py

@@ -0,0 +1,61 @@
+from ..datatypes import Truth
+from ..datatypes import TruthsDict
+from .mixdb import MixtureDatabase
+
+
+def truth_function(mixdb: MixtureDatabase, m_id: int) -> TruthsDict:
+    from ..datatypes import TruthDict
+    from . import truth_functions
+
+    result: TruthsDict = {}
+    for category, source in mixdb.mixture(m_id).all_sources.items():
+        truth: TruthDict = {}
+        source_file = mixdb.source_file(source.file_id)
+        for name, config in source_file.truth_configs.items():
+            try:
+                truth[name] = getattr(truth_functions, config.function)(mixdb, m_id, category, config.config)
+            except AttributeError as e:
+                raise AttributeError(f"Unsupported truth function: {config.function}") from e
+            except Exception as e:
+                raise RuntimeError(f"Error in truth function '{config.function}': {e}") from e
+
+        if truth:
+            result[category] = truth
+
+    return result
+
+
+def get_class_indices_for_mixid(mixdb: MixtureDatabase, mixid: int) -> list[int]:
+    """Get a list of class indices for a given mixid."""
+    indices: list[int] = []
+    for source_id in [source.file_id for source in mixdb.mixture(mixid).all_sources.values()]:
+        indices.append(*mixdb.source_file(source_id).class_indices)
+
+    return sorted(set(indices))
+
+
+def truth_stride_reduction(truth: Truth, function: str) -> Truth:
+    """Reduce stride dimension of truth.
+
+    :param truth: Truth data [frames, stride, truth_parameters]
+    :param function: Truth stride reduction function name
+    :return: Stride reduced truth data [frames, stride or 1, truth_parameters]
+    """
+    import numpy as np
+
+    if truth.ndim != 3:
+        raise ValueError("Invalid truth shape")
+
+    if function == "none":
+        return truth
+
+    if function == "max":
+        return np.max(truth, axis=1, keepdims=True)
+
+    if function == "mean":
+        return np.mean(truth, axis=1, keepdims=True)
+
+    if function == "first":
+        return truth[:, 0, :].reshape((truth.shape[0], 1, truth.shape[2]))
+
+    raise ValueError(f"Invalid truth stride reduction function: {function}")

sonusai/mixture/truth_functions/__init__.py

@@ -0,0 +1,45 @@
+# SonusAI truth functions
+# ruff: noqa: F401
+
+from .crm import crm
+from .crm import crm_parameters
+from .crm import crm_validate
+from .crm import crmp
+from .crm import crmp_parameters
+from .crm import crmp_validate
+from .energy import energy_f
+from .energy import energy_f_parameters
+from .energy import energy_f_validate
+from .energy import energy_t
+from .energy import energy_t_parameters
+from .energy import energy_t_validate
+from .energy import mapped_snr_f
+from .energy import mapped_snr_f_parameters
+from .energy import mapped_snr_f_validate
+from .energy import snr_f
+from .energy import snr_f_parameters
+from .energy import snr_f_validate
+from .file import file
+from .file import file_parameters
+from .file import file_validate
+from .metadata import metadata
+from .metadata import metadata_parameters
+from .metadata import metadata_validate
+from .metrics import metrics
+from .metrics import metrics_parameters
+from .metrics import metrics_validate
+from .phoneme import phoneme
+from .phoneme import phoneme_parameters
+from .phoneme import phoneme_validate
+from .sed import sed
+from .sed import sed_parameters
+from .sed import sed_validate
+from .target import target_f
+from .target import target_f_parameters
+from .target import target_f_validate
+from .target import target_mixture_f
+from .target import target_mixture_f_parameters
+from .target import target_mixture_f_validate
+from .target import target_swin_f
+from .target import target_swin_f_parameters
+from .target import target_swin_f_validate

sonusai/mixture/truth_functions/crm.py

@@ -0,0 +1,105 @@
+from ...datatypes import Truth
+from ..mixdb import MixtureDatabase
+
+
+def _core(mixdb: MixtureDatabase, m_id: int, category: str, parameters: int, polar: bool) -> Truth:
+    import numpy as np
+    import torch
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+    from pyaaware import feature_inverse_transform_config
+
+    source_audio = torch.from_numpy(mixdb.mixture_sources(m_id)[category])
+    t_ft = ForwardTransform(**feature_forward_transform_config(mixdb.feature))
+    n_ft = ForwardTransform(**feature_forward_transform_config(mixdb.feature))
+
+    frames = t_ft.frames(source_audio)
+    if mixdb.mixture(m_id).all_sources[category].snr_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
+
+    noise_audio = torch.from_numpy(mixdb.mixture_noise(m_id))
+
+    frame_size = feature_inverse_transform_config(mixdb.feature)["overlap"]
+
+    frames = len(source_audio) // frame_size
+    truth = np.empty((frames, t_ft.bins * 2), dtype=np.float32)
+    for frame in range(frames):
+        offset = frame * frame_size
+        target_f = t_ft.execute(source_audio[offset : offset + frame_size])[0].numpy().astype(np.complex64)
+        noise_f = n_ft.execute(noise_audio[offset : offset + frame_size])[0].numpy().astype(np.complex64)
+        mixture_f = target_f + noise_f
+
+        crm_data = np.empty(target_f.shape, dtype=np.complex64)
+        with np.nditer(target_f, flags=["multi_index"], op_flags=[["readwrite"]]) as it:
+            for _ in it:
+                num = target_f[it.multi_index]
+                den = mixture_f[it.multi_index]
+                if num == 0:
+                    crm_data[it.multi_index] = 0
+                elif den == 0:
+                    crm_data[it.multi_index] = complex(np.inf, np.inf)
+                else:
+                    crm_data[it.multi_index] = num / den
+
+        truth[frame, : t_ft.bins] = np.absolute(crm_data) if polar else np.real(crm_data)
+        truth[frame, t_ft.bins :] = np.angle(crm_data) if polar else np.imag(crm_data)
+
+    return truth
+
+
+def crm_validate(_config: dict) -> None:
+    pass
+
+
+def crm_parameters(feature: str, _num_classes: int, _config: dict) -> int:
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    return ForwardTransform(**feature_forward_transform_config(feature)).bins * 2
+
+
+def crm(mixdb: MixtureDatabase, m_id: int, category: str, _config: dict) -> Truth:
+    """Complex ratio mask truth generation function
+
+    Calculates the true complex ratio mask (CRM) truth which is a complex number
+    per bin = Mr + j*Mi. For a given noisy STFT bin value Y, it is used as
+
+    (Mr*Yr + Mi*Yi) / (Yr^2 + Yi^2) + j*(Mi*Yr - Mr*Yi)/ (Yr^2 + Yi^2)
+
+    Output shape: [:, 2 * bins]
+    """
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        category=category,
+        parameters=crm_parameters(mixdb.feature, mixdb.num_classes, _config),
+        polar=False,
+    )
+
+
+def crmp_validate(_config: dict) -> None:
+    pass
+
+
+def crmp_parameters(feature: str, _num_classes: int, _config: dict) -> int:
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    return ForwardTransform(**feature_forward_transform_config(feature)).bins * 2
+
+
+def crmp(mixdb: MixtureDatabase, m_id: int, category: str, _config: dict) -> Truth:
+    """Complex ratio mask polar truth generation function
+
+    Same as the crm function except the results are magnitude and phase
+    instead of real and imaginary.
+
+    Output shape: [:, bins]
+    """
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        category=category,
+        parameters=crmp_parameters(mixdb.feature, mixdb.num_classes, _config),
+        polar=True,
+    )

sonusai/mixture/truth_functions/energy.py

@@ -0,0 +1,222 @@
+import numpy as np
+
+from ...datatypes import Truth
+from ...utils.load_object import load_object
+from ..mixdb import MixtureDatabase
+
+
+def _core(
+    mixdb: MixtureDatabase,
+    m_id: int,
+    category: str,
+    config: dict,
+    parameters: int,
+    mapped: bool,
+    snr: bool,
+    use_cache: bool = True,
+) -> Truth:
+    from os.path import join
+
+    import torch
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    from ...utils.energy_f import compute_energy_f
+
+    source_audio = mixdb.mixture_sources(m_id)[category]
+    ft = ForwardTransform(**feature_forward_transform_config(mixdb.feature))
+
+    frames = ft.frames(torch.from_numpy(source_audio))
+
+    if mixdb.mixture(m_id).all_sources[category].snr_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
+
+    noise_audio = mixdb.mixture_noise(m_id)
+
+    source_energy = compute_energy_f(time_domain=source_audio, transform=ft)
+    noise_energy = None
+    if snr:
+        noise_energy = compute_energy_f(time_domain=noise_audio, transform=ft)
+
+    frames = len(source_energy)
+    truth = np.empty((frames, ft.bins), dtype=np.float32)
+    for frame in range(frames):
+        tmp = source_energy[frame]
+
+        if noise_energy is not None:
+            old_err = np.seterr(divide="ignore", invalid="ignore")
+            tmp /= noise_energy[frame]
+            np.seterr(**old_err)
+
+        tmp = np.nan_to_num(tmp, nan=-np.inf, posinf=np.inf, neginf=-np.inf)
+
+        if mapped:
+            snr_db_mean = load_object(join(mixdb.location, config["snr_db_mean"]), use_cache)
+            snr_db_std = load_object(join(mixdb.location, config["snr_db_std"]), use_cache)
+            tmp = _calculate_mapped_snr_f(tmp, snr_db_mean, snr_db_std)
+
+        truth[frame] = tmp
+
+    return truth
+
+
+def _calculate_mapped_snr_f(truth_f: np.ndarray, snr_db_mean: np.ndarray, snr_db_std: np.ndarray) -> np.ndarray:
+    """Calculate mapped SNR from standard SNR energy per bin/class."""
+    import scipy.special as sc
+
+    old_err = np.seterr(divide="ignore", invalid="ignore")
+    num = 10 * np.log10(np.double(truth_f)) - np.double(snr_db_mean)
+    den = np.double(snr_db_std) * np.sqrt(2)
+    q = num / den
+    q = np.nan_to_num(q, nan=-np.inf, posinf=np.inf, neginf=-np.inf)
+    result = 0.5 * (1 + sc.erf(q))
+    np.seterr(**old_err)
+
+    return result.astype(np.float32)
+
+
+def energy_f_validate(_config: dict) -> None:
+    pass
+
+
+def energy_f_parameters(feature: str, _num_classes: int, _config: dict) -> int:
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    return ForwardTransform(**feature_forward_transform_config(feature)).bins
+
+
+def energy_f(mixdb: MixtureDatabase, m_id: int, category: str, config: dict, use_cache: bool = True) -> Truth:
+    """Frequency domain energy truth generation function
+
+    Calculates the true energy per bin:
+
+    Ti^2 + Tr^2
+
+    where T is the target STFT bin values.
+
+    Output shape: [:, bins]
+    """
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        category=category,
+        config=config,
+        parameters=energy_f_parameters(mixdb.feature, mixdb.num_classes, config),
+        mapped=False,
+        snr=False,
+        use_cache=use_cache,
+    )
+
+
+def snr_f_validate(_config: dict) -> None:
+    pass
+
+
+def snr_f_parameters(feature: str, _num_classes: int, _config: dict) -> int:
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    return ForwardTransform(**feature_forward_transform_config(feature)).bins
+
+
+def snr_f(mixdb: MixtureDatabase, m_id: int, category: str, config: dict, use_cache: bool = True) -> Truth:
+    """Frequency domain SNR truth function documentation
+
+    Calculates the true SNR per bin:
+
+    (Ti^2 + Tr^2) / (Ni^2 + Nr^2)
+
+    where T is the target and N is the noise STFT bin values.
+
+    Output shape: [:, bins]
+    """
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        category=category,
+        config=config,
+        parameters=snr_f_parameters(mixdb.feature, mixdb.num_classes, config),
+        mapped=False,
+        snr=True,
+        use_cache=use_cache,
+    )
+
+
+def mapped_snr_f_validate(config: dict) -> None:
+    if len(config) == 0:
+        raise AttributeError("mapped_snr_f truth function is missing config")
+
+    for parameter in ("snr_db_mean", "snr_db_std"):
+        if parameter not in config:
+            raise AttributeError(f"mapped_snr_f truth function is missing required '{parameter}'")
+
+
+def mapped_snr_f_parameters(feature: str, _num_classes: int, _config: dict) -> int:
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    return ForwardTransform(**feature_forward_transform_config(feature)).bins
+
+
+def mapped_snr_f(mixdb: MixtureDatabase, m_id: int, category: str, config: dict, use_cache: bool = True) -> Truth:
+    """Frequency domain mapped SNR truth function documentation
+
+    Output shape: [:, bins]
+    """
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        category=category,
+        config=config,
+        parameters=mapped_snr_f_parameters(mixdb.feature, mixdb.num_classes, config),
+        mapped=True,
+        snr=True,
+        use_cache=use_cache,
+    )
+
+
+def energy_t_validate(_config: dict) -> None:
+    pass
+
+
+def energy_t_parameters(_feature: str, _num_classes: int, _config: dict) -> int:
+    return 1
+
+
+def energy_t(mixdb: MixtureDatabase, m_id: int, category: str, _config: dict) -> Truth:
+    """Time domain energy truth function documentation
+
+    Calculates the true time domain energy of each frame:
+
+    For OLS:
+        sum(x[0:N-1]^2) / N
+
+    For OLA:
+        sum(x[0:R-1]^2) / R
+
+    where x is the target time domain data,
+    N is the size of the transform, and
+    R is the number of new samples in the frame.
+
+    Output shape: [:, 1]
+
+    Note: feature transforms can be defined to use a subset of all bins,
+    i.e., subset of 0:128 for N=256 could be 0:127 or 1:128. energy_t
+    will reflect the total energy over all bins regardless of the feature
+    transform config.
+    """
+    import torch
+    from pyaaware import ForwardTransform
+    from pyaaware import feature_forward_transform_config
+
+    source_audio = torch.from_numpy(mixdb.mixture_sources(m_id)[category])
+
+    ft = ForwardTransform(**feature_forward_transform_config(mixdb.feature))
+
+    frames = ft.frames(source_audio)
+    parameters = energy_f_parameters(mixdb.feature, mixdb.num_classes, _config)
+    if mixdb.mixture(m_id).all_sources[category].snr_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
+
+    return ft.execute_all(source_audio)[1].numpy()

sonusai/mixture/truth_functions/file.py

@@ -0,0 +1,48 @@
+from ...datatypes import Truth
+from ..mixdb import MixtureDatabase
+
+
+def file_validate(config: dict) -> None:
+    import h5py
+
+    if len(config) == 0:
+        raise AttributeError("file truth function is missing config")
+
+    if "file" not in config:
+        raise AttributeError("file truth function is missing required 'file'")
+
+    with h5py.File(config["file"], "r") as f:
+        if "truth_f" not in f:
+            raise ValueError("Truth file does not contain truth_f dataset")
+
+
+def file_parameters(_feature: str, _num_classes: int, config: dict) -> int:
+    import h5py
+    import numpy as np
+
+    with h5py.File(config["file"], "r") as f:
+        truth = np.array(f["truth_f"])
+
+    return truth.shape[-1]
+
+
+def file(mixdb: MixtureDatabase, m_id: int, category: str, config: dict) -> Truth:
+    """file truth function documentation"""
+    import h5py
+    import numpy as np
+    from pyaaware import feature_inverse_transform_config
+
+    source_audio = mixdb.mixture_sources(m_id)[category]
+
+    frame_size = feature_inverse_transform_config(mixdb.feature)["overlap"]
+
+    with h5py.File(config["file"], "r") as f:
+        truth = np.array(f["truth_f"])
+
+    if truth.ndim != 2:
+        raise ValueError("Truth file data is not 2 dimensions")
+
+    if truth.shape[0] != len(source_audio) // frame_size:
+        raise ValueError("Truth file does not contain the right amount of frames")
+
+    return truth

sonusai/mixture/truth_functions/metadata.py

@@ -0,0 +1,24 @@
+from ...datatypes import Truth
+from ..mixdb import MixtureDatabase
+
+
+def metadata_validate(config: dict) -> None:
+    if len(config) == 0:
+        raise AttributeError("metadata truth function is missing config")
+
+    parameters = ["tier"]
+    for parameter in parameters:
+        if parameter not in config:
+            raise AttributeError(f"metadata truth function is missing required '{parameter}'")
+
+
+def metadata_parameters(_feature: str, _num_classes: int, _config: dict) -> int | None:
+    return None
+
+
+def metadata(mixdb: MixtureDatabase, m_id: int, category: str, config: dict) -> Truth:
+    """Metadata truth generation function
+
+    Retrieves metadata from target.
+    """
+    return mixdb.mixture_speech_metadata(m_id, config["tier"])[category]

sonusai/mixture/truth_functions/metrics.py

@@ -0,0 +1,28 @@
+from ...datatypes import Truth
+from ..mixdb import MixtureDatabase
+
+
+def metrics_validate(config: dict) -> None:
+    if len(config) == 0:
+        raise AttributeError("metrics truth function is missing config")
+
+    parameters = ["metric"]
+    for parameter in parameters:
+        if parameter not in config:
+            raise AttributeError(f"metrics truth function is missing required '{parameter}'")
+
+
+def metrics_parameters(_feature: str, _num_classes: int, _config: dict) -> int | None:
+    return None
+
+
+def metrics(mixdb: MixtureDatabase, m_id: int, category: str, config: dict) -> Truth:
+    """Metadata truth generation function
+
+    Retrieves metrics from target.
+    """
+    if not isinstance(config["metric"], list):
+        m = [config["metric"]]
+    else:
+        m = config["metric"]
+    return mixdb.mixture_metrics(m_id, m)[m[0]][category]

sonusai/mixture/truth_functions/phoneme.py

@@ -0,0 +1,18 @@
+from ...datatypes import Truth
+from ..mixdb import MixtureDatabase
+
+
+def phoneme_validate(_config: dict) -> None:
+    raise NotImplementedError("Truth function phoneme is not supported yet")
+
+
+def phoneme_parameters(_feature: str, _num_classes: int, _config: dict) -> int:
+    raise NotImplementedError("Truth function phoneme is not supported yet")
+
+
+def phoneme(_mixdb: MixtureDatabase, _m_id: int, _category: str, _config: dict) -> Truth:
+    """Read in .txt transcript and run a Python function to generate text grid data
+    (indicating which phonemes are active). Then generate truth based on this data and put
+    in the correct classes based on the index in the config.
+    """
+    raise NotImplementedError("Truth function phoneme is not supported yet")