sonusai 0.18.9__py3-none-any.whl → 0.19.6__py3-none-any.whl

sonusai/mixture/truth.py

@@ -1,50 +1,72 @@
 from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.datatypes import TruthFunctionConfig
+from sonusai.mixture.datatypes import TruthConfig
 from sonusai.mixture.mixdb import MixtureDatabase
 
 
-def truth_function(target_audio: AudioT,
-                   noise_audio: AudioT,
-                   mixture_audio: AudioT,
-                   config: TruthFunctionConfig) -> Truth:
-    from sonusai import SonusAIError
+def truth_function(
+    target_audio: AudioT,
+    noise_audio: AudioT,
+    mixture_audio: AudioT,
+    config: TruthConfig,
+    feature: str,
+    num_classes: int,
+    class_indices: list[int],
+    target_gain: float,
+) -> Truth:
     from sonusai.mixture import truth_functions
-    from .truth_functions.data import Data
 
-    data = Data(target_audio, noise_audio, mixture_audio, config)
-    if data.config.target_gain == 0:
-        return data.truth
+    from .truth_functions.datatypes import TruthFunctionConfig
+    from .truth_functions.datatypes import TruthFunctionData
+
+    t_config = TruthFunctionConfig(
+        feature=feature,
+        num_classes=num_classes,
+        class_indices=class_indices,
+        target_gain=target_gain,
+        config=config.config,
+    )
+    t_data = TruthFunctionData(target_audio, noise_audio, mixture_audio)
 
     try:
-        return getattr(truth_functions, data.config.function)(data)
-    except AttributeError:
-        raise SonusAIError(f'Unsupported truth function: {data.config.function}')
+        return getattr(truth_functions, config.function)(t_data, t_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
 
 
 def get_truth_indices_for_mixid(mixdb: MixtureDatabase, mixid: int) -> list[int]:
     """Get a list of truth indices for a given mixid."""
-    from .targets import get_truth_indices_for_target
-
     indices: list[int] = []
     for target_id in [target.file_id for target in mixdb.mixture(mixid).targets]:
-        indices.append(*get_truth_indices_for_target(mixdb.target_file(target_id)))
+        indices.append(*mixdb.target_file(target_id).class_indices)
 
-    return sorted(list(set(indices)))
+    return sorted(set(indices))
 
 
-def truth_reduction(x: Truth, func: str) -> Truth:
+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
 
-    from sonusai import SonusAIError
+    if truth.ndim != 3:
+        raise ValueError("Invalid truth shape")
+
+    if function == "none":
+        return truth
 
-    if func == 'max':
-        return np.max(x, axis=0)
+    if function == "max":
+        return np.max(truth, axis=1, keepdims=True)
 
-    if func == 'mean':
-        return np.mean(x, axis=0)
+    if function == "mean":
+        return np.mean(truth, axis=1, keepdims=True)
 
-    if func == 'index0':
-        return np.squeeze(x[0, :])
+    if function == "first":
+        return truth[:, 0, :].reshape((truth.shape[0], 1, truth.shape[2]))
 
-    raise SonusAIError(f'Invalid truth reduction function: {func}')
+    raise ValueError(f"Invalid truth stride reduction function: {function}")

sonusai/mixture/truth_functions/__init__.py

@@ -1,13 +1,39 @@
 # 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 .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

@@ -1,25 +1,26 @@
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.data import Data
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
 
 
-def _core(data: Data, polar: bool) -> Truth:
+def _core(data: TruthFunctionData, config: TruthFunctionConfig, polar: bool) -> Truth:
     import numpy as np
 
-    from sonusai import SonusAIError
-
-    if data.config.num_classes != data.target_fft.bins:
-        raise SonusAIError(f'Invalid num_classes for crm truth: {data.config.num_classes}')
-
-    if data.target_fft.bins != data.noise_fft.bins:
-        raise SonusAIError('Transform size mismatch for crm truth')
-
-    for offset in data.offsets:
-        target_f = data.target_fft.execute(data.target_audio[offset:offset + data.frame_size]).astype(np.complex64)
-        noise_f = data.noise_fft.execute(data.noise_audio[offset:offset + data.frame_size]).astype(np.complex64)
+    if config.target_fft.bins != config.noise_fft.bins:
+        raise ValueError("Transform size mismatch for crm truth")
+
+    frames = len(data.target_audio) // config.frame_size
+    truth = np.empty((frames, config.target_fft.bins * 2), dtype=np.float32)
+    for frame in range(frames):
+        offset = frame * config.frame_size
+        target_f = config.target_fft.execute(data.target_audio[offset : offset + config.frame_size]).astype(
+            np.complex64
+        )
+        noise_f = config.noise_fft.execute(data.noise_audio[offset : offset + config.frame_size]).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:
+        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]
@@ -30,44 +31,61 @@ def _core(data: Data, polar: bool) -> Truth:
                 else:
                     crm_data[it.multi_index] = num / den
 
