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

sonusai/mixture/constants.py

@@ -1,49 +1,60 @@
 import re
-
 from importlib.resources import as_file
 from importlib.resources import files
 
 REQUIRED_CONFIGS = [
-    'asr_configs',
-    'class_balancing',
-    'class_balancing_augmentation',
-    'class_labels',
-    'class_weights_threshold',
-    'feature',
-    'impulse_responses',
-    'noise_augmentations',
-    'noise_mix_mode',
-    'noises',
-    'num_classes',
-    'random_snrs',
-    'seed',
-    'snrs',
-    'spectral_masks',
-    'target_augmentations',
-    'target_level_type',
-    'targets',
-    'truth_mode',
-    'truth_reduction_function',
-    'truth_settings',
+    "asr_configs",
+    "class_balancing",
+    "class_balancing_augmentation",
+    "class_indices",
+    "class_labels",
+    "class_weights_threshold",
+    "feature",
+    "impulse_responses",
+    "noise_augmentations",
+    "noise_mix_mode",
+    "noises",
+    "num_classes",
+    "random_snrs",
+    "seed",
+    "snrs",
+    "spectral_masks",
+    "target_augmentations",
+    "target_level_type",
+    "targets",
+    "truth_configs",
 ]
 OPTIONAL_CONFIGS: list[str] = []
 VALID_CONFIGS = REQUIRED_CONFIGS + OPTIONAL_CONFIGS
-VALID_TRUTH_SETTINGS = ['function', 'config', 'index']
-VALID_AUGMENTATIONS = ['normalize', 'gain', 'pitch', 'tempo', 'eq1', 'eq2', 'eq3', 'lpf', 'ir', 'mixup']
-VALID_NOISE_MIX_MODES = ['exhaustive', 'non-exhaustive', 'non-combinatorial']
-RAND_PATTERN = re.compile(r'rand\(([-+]?(\d+(\.\d*)?|\.\d+)),\s*([-+]?(\d+(\.\d*)?|\.\d+))\)')
+REQUIRED_TRUTH_CONFIGS = ["function", "stride_reduction"]
+REQUIRED_ASR_CONFIGS = ["engine"]
+VALID_AUGMENTATIONS = [
+    "normalize",
+    "gain",
+    "pitch",
+    "tempo",
+    "eq1",
+    "eq2",
+    "eq3",
+    "lpf",
+    "ir",
+    "mixup",
+]
+VALID_NOISE_MIX_MODES = ["exhaustive", "non-exhaustive", "non-combinatorial"]
+RAND_PATTERN = re.compile(r"rand\(([-+]?(\d+(\.\d*)?|\.\d+)),\s*([-+]?(\d+(\.\d*)?|\.\d+))\)")
 SAMPLE_RATE = 16000
 BIT_DEPTH = 32
-ENCODING = 'floating-point'
+ENCODING = "floating-point"
 CHANNEL_COUNT = 1
 SAMPLE_BYTES = BIT_DEPTH // 8
 FLOAT_BYTES = 4
+MIXDB_VERSION = 2
 
-with as_file(files('sonusai.data').joinpath('whitenoise.wav')) as path:
+with as_file(files("sonusai.data").joinpath("whitenoise.wav")) as path:
     DEFAULT_NOISE = str(path)
 
-with as_file(files('sonusai.data').joinpath('genmixdb.yml')) as path:
+with as_file(files("sonusai.data").joinpath("genmixdb.yml")) as path:
     DEFAULT_CONFIG = str(path)
 
-with as_file(files('sonusai.data').joinpath('speech_ma01_01.wav')) as path:
+with as_file(files("sonusai.data").joinpath("speech_ma01_01.wav")) as path:
     DEFAULT_SPEECH = str(path)

sonusai/mixture/data_io.py

