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

sonusai/mixture/soundfile_audio.py

@@ -32,7 +32,7 @@ def _raw_read(name: str | Path) -> tuple[AudioT, int]:
         else:
             raise OSError(f"Error reading {name}: {e}") from e
 
-    return np.squeeze(raw[:, 0]), sample_rate
+    return np.squeeze(raw[:, 0].astype(np.float32)), sample_rate
 
 
 def get_sample_rate(name: str | Path) -> int:

sonusai/mixture/sox_audio.py

@@ -207,7 +207,7 @@ class Transformer(SoxTransformer):
 
         return self
 
-    def build(
+    def build(  # pyright: ignore [reportIncompatibleMethodOverride]
         self,
         input_filepath: str | Path | None = None,
         output_filepath: str | Path | None = None,
@@ -320,11 +320,11 @@ class Transformer(SoxTransformer):
         logger.info("Created %s with effects: %s", output_filepath, " ".join(self.effects_log))
 
         if return_output:
-            return status, out, err
+            return status, out, err  # pyright: ignore [reportReturnType]
 
         return True, None, None
 
-    def build_array(
+    def build_array(  # pyright: ignore [reportIncompatibleMethodOverride]
         self,
         input_filepath: str | Path | None = None,
         input_array: np.ndarray | None = None,
@@ -465,7 +465,7 @@ class Transformer(SoxTransformer):
         if status != 0:
             raise SoxError(f"Stdout: {out}\nStderr: {err}")
 
-        out = np.frombuffer(out, dtype=encoding_out)
+        out = np.frombuffer(out, dtype=encoding_out)  # pyright: ignore [reportArgumentType, reportCallIssue]
         if output_format["channels"] > 1:
             out = out.reshape(
                 (output_format["channels"], int(len(out) / output_format["channels"])),

sonusai/mixture/sox_augmentation.py

@@ -118,7 +118,7 @@ def apply_impulse_response(audio: AudioT, ir: ImpulseResponseData) -> AudioT:
     # Apply IR and convert back to global sample rate
     tfm = Transformer()
     tfm.set_output_format(rate=SAMPLE_RATE)
-    tfm.fir(coefficients=temp.name)
+    tfm.fir(coefficients=temp.name)  # pyright: ignore [reportArgumentType]
     try:
         audio_out = tfm.build_array(input_array=audio_out, sample_rate_in=ir.sample_rate)
     except Exception as e:

sonusai/mixture/target_class_balancing.py

@@ -1,19 +1,17 @@
 from sonusai.mixture.datatypes import AugmentationRule
-from sonusai.mixture.datatypes import AugmentationRules
-from sonusai.mixture.datatypes import AugmentedTargets
+from sonusai.mixture.datatypes import AugmentedTarget
 from sonusai.mixture.datatypes import TargetFile
-from sonusai.mixture.datatypes import TargetFiles
 
 
 def balance_targets(
-    augmented_targets: AugmentedTargets,
-    targets: TargetFiles,
-    target_augmentations: AugmentationRules,
+    augmented_targets: list[AugmentedTarget],
+    targets: list[TargetFile],
+    target_augmentations: list[AugmentationRule],
     class_balancing_augmentation: AugmentationRule,
     num_classes: int,
     num_ir: int,
     mixups: list[int] | None = None,
-) -> tuple[AugmentedTargets, AugmentationRules]:
+) -> tuple[list[AugmentedTarget], list[AugmentationRule]]:
     import math
 
     from .augmentation import get_mixups
@@ -64,15 +62,15 @@ def balance_targets(
 
 
 def _get_unused_balancing_augmentation(
-    augmented_targets: AugmentedTargets,
-    targets: TargetFiles,
-    target_augmentations: AugmentationRules,
+    augmented_targets: list[AugmentedTarget],
+    targets: list[TargetFile],
+    target_augmentations: list[AugmentationRule],
     class_balancing_augmentation: AugmentationRule,
     target_id: int,
     mixup: int,
     num_ir: int,
     first_cba_id: int,
-) -> tuple[int, AugmentationRules]:
+) -> tuple[int, list[AugmentationRule]]:
     """Get an unused balancing augmentation for a given target file index"""
     from dataclasses import asdict
 

sonusai/mixture/targets.py

@@ -1,21 +1,20 @@
-from sonusai.mixture.datatypes import AugmentationRules
+from sonusai.mixture.datatypes import AugmentationRule
 from sonusai.mixture.datatypes import AugmentedTarget
-from sonusai.mixture.datatypes import AugmentedTargets
-from sonusai.mixture.datatypes import TargetFiles
+from sonusai.mixture.datatypes import TargetFile
 
 
 def get_augmented_targets(
-    target_files: TargetFiles,
-    target_augmentations: AugmentationRules,
+    target_files: list[TargetFile],
+    target_augmentations: list[AugmentationRule],
     mixups: list[int] | None = None,
-) -> AugmentedTargets:
+) -> list[AugmentedTarget]:
     from .augmentation import get_augmentation_indices_for_mixup
     from .augmentation import get_mixups
 
     if mixups is None:
         mixups = get_mixups(target_augmentations)
 
-    augmented_targets: AugmentedTargets = []
+    augmented_targets: list[AugmentedTarget] = []
     for mixup in mixups:
         augmentation_indices = get_augmentation_indices_for_mixup(target_augmentations, mixup)
         for target_index in range(len(target_files)):
@@ -30,15 +29,17 @@ def get_augmented_targets(
     return augmented_targets
 
 
-def get_class_index_for_augmented_target(augmented_target: AugmentedTarget, targets: TargetFiles) -> list[int]:
+def get_class_index_for_augmented_target(augmented_target: AugmentedTarget, targets: list[TargetFile]) -> list[int]:
     return targets[augmented_target.target_id].class_indices
 
 
-def get_mixup_for_augmented_target(augmented_target: AugmentedTarget, augmentations: AugmentationRules) -> int:
+def get_mixup_for_augmented_target(augmented_target: AugmentedTarget, augmentations: list[AugmentationRule]) -> int:
     return augmentations[augmented_target.target_augmentation_id].mixup
 
 
-def get_target_ids_for_class_index(targets: TargetFiles, class_index: int, allow_multiple: bool = False) -> list[int]:
+def get_target_ids_for_class_index(
+    targets: list[TargetFile], class_index: int, allow_multiple: bool = False
+) -> list[int]:
     """Get a list of target indices containing the given class index.
 
     If allow_multiple is True, then include targets that contain multiple class indices.
@@ -55,9 +56,9 @@ def get_target_ids_for_class_index(targets: TargetFiles, class_index: int, allow
 
 
 def get_augmented_target_ids_for_class_index(
-    augmented_targets: AugmentedTargets,
-    targets: TargetFiles,
-    augmentations: AugmentationRules,
+    augmented_targets: list[AugmentedTarget],
+    targets: list[TargetFile],
+    augmentations: list[AugmentationRule],
     class_index: int,
     mixup: int,
     allow_multiple: bool = False,
@@ -79,9 +80,9 @@ def get_augmented_target_ids_for_class_index(
 
 
 def get_augmented_target_ids_by_class(
-    augmented_targets: AugmentedTargets,
-    targets: TargetFiles,
-    target_augmentations: AugmentationRules,
+    augmented_targets: list[AugmentedTarget],
+    targets: list[TargetFile],
+    target_augmentations: list[AugmentationRule],
     mixup: int,
     num_classes: int,
 ) -> list[list[int]]:
@@ -99,7 +100,9 @@ def get_augmented_target_ids_by_class(
     return indices
 
 
-def get_target_augmentations_for_mixup(target_augmentations: AugmentationRules, mixup: int) -> AugmentationRules:
+def get_target_augmentations_for_mixup(
+    target_augmentations: list[AugmentationRule], mixup: int
+) -> list[AugmentationRule]:
     """Get target augmentations for a given mixup value
 
     :param target_augmentations: List of target augmentation rules
@@ -110,9 +113,9 @@ def get_target_augmentations_for_mixup(target_augmentations: AugmentationRules,
 
 
 def get_augmented_target_ids_for_mixup(
-    augmented_targets: AugmentedTargets,
-    targets: TargetFiles,
-    target_augmentations: AugmentationRules,
+    augmented_targets: list[AugmentedTarget],
+    targets: list[TargetFile],
+    target_augmentations: list[AugmentationRule],
     mixup: int,
     num_classes: int,
 ) -> list[list[int]]:

sonusai/mixture/torchaudio_audio.py

@@ -4,10 +4,16 @@ from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import ImpulseResponseData
 
 
-def read_impulse_response(name: str | Path) -> ImpulseResponseData:
+def read_impulse_response(
+    name: str | Path,
+    delay_compensation: bool = True,
+    normalize: bool = True,
+) -> ImpulseResponseData:
     """Read impulse response data using torchaudio
 
     :param name: File name
+    :param delay_compensation: Apply delay compensation
+    :param normalize: Apply normalization
     :return: ImpulseResponseData object
     """
     import numpy as np
@@ -28,14 +34,19 @@ def read_impulse_response(name: str | Path) -> ImpulseResponseData:
             raise OSError(f"Error reading {name}: {e}") from e
 
     raw = torch.squeeze(raw[0, :])
-    offset = torch.argmax(raw)
-    raw = raw[offset:]
-    # Inexplicably, torch.linalg.vector_norm() causes multiprocessing contexts to hang.
-    # Use np.linalg.norm() instead.
-    # raw = raw / torch.linalg.vector_norm(raw)
+
+    if delay_compensation:
+        offset = torch.argmax(raw)
+        raw = raw[offset:]
 
     data = np.array(raw).astype(np.float32)
-    data = data / np.linalg.norm(data)
+
+    if normalize:
+        # Inexplicably,
+        #   data = data / torch.linalg.vector_norm(data)
+        # causes multiprocessing contexts to hang.
+        # Use np.linalg.norm() instead.
+        data = data / np.linalg.norm(data)
 
     return ImpulseResponseData(name=str(name), sample_rate=sample_rate, data=data)
 

sonusai/mixture/torchaudio_augmentation.py

@@ -20,10 +20,9 @@ def apply_augmentation(audio: AudioT, augmentation: Augmentation, frame_length:
 
     effects: list[list[str]] = []
 
-    # TODO
-    #  Always normalize and remove normalize from list of available augmentations
-    #  Normalize to globally set level (should this be a global config parameter,
-    #  or hard-coded into the script?)
+    # TODO: Always normalize and remove normalize from list of available augmentations
+    # Normalize to globally set level (should this be a global config parameter, or hard-coded into the script?)
+    # TODO: Support all sox effects supported by torchaudio (torchaudio.sox_effects.effect_names())
     if augmentation.normalize is not None:
         effects.append(["norm", str(augmentation.normalize)])
 

sonusai/mixture/truth.py

@@ -1,39 +1,26 @@
-from sonusai.mixture.datatypes import AudioT
-from sonusai.mixture.datatypes import Truth
-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: TruthConfig,
-    feature: str,
-    num_classes: int,
-    class_indices: list[int],
-    target_gain: float,
-) -> Truth:
+from sonusai.mixture import MixtureDatabase
+from sonusai.mixture import Truth
+
+
+def truth_function(mixdb: MixtureDatabase, m_id: int) -> list[Truth]:
+    from sonusai.mixture import TruthDict
     from sonusai.mixture import truth_functions
 
-    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, 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
+    result: list[Truth] = []
+    for target_index in range(len(mixdb.mixture(m_id).targets)):
+        truth: TruthDict = {}
+        target_file = mixdb.target_file(mixdb.mixture(m_id).targets[target_index].file_id)
+        for name, config in target_file.truth_configs.items():
+            try:
+                truth[name] = getattr(truth_functions, config.function)(mixdb, m_id, target_index, 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
+
+        result.append(truth)
+
+    return result
 
 
 def get_truth_indices_for_mixid(mixdb: MixtureDatabase, mixid: int) -> list[int]:

sonusai/mixture/truth_functions/__init__.py

@@ -22,6 +22,12 @@ 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

sonusai/mixture/truth_functions/crm.py

@@ -1,22 +1,32 @@
-from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
-from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
+from sonusai.mixture import MixtureDatabase
+from sonusai.mixture import Truth
 
 
-def _core(data: TruthFunctionData, config: TruthFunctionConfig, polar: bool) -> Truth:
+def _core(mixdb: MixtureDatabase, m_id: int, target_index: int, 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
 
-    if config.target_fft.bins != config.noise_fft.bins:
-        raise ValueError("Transform size mismatch for crm truth")
+    target_audio = torch.from_numpy(mixdb.mixture_targets(m_id)[target_index])
+    t_ft = ForwardTransform(**feature_forward_transform_config(mixdb.feature))
+    n_ft = ForwardTransform(**feature_forward_transform_config(mixdb.feature))
 
-    frames = len(data.target_audio) // config.frame_size
-    truth = np.empty((frames, config.target_fft.bins * 2), dtype=np.float32)
+    frames = t_ft.frames(target_audio)
+    if mixdb.mixture(m_id).target_gain(target_index) == 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(target_audio) // frame_size
+    truth = np.empty((frames, t_ft.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)
+        offset = frame * frame_size
+        target_f = t_ft.execute(target_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)
@@ -31,8 +41,8 @@ def _core(data: TruthFunctionData, config: TruthFunctionConfig, polar: bool) ->
                 else:
                     crm_data[it.multi_index] = num / den
 
-        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)
+        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
 
@@ -41,11 +51,14 @@ def crm_validate(_config: dict) -> None:
     pass
 
 
-def crm_parameters(config: TruthFunctionConfig) -> int:
-    return config.target_fft.bins * 2
+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(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
+
+def crm(mixdb: MixtureDatabase, m_id: int, target_index: int, _config: dict) -> Truth:
     """Complex ratio mask truth generation function
 
     Calculates the true complex ratio mask (CRM) truth which is a complex number
@@ -55,25 +68,27 @@ def crm(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
 
     Output shape: [:, 2 * bins]
     """
-    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)
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        target_index=target_index,
+        parameters=crm_parameters(mixdb.feature, mixdb.num_classes, _config),
+        polar=False,
+    )
 
 
 def crmp_validate(_config: dict) -> None:
     pass
 
 
-def crmp_parameters(config: TruthFunctionConfig) -> int:
-    return config.target_fft.bins * 2
+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(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
+
+def crmp(mixdb: MixtureDatabase, m_id: int, target_index: int, _config: dict) -> Truth:
     """Complex ratio mask polar truth generation function
 
     Same as the crm function except the results are magnitude and phase
@@ -81,11 +96,10 @@ def crmp(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
 
     Output shape: [:, bins]
     """
-    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)
+    return _core(
+        mixdb=mixdb,
+        m_id=m_id,
+        target_index=target_index,
+        parameters=crmp_parameters(mixdb.feature, mixdb.num_classes, _config),
+        polar=True,
+    )