-        indices = slice(offset, offset + data.frame_size)
+        truth[frame, : config.target_fft.bins] = np.absolute(crm_data) if polar else np.real(crm_data)
+        truth[frame, config.target_fft.bins :] = np.angle(crm_data) if polar else np.imag(crm_data)
+
+    return truth
 
-        def c1(c_data: np.ndarray, is_polar: bool) -> np.ndarray:
-            if is_polar:
-                return np.absolute(c_data)
-            return np.real(c_data)
 
-        def c2(c_data: np.ndarray, is_polar: bool) -> np.ndarray:
-            if is_polar:
-                return np.angle(c_data)
-            return np.imag(c_data)
+def crm_validate(_config: dict) -> None:
+    pass
 
-        for index in data.zero_based_indices:
-            data.truth[indices, index:index + data.target_fft.bins] = c1(crm_data, polar)
-            data.truth[indices, (index + data.target_fft.bins):(index + 2 * data.target_fft.bins)] = c2(crm_data, polar)
 
-    return data.truth
+def crm_parameters(config: TruthFunctionConfig) -> int:
+    return config.target_fft.bins * 2
 
 
-def crm(data: Data) -> Truth:
+def crm(data: TruthFunctionData, config: TruthFunctionConfig) -> 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
+    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)
+    (Mr*Yr + Mi*Yi) / (Yr^2 + Yi^2) + j*(Mi*Yr - Mr*Yi)/ (Yr^2 + Yi^2)
 
-Output shape: [:, bins]
+    Output shape: [:, 2 * bins]
     """
-    return _core(data=data, polar=False)
+    import numpy as np
+
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = crm_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
+
+    return _core(data=data, config=config, polar=False)
 
 
-def crmp(data: Data) -> Truth:
+def crmp_validate(_config: dict) -> None:
+    pass
+
+
+def crmp_parameters(config: TruthFunctionConfig) -> int:
+    return config.target_fft.bins * 2
+
+
+def crmp(data: TruthFunctionData, config: TruthFunctionConfig) -> 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.
+    Same as the crm function except the results are magnitude and phase
+    instead of real and imaginary.
 
-Output shape: [:, bins]
+    Output shape: [:, bins]
     """
-    return _core(data=data, polar=True)
+    import numpy as np
+
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = crmp_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
+
+    return _core(data=data, config=config, polar=True)

sonusai/mixture/truth_functions/datatypes.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+
+from sonusai.mixture.datatypes import AudioT
+
+
+class TruthFunctionConfig:
+    def __init__(self, feature: str, num_classes: int, class_indices: list[int], target_gain: float, config: dict):
+        from pyaaware import ForwardTransform
+        from pyaaware import InverseTransform
+        from pyaaware import feature_forward_transform_config
+        from pyaaware import feature_inverse_transform_config
+        from pyaaware import feature_parameters
+
+        self.feature = feature
+        self.num_classes = num_classes
+        self.class_indices = class_indices
+        self.target_gain = target_gain
+        self.config = config
+
+        self.feature_parameters = feature_parameters(feature)
+        ft_config = feature_forward_transform_config(feature)
+        it_config = feature_inverse_transform_config(feature)
+
+        self.ttype = it_config["ttype"]
+        self.frame_size = it_config["overlap"]
+
+        self.target_fft = ForwardTransform(**ft_config)
+        self.noise_fft = ForwardTransform(**ft_config)
+        self.mixture_fft = ForwardTransform(**ft_config)
+        self.swin = InverseTransform(**it_config).window
+
+
+@dataclass
+class TruthFunctionData:
+    target_audio: AudioT
+    noise_audio: AudioT
+    mixture_audio: AudioT

sonusai/mixture/truth_functions/energy.py

@@ -1,69 +1,43 @@
 import numpy as np
 
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.data import Data
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
 
 
-def _core(data: Data, mapped: bool, snr: bool) -> Truth:
-    from sonusai import SonusAIError
+def _core(data: TruthFunctionData, config: TruthFunctionConfig, mapped: bool, snr: bool) -> Truth:
     from sonusai.utils import compute_energy_f
 