@@ -0,0 +1,155 @@
+from typing import Any
+
+
+def _get_hdf5_name(location: str, index: str) -> str:
+    from os.path import join
+
+    return join(location, index + ".h5")
+
+
+def _get_pickle_name(location: str, index: str, item: str) -> str:
+    from os.path import join
+
+    return join(location, index, item + ".pkl")
+
+
+def read_hdf5_data(location: str, index: str, items: list[str] | str) -> Any:
+    """Read mixture, target, or noise data from an HDF5 file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: String(s) of data to retrieve
+    :return: Data (or tuple of data)
+    """
+    from os.path import exists
+    from typing import Any
+
+    import h5py
+    import numpy as np
+
+    def _get_dataset(file: h5py.File, d_name: str) -> Any:
+        if d_name in file:
+            data = np.array(file[d_name])
+            if data.size == 1:
+                item = data.item()
+                if isinstance(item, bytes):
+                    return item.decode("utf-8")
+                return item
+            return data
+        return None
+
+    if not isinstance(items, list):
+        items = [items]
+
+    h5_name = _get_hdf5_name(location, index)
+    if exists(h5_name):
+        try:
+            with h5py.File(h5_name, "r") as f:
+                result = [_get_dataset(f, item) for item in items]
+        except Exception as e:
+            raise OSError(f"Error reading {h5_name}: {e}") from e
+    else:
+        result = [None for _ in items]
+
+    if len(items) == 1:
+        result = result[0]
+
+    return result
+
+
+def write_hdf5_data(location: str, index: str, items: list[tuple[str, Any]] | tuple[str, Any]) -> None:
+    """Write mixture, target, or noise data to an HDF5 file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: Tuple(s) of (name, data)
+    """
+    import h5py
+
+    if not isinstance(items, list):
+        items = [items]
+
+    h5_name = _get_hdf5_name(location, index)
+    with h5py.File(h5_name, "a") as f:
+        for item in items:
+            if item[0] in f:
+                del f[item[0]]
+            f.create_dataset(name=item[0], data=item[1])
+
+
+def read_pickle_data(location: str, index: str, items: list[str] | str) -> Any:
+    """Read mixture, target, or noise data from a pickle file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: String(s) of data to retrieve
+    :return: Data (or tuple of data)
+    """
+    import pickle
+    from os.path import exists
+    from typing import Any
+
+    if not isinstance(items, list):
+        items = [items]
+
+    result: list[Any] = []
+    for item in items:
+        pkl_name = _get_pickle_name(location, index, item)
+        if exists(pkl_name):
+            with open(pkl_name, "rb") as f:
+                result.append(pickle.load(f))  # noqa: S301
+        else:
+            result.append(None)
+
+    if len(items) == 1:
+        result = result[0]
+
+    return result
+
+
+def write_pickle_data(location: str, index: str, items: list[tuple[str, Any]] | tuple[str, Any]) -> None:
+    """Write mixture, target, or noise data to a pickle file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: Tuple(s) of (name, data)
+    """
+    import pickle
+    from os import makedirs
+    from os.path import join
+
+    if not isinstance(items, list):
+        items = [items]
+
+    makedirs(join(location, index), exist_ok=True)
+    for item in items:
+        pkl_name = _get_pickle_name(location, index, item[0])
+        with open(pkl_name, "wb") as f:
+            f.write(pickle.dumps(item[1]))
+
+
+def read_cached_data(location: str, name: str, index: str, items: list[str] | str) -> Any:
+    """Read cached data from a file
+
+    :param location: Location of the mixture database
+    :param name: Data name ('mixture', 'target', or 'noise')
+    :param index: Data index (mixture, target, or noise ID)
+    :param items: String(s) of data to retrieve
+    :return: Data (or tuple of data)
+    """
+    from os.path import join
+
+    return read_pickle_data(join(location, name), index, items)
+
+
+def write_cached_data(location: str, name: str, index: str, items: list[tuple[str, Any]] | tuple[str, Any]) -> None:
+    """Write mixture data to a file
+
+    :param location: Location of the mixture database
+    :param name: Data name ('mixture', 'target', or 'noise')
+    :param index: Data index (mixture, target, or noise ID)
+    :param items: Tuple(s) of (name, data)
+    """
+    from os.path import join
+
+    write_pickle_data(join(location, name), index, items)

sonusai/mixture/datatypes.py

@@ -1,8 +1,8 @@
+from collections.abc import Iterable
 from dataclasses import dataclass
+from dataclasses import field
 from typing import Any
-from typing import Iterable
 from typing import NamedTuple
-from typing import Optional
 from typing import SupportsIndex
 from typing import TypeAlias
 
