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

sonusai/mixture/datatypes.py

@@ -12,16 +12,12 @@ from dataclasses_json import DataClassJsonMixin
 from praatio.utilities.constants import Interval
 
 AudioT: TypeAlias = npt.NDArray[np.float32]
-AudiosT: TypeAlias = list[AudioT]
 
-ListAudiosT: TypeAlias = list[AudiosT]
-
-Truth: TypeAlias = npt.NDArray[np.float32]
+Truth: TypeAlias = Any
 TruthDict: TypeAlias = dict[str, Truth]
 Segsnr: TypeAlias = npt.NDArray[np.float32]
 
 AudioF: TypeAlias = npt.NDArray[np.complex64]
-AudiosF: TypeAlias = list[AudioF]
 
 EnergyT: TypeAlias = npt.NDArray[np.float32]
 EnergyF: TypeAlias = npt.NDArray[np.float32]
@@ -92,9 +88,6 @@ class AugmentationRule(DataClassSonusAIMixin):
     mixup: int = 1
 
 
-AugmentationRules: TypeAlias = list[AugmentationRule]
-
-
 @dataclass
 class Augmentation(DataClassSonusAIMixin):
     normalize: float | None = None
@@ -108,9 +101,6 @@ class Augmentation(DataClassSonusAIMixin):
     ir: int | None = None
 
 
-Augmentations: TypeAlias = list[Augmentation]
-
-
 @dataclass(frozen=True)
 class UniversalSNRGenerator:
     is_random: bool
@@ -159,18 +149,12 @@ class TargetFile(DataClassSonusAIMixin):
         return self.samples / SAMPLE_RATE
 
 
-TargetFiles: TypeAlias = list[TargetFile]
-
-
 @dataclass
 class AugmentedTarget(DataClassSonusAIMixin):
     target_id: int
     target_augmentation_id: int
 
 
-AugmentedTargets: TypeAlias = list[AugmentedTarget]
-
-
 @dataclass
 class NoiseFile(DataClassSonusAIMixin):
     name: str
@@ -183,7 +167,6 @@ class NoiseFile(DataClassSonusAIMixin):
         return self.samples / SAMPLE_RATE
 
 
-NoiseFiles: TypeAlias = list[NoiseFile]
 ClassCount: TypeAlias = list[int]
 
 GeneralizedIDs: TypeAlias = str | int | list[int] | range
@@ -191,11 +174,11 @@ GeneralizedIDs: TypeAlias = str | int | list[int] | range
 
 @dataclass
 class GenMixData:
-    targets: AudiosT | None = None
+    targets: list[AudioT] | None = None
     target: AudioT | None = None
     noise: AudioT | None = None
     mixture: AudioT | None = None
-    truth_t: TruthDict | None = None
+    truth_t: list[TruthDict] | None = None
     segsnr_t: Segsnr | None = None
 
 
@@ -223,9 +206,6 @@ class ImpulseResponseFile:
     tags: list[str]
 
 
-ImpulseResponseFiles: TypeAlias = list[ImpulseResponseFile]
-
-
 @dataclass(frozen=True)
 class SpectralMask(DataClassSonusAIMixin):
     f_max_width: int
@@ -235,23 +215,24 @@ class SpectralMask(DataClassSonusAIMixin):
     t_max_percent: int
 
 
-SpectralMasks: TypeAlias = list[SpectralMask]
-
-
 @dataclass(frozen=True)
 class TruthParameter(DataClassSonusAIMixin):
     name: str
-    parameters: int
-
-
-TruthParameters: TypeAlias = list[TruthParameter]
+    parameters: int | None
 
 
 @dataclass
 class Target(DataClassSonusAIMixin):
     file_id: int
     augmentation: Augmentation
-    gain: float = 1.0
+
+    @property
+    def gain(self) -> float:
+        # gain is used to back out the gain augmentation in order to return the target audio
+        # to its normalized level when calculating truth (if needed).
+        if self.augmentation.gain is None:
+            return 1.0
+        return round(10 ** (self.augmentation.gain / 20), ndigits=5)
 
 
 Targets: TypeAlias = list[Target]
@@ -261,14 +242,14 @@ Targets: TypeAlias = list[Target]
 class Noise(DataClassSonusAIMixin):
     file_id: int
     augmentation: Augmentation
-    offset: int = 0
 
 
 @dataclass
 class Mixture(DataClassSonusAIMixin):
     name: str