-    snr_db_mean = None
-    snr_db_std = None
-    if mapped:
-        if data.config.config is None:
-            raise SonusAIError('Truth function mapped SNR missing config')
-
-        parameters = ['snr_db_mean', 'snr_db_std']
-        for parameter in parameters:
-            if parameter not in data.config.config:
-                raise SonusAIError(f'Truth function mapped_snr_f config missing required parameter: {parameter}')
-
-        snr_db_mean = data.config.config['snr_db_mean']
-        if len(snr_db_mean) != data.target_fft.bins:
-            raise SonusAIError(f'Truth function mapped_snr_f snr_db_mean does not have {data.target_fft.bins} elements')
-
-        snr_db_std = data.config.config['snr_db_std']
-        if len(snr_db_std) != data.target_fft.bins:
-            raise SonusAIError(f'Truth function mapped_snr_f snr_db_std does not have {data.target_fft.bins} elements')
-
-    for index in data.zero_based_indices:
-        if index + data.target_fft.bins > data.config.num_classes:
-            raise SonusAIError('Truth index exceeds the number of classes')
-
-    target_energy = compute_energy_f(time_domain=data.target_audio, transform=data.target_fft)
+    target_energy = compute_energy_f(time_domain=data.target_audio, transform=config.target_fft)
     noise_energy = None
     if snr:
-        noise_energy = compute_energy_f(time_domain=data.noise_audio, transform=data.noise_fft)
-
-    if len(target_energy) != len(data.offsets):
-        raise SonusAIError(f'Number of frames in target_energy, {len(target_energy)},'
-                           f' is not number of frames in truth, {len(data.offsets)}')
+        noise_energy = compute_energy_f(time_domain=data.noise_audio, transform=config.noise_fft)
 
-    for idx, offset in enumerate(data.offsets):
-        tmp = target_energy[idx]
+    frames = len(target_energy)
+    truth = np.empty((frames, config.target_fft.bins), dtype=np.float32)
+    for frame in range(frames):
+        tmp = target_energy[frame]
 
-        if snr:
-            old_err = np.seterr(divide='ignore', invalid='ignore')
-            tmp /= noise_energy[idx]
+        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:
-            tmp = _calculate_mapped_snr_f(tmp, snr_db_mean, snr_db_std)
+            tmp = _calculate_mapped_snr_f(tmp, config.config["snr_db_mean"], config.config["snr_db_std"])
 
-        for index in data.zero_based_indices:
-            data.truth[offset:offset + data.frame_size, index:index + data.target_fft.bins] = tmp
+        truth[frame] = tmp
 