@@ -17,6 +17,7 @@ AudiosT: TypeAlias = list[AudioT]
 ListAudiosT: TypeAlias = list[AudiosT]
 
 Truth: TypeAlias = npt.NDArray[np.float32]
+TruthDict: TypeAlias = dict[str, Truth]
 Segsnr: TypeAlias = npt.NDArray[np.float32]
 
 AudioF: TypeAlias = npt.NDArray[np.complex64]
@@ -35,7 +36,7 @@ Json: TypeAlias = dict | list | str | int | float | bool | None
 
 class DataClassSonusAIMixin(DataClassJsonMixin):
     def __str__(self):
-        return f'{self.to_dict()}'
+        return f"{self.to_dict()}"
 
     # Override DataClassJsonMixin to remove dictionary keys with values of None
     def to_dict(self, encode_json=False) -> dict[str, Json]:
@@ -58,22 +59,22 @@ class DataClassSonusAIMixin(DataClassJsonMixin):
 
 
 @dataclass(frozen=True)
-class TruthSetting(DataClassSonusAIMixin):
-    config: Optional[dict] = None
-    function: Optional[str] = None
-    index: Optional[list[int]] = None
+class TruthConfig(DataClassSonusAIMixin):
+    function: str
+    stride_reduction: str
+    config: dict = field(default_factory=dict)
 
     def __hash__(self):
         return hash(self.to_json())
 
     def __eq__(self, other):
-        return isinstance(other, TruthSetting) and hash(self) == hash(other)
+        return isinstance(other, TruthConfig) and hash(self) == hash(other)
 
 
-TruthSettings: TypeAlias = list[TruthSetting]
+TruthConfigs: TypeAlias = dict[str, TruthConfig]
 NumberStr: TypeAlias = float | int | str
-OptionalNumberStr: TypeAlias = Optional[NumberStr]
-OptionalListNumberStr: TypeAlias = Optional[list[NumberStr]]
+OptionalNumberStr: TypeAlias = NumberStr | None
+OptionalListNumberStr: TypeAlias = list[NumberStr] | None
 EQ: TypeAlias = tuple[float | int, float | int, float | int]
 
 
@@ -88,7 +89,7 @@ class AugmentationRule(DataClassSonusAIMixin):
     eq3: OptionalListNumberStr = None
     lpf: OptionalNumberStr = None
     ir: OptionalNumberStr = None
-    mixup: Optional[int] = 1
+    mixup: int = 1
 
 
 AugmentationRules: TypeAlias = list[AugmentationRule]
@@ -96,15 +97,15 @@ AugmentationRules: TypeAlias = list[AugmentationRule]
 
 @dataclass
 class Augmentation(DataClassSonusAIMixin):
-    normalize: Optional[float] = None
-    pitch: Optional[float] = None
-    tempo: Optional[float] = None
-    gain: Optional[float] = None
-    eq1: Optional[EQ] = None
-    eq2: Optional[EQ] = None
-    eq3: Optional[EQ] = None
-    lpf: Optional[float] = None
-    ir: Optional[int] = None
+    normalize: float | None = None
+    pitch: float | None = None
+    tempo: float | None = None
+    gain: float | None = None
+    eq1: EQ | None = None
+    eq2: EQ | None = None
+    eq3: EQ | None = None
+    lpf: float | None = None
+    ir: int | None = None
 
 
 Augmentations: TypeAlias = list[Augmentation]
@@ -145,10 +146,11 @@ Speaker: TypeAlias = dict[str, str]
 class TargetFile(DataClassSonusAIMixin):
     name: str
     samples: int
-    truth_settings: TruthSettings
-    class_balancing_augmentation: Optional[AugmentationRule] = None
-    level_type: Optional[str] = None
-    speaker_id: Optional[int] = None
+    class_indices: list[int]
+    truth_configs: TruthConfigs
+    class_balancing_augmentation: AugmentationRule | None = None
+    level_type: str | None = None
+    speaker_id: int | None = None
 
     @property
     def duration(self) -> float:
@@ -187,32 +189,21 @@ ClassCount: TypeAlias = list[int]
 GeneralizedIDs: TypeAlias = str | int | list[int] | range
 
 