-    targets: Targets
+    targets: list[Target]
     noise: Noise
+    noise_offset: int
     samples: int
     snr: UniversalSNR
     spectral_mask_id: int
@@ -288,8 +269,16 @@ class Mixture(DataClassSonusAIMixin):
     def target_augmentations(self) -> list[Augmentation]:
         return [target.augmentation for target in self.targets]
 
+    @property
+    def is_noise_only(self) -> bool:
+        return self.snr < -96
+
+    @property
+    def is_target_only(self) -> bool:
+        return self.snr > 96
 
-Mixtures: TypeAlias = list[Mixture]
+    def target_gain(self, target_index: int) -> float:
+        return (self.targets[target_index].gain if not self.is_noise_only else 0) * self.target_snr_gain
 
 
 @dataclass(frozen=True)
@@ -304,7 +293,7 @@ class TransformConfig:
 @dataclass(frozen=True)
 class FeatureGeneratorConfig:
     feature_mode: str
-    truth_parameters: dict[str, int]
+    truth_parameters: dict[str, int | None]
 
 
 @dataclass(frozen=True)
@@ -328,13 +317,13 @@ class MixtureDatabaseConfig(DataClassSonusAIMixin):
     class_labels: list[str]
     class_weights_threshold: list[float]
     feature: str
-    impulse_response_files: ImpulseResponseFiles
-    mixtures: Mixtures
+    impulse_response_files: list[ImpulseResponseFile]
+    mixtures: list[Mixture]
     noise_mix_mode: str
-    noise_files: NoiseFiles
+    noise_files: list[NoiseFile]
     num_classes: int
-    spectral_masks: SpectralMasks
-    target_files: TargetFiles
+    spectral_masks: list[SpectralMask]
+    target_files: list[TargetFile]
 
 
 SpeechMetadata: TypeAlias = str | list[Interval] | None

sonusai/mixture/db_datatypes.py

@@ -35,7 +35,7 @@ SpectralMaskRecord = namedtuple(
     ["id", "f_max_width", "f_num", "t_max_width", "t_num", "t_max_percent"],
 )
 
-TargetRecord = namedtuple("TargetRecord", ["id", "file_id", "augmentation", "gain"])
+TargetRecord = namedtuple("TargetRecord", ["id", "file_id", "augmentation"])
 
 MixtureRecord = namedtuple(
     "MixtureRecord",

sonusai/mixture/feature.py

@@ -12,7 +12,6 @@ def get_feature_from_audio(
     :param feature_mode: Feature mode
     :return: Feature data [frames, strides, feature_parameters]
     """
-    import numpy as np
     from pyaaware import FeatureGenerator
 
     from .datatypes import TransformConfig
@@ -31,33 +30,14 @@ def get_feature_from_audio(
         ),
     )
 
-    transform_frames = audio_f.shape[0]
-    feature_frames = transform_frames // (fg.decimation * fg.step)
-    feature = np.empty((feature_frames, fg.stride, fg.feature_parameters), dtype=np.float32)
-
-    feature_frame = 0
-    for transform_frame in range(transform_frames):
-        fg.execute(audio_f[transform_frame])
-
-        if fg.eof():
-            feature[feature_frame] = fg.feature()
-            feature_frame += 1
+    return fg.execute_all(audio_f)[0]
 
-    return feature
 
-
-def get_audio_from_feature(
-    feature: Feature,
-    feature_mode: str,
-    num_classes: int | None = 1,
-    truth_mutex: bool | None = False,
-) -> AudioT:
+def get_audio_from_feature(feature: Feature, feature_mode: str) -> AudioT:
     """Apply inverse transform to feature data to generate audio data
 
     :param feature: Feature data [frames, stride=1, feature_parameters]
     :param feature_mode: Feature mode
-    :param num_classes: Number of classes
-    :param truth_mutex: Whether to calculate 'other' label
     :return: Audio data [samples]
     """
     import numpy as np
@@ -75,7 +55,7 @@ def get_audio_from_feature(
     if feature.shape[1] != 1:
         raise ValueError("Strided feature data is not supported for audio extraction; stride must be 1.")
 
-    fg = FeatureGenerator(feature_mode=feature_mode, num_classes=num_classes, truth_mutex=truth_mutex)
+    fg = FeatureGenerator(feature_mode=feature_mode)
 
     feature_complex = unstack_complex(feature.squeeze())
     if feature_mode[0:1] == "h":