-    return data.truth
+    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')
+    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
@@ -74,7 +48,15 @@ def _calculate_mapped_snr_f(truth_f: np.ndarray, snr_db_mean: np.ndarray, snr_db
     return result.astype(np.float32)
 
 
-def energy_f(data: Data) -> Truth:
+def energy_f_validate(_config: dict) -> None:
+    pass
+
+
+def energy_f_parameters(config: TruthFunctionConfig) -> int:
+    return config.target_fft.bins
+
+
+def energy_f(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
     """Frequency domain energy truth generation function
 
     Calculates the true energy per bin:
@@ -85,10 +67,23 @@ def energy_f(data: Data) -> Truth:
 
     Output shape: [:, bins]
     """
-    return _core(data=data, mapped=False, snr=False)
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = energy_f_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
+
+    return _core(data=data, config=config, mapped=False, snr=False)
+
+
+def snr_f_validate(_config: dict) -> None:
+    pass
+
+
+def snr_f_parameters(config: TruthFunctionConfig) -> int:
+    return config.target_fft.bins
 
 
-def snr_f(data: Data) -> Truth:
+def snr_f(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
     """Frequency domain SNR truth function documentation
 
     Calculates the true SNR per bin:
@@ -99,18 +94,54 @@ def snr_f(data: Data) -> Truth:
 
     Output shape: [:, bins]
     """
-    return _core(data=data, mapped=False, snr=True)
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = snr_f_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
 
+    return _core(data=data, config=config, mapped=False, snr=True)
 
-def mapped_snr_f(data: Data) -> Truth:
+
+def mapped_snr_f_validate(config: TruthFunctionConfig) -> None:
+    if len(config.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.config:
+            raise AttributeError(f"mapped_snr_f truth function is missing required '{parameter}'")
+
+        if len(config.config[parameter]) != config.target_fft.bins:
+            raise ValueError(
+                f"mapped_snr_f truth function '{parameter}' does not have {config.target_fft.bins} elements"
+            )
+
+
+def mapped_snr_f_parameters(config: TruthFunctionConfig) -> int:
+    return config.target_fft.bins
+
+
+def mapped_snr_f(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
     """Frequency domain mapped SNR truth function documentation
 
     Output shape: [:, bins]
     """
-    return _core(data=data, mapped=True, snr=True)
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = mapped_snr_f_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
 
+    return _core(data=data, config=config, mapped=True, snr=True)
 
-def energy_t(data: Data) -> Truth:
+
+def energy_t_validate(_config: dict) -> None:
+    pass
+
+
+def energy_t_parameters(_config: TruthFunctionConfig) -> int:
+    return 1
+
+
+def energy_t(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
     """Time domain energy truth function documentation
 
     Calculates the true time domain energy of each frame:
@@ -134,14 +165,9 @@ def energy_t(data: Data) -> Truth:
     """
     import torch
 
-    from sonusai import SonusAIError
-
-    target_energy = data.target_fft.execute_all(torch.from_numpy(data.target_audio))[1].numpy()
-    if len(target_energy) != len(data.offsets):
-        raise SonusAIError(f'Number of frames in target_energy, {len(target_energy)},'
-                           f' is not number of frames in truth, {len(data.offsets)}')
-
-    for offset in data.offsets:
-        data.truth[offset:offset + data.frame_size, data.zero_based_indices] = np.float32(target_energy)
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = energy_t_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
 
-    return data.truth
+    return config.target_fft.execute_all(torch.from_numpy(data.target_audio))[1].numpy()

sonusai/mixture/truth_functions/file.py

@@ -1,44 +1,44 @@
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.data import Data
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
 
 
-def file(data: Data) -> Truth:
-    """file truth function documentation
-    """
+def file_validate(config: dict) -> None:
     import h5py
-    import numpy as np
 
-    from sonusai import SonusAIError
+    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")
+
 
-    if data.config.config is None:
-        raise SonusAIError('Truth function file missing config')
+def file_parameters(config: TruthFunctionConfig) -> int:
+    import h5py
+    import numpy as np
 
-    parameters = ['file']
-    for parameter in parameters:
-        if 'file' not in data.config.config:
-            raise SonusAIError(f'Truth function file config missing required parameter: {parameter}')
+    with h5py.File(config.config["file"], "r") as f:
+        truth = np.array(f["truth_f"])
 
-    with h5py.File(data.config.config['file'], 'r') as f:
-        if 'truth_t' not in f:
-            raise SonusAIError('Truth file does not contain truth_t dataset')
-        truth_in = np.array(f['truth_t'])
+    return truth.shape[-1]
 
-    if truth_in.ndim != 2:
-        raise SonusAIError('Truth file data is not 2 dimensions')
 
-    if truth_in.shape[0] != len(data.target_audio):
-        raise SonusAIError('Truth file does not contain the right amount of samples')
+def file(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
+    """file truth function documentation"""
+    import h5py
+    import numpy as np
 
-    if len(data.zero_based_indices) > 1:
-        if len(data.zero_based_indices) != truth_in.shape[1]:
-            raise SonusAIError('Truth file does not contain the right amount of classes')
+    with h5py.File(config.config["file"], "r") as f:
+        truth = np.array(f["truth_f"])
 
-        data.truth[:, data.zero_based_indices] = truth_in
-    else:
-        index = data.zero_based_indices[0]
-        if index + truth_in.shape[1] > data.config.num_classes:
-            raise SonusAIError('Truth file contains too many classes')
+    if truth.ndim != 2:
+        raise ValueError("Truth file data is not 2 dimensions")
 
-        data.truth[:, index:index + truth_in.shape[1]] = truth_in
+    if truth.shape[0] != len(data.target_audio) // config.frame_size:
+        raise ValueError("Truth file does not contain the right amount of frames")
 
-    return data.truth
+    return truth

sonusai/mixture/truth_functions/phoneme.py

@@ -1,12 +1,19 @@
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.data import Data
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
 
 
-def phoneme(_data: Data) -> Truth:
+def phoneme_validate(_config: dict) -> None:
+    raise NotImplementedError("Truth function phoneme is not supported yet")
+
+
+def phoneme_parameters(_config: TruthFunctionConfig) -> int:
+    raise NotImplementedError("Truth function phoneme is not supported yet")
+
+
+def phoneme(_data: TruthFunctionData, _config: TruthFunctionConfig) -> 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.
+    (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.
     """
-    from sonusai import SonusAIError
-
-    raise SonusAIError('Truth function phoneme is not supported yet')
+    raise NotImplementedError("Truth function phoneme is not supported yet")