-@dataclass(frozen=True)
-class TruthFunctionConfig(DataClassSonusAIMixin):
-    feature: str
-    mutex: bool
-    num_classes: int
-    target_gain: float
-    config: Optional[dict] = None
-    function: Optional[str] = None
-    index: Optional[list[int]] = None
-
-
 @dataclass
 class GenMixData:
-    targets: Optional[AudiosT] = None
-    target: Optional[AudioT] = None
-    noise: Optional[AudioT] = None
-    mixture: Optional[AudioT] = None
-    truth_t: Optional[Truth] = None
-    segsnr_t: Optional[Segsnr] = None
+    targets: AudiosT | None = None
+    target: AudioT | None = None
+    noise: AudioT | None = None
+    mixture: AudioT | None = None
+    truth_t: TruthDict | None = None
+    segsnr_t: Segsnr | None = None
 
 
 @dataclass
 class GenFTData:
-    feature: Optional[Feature] = None
-    truth_f: Optional[Truth] = None
-    segsnr: Optional[Segsnr] = None
+    feature: Feature | None = None
+    truth_f: TruthDict | None = None
+    segsnr: Segsnr | None = None
 
 
 @dataclass
@@ -226,7 +217,13 @@ class ImpulseResponseData:
         return len(self.data)
 
 
-ImpulseResponseFiles: TypeAlias = list[str]
+@dataclass
+class ImpulseResponseFile:
+    file: str
+    tags: list[str]
+
+
+ImpulseResponseFiles: TypeAlias = list[ImpulseResponseFile]
 
 
 @dataclass(frozen=True)
@@ -241,11 +238,20 @@ class SpectralMask(DataClassSonusAIMixin):
 SpectralMasks: TypeAlias = list[SpectralMask]
 
 
+@dataclass(frozen=True)
+class TruthParameter(DataClassSonusAIMixin):
+    name: str
+    parameters: int
+
+
+TruthParameters: TypeAlias = list[TruthParameter]
+
+
 @dataclass
 class Target(DataClassSonusAIMixin):
-    file_id: Optional[int] = None
-    augmentation: Optional[Augmentation] = None
-    gain: Optional[float] = None
+    file_id: int
+    augmentation: Augmentation
+    gain: float = 1.0
 
 
 Targets: TypeAlias = list[Target]
@@ -253,22 +259,22 @@ Targets: TypeAlias = list[Target]
 
 @dataclass
 class Noise(DataClassSonusAIMixin):
-    file_id: Optional[int] = None
-    augmentation: Optional[Augmentation] = None
-    offset: Optional[int] = None
+    file_id: int
+    augmentation: Augmentation
+    offset: int = 0
 
 
 @dataclass
 class Mixture(DataClassSonusAIMixin):
-    name: Optional[str] = None
-    noise: Optional[Noise] = None
-    noise_snr_gain: Optional[float] = None
-    samples: Optional[int] = None
-    snr: Optional[UniversalSNR] = None
-    spectral_mask_id: Optional[int] = None
-    spectral_mask_seed: Optional[int] = None
-    target_snr_gain: Optional[float] = None
-    targets: Optional[Targets] = None
+    name: str
+    targets: Targets
+    noise: Noise
+    samples: int
+    snr: UniversalSNR
+    spectral_mask_id: int
+    spectral_mask_seed: int
+    target_snr_gain: float = 1.0
+    noise_snr_gain: float = 1.0
 
     @property
     def noise_id(self) -> int:
@@ -288,8 +294,8 @@ Mixtures: TypeAlias = list[Mixture]
 
 @dataclass(frozen=True)
 class TransformConfig:
-    N: int
-    R: int
+    length: int
+    overlap: int
     bin_start: int
     bin_end: int
     ttype: str
@@ -298,8 +304,7 @@ class TransformConfig:
 @dataclass(frozen=True)
 class FeatureGeneratorConfig:
     feature_mode: str
-    num_classes: int
-    truth_mutex: bool
+    truth_parameters: dict[str, int]
 
 
 @dataclass(frozen=True)
@@ -318,57 +323,55 @@ ASRConfigs: TypeAlias = dict[str, dict[str, Any]]
 
 @dataclass
 class MixtureDatabaseConfig(DataClassSonusAIMixin):
-    asr_configs: Optional[ASRConfigs] = None
-    class_balancing: Optional[bool] = False
-    class_labels: Optional[list[str]] = None
-    class_weights_threshold: Optional[list[float]] = None
-    feature: Optional[str] = None
-    impulse_response_files: Optional[ImpulseResponseFiles] = None
-    mixtures: Optional[Mixtures] = None
-    noise_mix_mode: Optional[str] = 'exhaustive'
-    noise_files: Optional[NoiseFiles] = None
-    num_classes: Optional[int] = None
-    spectral_masks: Optional[SpectralMasks] = None
-    target_files: Optional[TargetFiles] = None
-    truth_mutex: Optional[bool] = None
-    truth_reduction_function: Optional[str] = None
+    asr_configs: ASRConfigs
+    class_balancing: bool
+    class_labels: list[str]
+    class_weights_threshold: list[float]
+    feature: str
+    impulse_response_files: ImpulseResponseFiles
+    mixtures: Mixtures
+    noise_mix_mode: str
+    noise_files: NoiseFiles
+    num_classes: int
+    spectral_masks: SpectralMasks
+    target_files: TargetFiles
 
 
 SpeechMetadata: TypeAlias = str | list[Interval] | None
 
 
 class SnrFMetrics(NamedTuple):
-    avg: Optional[float] = None
-    std: Optional[float] = None
-    db_avg: Optional[float] = None
-    db_std: Optional[float] = None
+    avg: float | None = None
+    std: float | None = None
+    db_avg: float | None = None
+    db_std: float | None = None
 
 
 class SnrFBinMetrics(NamedTuple):
-    avg: Optional[np.ndarray] = None
-    std: Optional[np.ndarray] = None
-    db_avg: Optional[np.ndarray] = None
-    db_std: Optional[np.ndarray] = None
+    avg: np.ndarray | None = None
+    std: np.ndarray | None = None
+    db_avg: np.ndarray | None = None
+    db_std: np.ndarray | None = None
 
 
 class SpeechMetrics(NamedTuple):
-    pesq: Optional[float] = None
-    csig: Optional[float] = None
-    cbak: Optional[float] = None
-    covl: Optional[float] = None
+    pesq: float | None = None
+    csig: float | None = None
+    cbak: float | None = None
+    covl: float | None = None
 
 
 class AudioStatsMetrics(NamedTuple):
-    dco: Optional[float] = None
-    min: Optional[float] = None
-    max: Optional[float] = None
-    pkdb: Optional[float] = None
-    lrms: Optional[float] = None
-    pkr: Optional[float] = None
-    tr: Optional[float] = None
-    cr: Optional[float] = None
-    fl: Optional[float] = None
-    pkc: Optional[float] = None
+    dco: float | None = None
+    min: float | None = None
+    max: float | None = None
+    pkdb: float | None = None
+    lrms: float | None = None
+    pkr: float | None = None
+    tr: float | None = None
+    cr: float | None = None
+    fl: float | None = None
+    pkc: float | None = None
 
 
 @dataclass
@@ -382,7 +385,7 @@ class MetricDocs(list[MetricDoc]):
     def __init__(self, __iterable: Iterable[MetricDoc]) -> None:
         super().__init__(item for item in __iterable)
 
-    def __setitem__(self, __key: SupportsIndex, __value: MetricDoc) -> None:  # type: ignore
+    def __setitem__(self, __key: SupportsIndex, __value: MetricDoc) -> None:  # type: ignore[override]
         super().__setitem__(__key, __value)
 
     def insert(self, __index: SupportsIndex, __object: MetricDoc) -> None:
@@ -406,16 +409,16 @@ class MetricDocs(list[MetricDoc]):
             if item.category not in categories:
                 categories.append(item.category)
 
-        result = ''
+        result = ""
         for category in categories:
-            result += f'{category}\n'
-            result += '-' * max_category_len + '\n'
+            result += f"{category}\n"
+            result += "-" * max_category_len + "\n"
             for item in [sub for sub in self if sub.category == category]:
-                result += f'  {item.name:<{max_name_len}}{item.description}\n'
-            result += '\n'
+                result += f"  {item.name:<{max_name_len}}{item.description}\n"
+            result += "\n"
 
         return result
 
     @property
     def names(self) -> set[str]:
-        return set(item.name for item in self)
+        return {item.name for item in self}