sonusai 0.20.3__py3-none-any.whl → 1.0.2__py3-none-any.whl

sonusai/{mixture/datatypes.py → datatypes.py}

@@ -15,6 +15,7 @@ AudioT: TypeAlias = npt.NDArray[np.float32]
 
 Truth: TypeAlias = Any
 TruthDict: TypeAlias = dict[str, Truth]
+TruthsDict: TypeAlias = dict[str, TruthDict]
 Segsnr: TypeAlias = npt.NDArray[np.float32]
 
 AudioF: TypeAlias = npt.NDArray[np.complex64]
@@ -68,63 +69,34 @@ class TruthConfig(DataClassSonusAIMixin):
 
 
 TruthConfigs: TypeAlias = dict[str, TruthConfig]
+TruthsConfigs: TypeAlias = dict[str, TruthConfigs]
+
+
 NumberStr: TypeAlias = float | int | str
 OptionalNumberStr: TypeAlias = NumberStr | None
 OptionalListNumberStr: TypeAlias = list[NumberStr] | None
-EQ: TypeAlias = tuple[float | int, float | int, float | int]
-
-
-@dataclass
-class AugmentationRuleEffects(DataClassSonusAIMixin):
-    normalize: OptionalNumberStr = None
-    pitch: OptionalNumberStr = None
-    tempo: OptionalNumberStr = None
-    gain: OptionalNumberStr = None
-    eq1: OptionalListNumberStr = None
-    eq2: OptionalListNumberStr = None
-    eq3: OptionalListNumberStr = None
-    lpf: OptionalNumberStr = None
-    ir: OptionalNumberStr = None
-
-
-@dataclass
-class AugmentationRule(DataClassSonusAIMixin):
-    pre: AugmentationRuleEffects
-    post: AugmentationRuleEffects | None = None
-    mixup: int = 1
 
 
-@dataclass
-class AugmentationEffects(DataClassSonusAIMixin):
-    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
+EffectList: TypeAlias = list[str]
 
 
 @dataclass
-class Augmentation(DataClassSonusAIMixin):
-    pre: AugmentationEffects
-    post: AugmentationEffects
+class Effects(DataClassSonusAIMixin):
+    pre: EffectList
+    post: EffectList = field(default_factory=EffectList)
 
 
-@dataclass(frozen=True)
 class UniversalSNRGenerator:
-    is_random: bool
-    _raw_value: float | str
+    def __init__(self, raw_value: float | str) -> None:
+        self._raw_value = str(raw_value)
+        self.is_random = isinstance(raw_value, str) and raw_value.startswith("sai_rand")
 
     @property
     def value(self) -> float:
-        if self.is_random:
-            from .augmentation import evaluate_random_rule
-
-            return float(evaluate_random_rule(str(self._raw_value)))
+        from sonusai.mixture.effects import evaluate_sai_random_float
 
+        if self.is_random:
+            return float(evaluate_sai_random_float(self._raw_value))
         return float(self._raw_value)
 
 
@@ -145,12 +117,14 @@ Speaker: TypeAlias = dict[str, str]
 
 
 @dataclass
-class TargetFile(DataClassSonusAIMixin):
+class SourceFile(DataClassSonusAIMixin):
+    category: str
+    class_indices: list[int]
     name: str
     samples: int
-    class_indices: list[int]
     truth_configs: TruthConfigs
-    class_balancing_augmentation: AugmentationRule | None = None
+    class_balancing_effect: EffectList | None = None
+    id: int = -1
     level_type: str | None = None
     speaker_id: int | None = None
 
@@ -162,21 +136,9 @@ class TargetFile(DataClassSonusAIMixin):
 
 
 @dataclass
-class AugmentedTarget(DataClassSonusAIMixin):
-    target_id: int
-    target_augmentation_id: int
-
-
-@dataclass
-class NoiseFile(DataClassSonusAIMixin):
-    name: str
-    samples: int
-
-    @property
-    def duration(self) -> float:
-        from .constants import SAMPLE_RATE
-
-        return self.samples / SAMPLE_RATE
+class EffectedFile(DataClassSonusAIMixin):
+    file_id: int
+    effect_id: int
 
 
 ClassCount: TypeAlias = list[int]
@@ -184,37 +146,6 @@ ClassCount: TypeAlias = list[int]
 GeneralizedIDs: TypeAlias = str | int | list[int] | range
 
 
-@dataclass
-class GenMixData:
-    targets: list[AudioT] | None = None
-    target: AudioT | None = None
-    noise: AudioT | None = None
-    mixture: AudioT | None = None
-    truth_t: list[TruthDict] | None = None
-    segsnr_t: Segsnr | None = None
-
-
-@dataclass
-class GenFTData:
-    feature: Feature | None = None
-    truth_f: TruthDict | None = None
-    segsnr: Segsnr | None = None
-
-
-@dataclass
-class ImpulseResponseData:
-    data: AudioT
-    sample_rate: int
-    delay: int
-
-
-@dataclass
-class ImpulseResponseFile:
-    file: str
-    tags: list[str]
-    delay: int
-
-
 @dataclass(frozen=True)
 class SpectralMask(DataClassSonusAIMixin):
     f_max_width: int
@@ -226,68 +157,70 @@ class SpectralMask(DataClassSonusAIMixin):
 
 @dataclass(frozen=True)
 class TruthParameter(DataClassSonusAIMixin):
+    category: str
     name: str
     parameters: int | None
 
 
 @dataclass
-class Target(DataClassSonusAIMixin):
+class Source(DataClassSonusAIMixin):
+    effects: Effects
     file_id: int
-    augmentation: Augmentation
-
-    @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.pre.gain is None:
-            return 1.0
-        return round(10 ** (self.augmentation.pre.gain / 20), ndigits=5)
+    pre_tempo: float = 1
+    repeat: bool = False
+    snr: UniversalSNR = field(default_factory=lambda: UniversalSNR(0))
+    snr_gain: float = 0
+    start: int = 0
 
 
-Targets: TypeAlias = list[Target]
-
-
-@dataclass
-class Noise(DataClassSonusAIMixin):
-    file_id: int
-    augmentation: Augmentation
+Sources: TypeAlias = dict[str, Source]
+SourcesAudioT: TypeAlias = dict[str, AudioT]
+SourcesAudioF: TypeAlias = dict[str, AudioF]
 
 
 @dataclass
 class Mixture(DataClassSonusAIMixin):
     name: str
-    targets: list[Target]
-    noise: Noise
-    noise_offset: int
     samples: int
-    snr: UniversalSNR
+    all_sources: Sources
     spectral_mask_id: int
     spectral_mask_seed: int
-    target_snr_gain: float = 1.0
-    noise_snr_gain: float = 1.0
+
+    @property
+    def all_source_ids(self) -> dict[str, int]:
+        return {category: source.file_id for category, source in self.all_sources.items()}
+
+    @property
+    def sources(self) -> Sources:
+        return {category: source for category, source in self.all_sources.items() if category != "noise"}
+
+    @property
+    def source_ids(self) -> dict[str, int]:
+        return {category: source.file_id for category, source in self.sources.items()}
+
+    @property
+    def noise(self) -> Source:
+        return self.all_sources["noise"]
 
     @property
     def noise_id(self) -> int:
         return self.noise.file_id
 
     @property
-    def target_ids(self) -> list[int]:
-        return [target.file_id for target in self.targets]
+    def source_effects(self) -> dict[str, Effects]:
+        return {category: source.effects for category, source in self.sources.items()}
 
     @property
-    def target_augmentations(self) -> list[Augmentation]:
-        return [target.augmentation for target in self.targets]
+    def noise_effects(self) -> Effects:
+        return self.noise.effects
 
     @property
     def is_noise_only(self) -> bool:
-        return self.snr < -96
+        return self.noise.snr < -96
 
     @property
-    def is_target_only(self) -> bool:
-        return self.snr > 96
-
-    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
+    def is_source_only(self) -> bool:
+        return self.noise.snr > 96
 
 
 @dataclass(frozen=True)
@@ -302,7 +235,7 @@ class TransformConfig:
 @dataclass(frozen=True)
 class FeatureGeneratorConfig:
     feature_mode: str
-    truth_parameters: dict[str, int | None]
+    truth_parameters: dict[str, dict[str, int | None]]
 
 
 @dataclass(frozen=True)
@@ -319,6 +252,37 @@ class FeatureGeneratorInfo:
 ASRConfigs: TypeAlias = dict[str, dict[str, Any]]
 
 
+@dataclass
+class GenMixData:
+    mixture: AudioT | None = None
+    truth_t: TruthsDict | None = None
+    segsnr_t: Segsnr | None = None
+    sources: SourcesAudioT | None = None
+    source: AudioT | None = None
+    noise: AudioT | None = None
+
+
+@dataclass
+class GenFTData:
+    feature: Feature | None = None
+    truth_f: TruthsDict | None = None
+    segsnr: Segsnr | None = None
+
+
+@dataclass
+class ImpulseResponseData:
+    data: AudioT
+    sample_rate: int
+    delay: int
+
+
+@dataclass
+class ImpulseResponseFile(DataClassSonusAIMixin):
+    name: str
+    tags: list[str]
+    delay: str | int = "auto"
+
+
 @dataclass
 class MixtureDatabaseConfig(DataClassSonusAIMixin):
     asr_configs: ASRConfigs
@@ -326,13 +290,11 @@ class MixtureDatabaseConfig(DataClassSonusAIMixin):
     class_labels: list[str]
     class_weights_threshold: list[float]
     feature: str
-    impulse_response_files: list[ImpulseResponseFile]
+    ir_files: list[ImpulseResponseFile]
     mixtures: list[Mixture]
-    noise_mix_mode: str
-    noise_files: list[NoiseFile]
     num_classes: int
+    source_files: dict[str, list[SourceFile]]
     spectral_masks: list[SpectralMask]
-    target_files: list[TargetFile]
 
 
 SpeechMetadata: TypeAlias = str | list[Interval] | None
@@ -353,7 +315,6 @@ class SnrFBinMetrics(NamedTuple):
 
 
 class SpeechMetrics(NamedTuple):
-    pesq: float | None = None
     csig: float | None = None
     cbak: float | None = None
     covl: float | None = None

sonusai/deprecated/plot.py

@@ -45,11 +45,10 @@ import signal
 
 import numpy as np
 from matplotlib import pyplot as plt
-
-from sonusai.mixture import AudioT
-from sonusai.mixture import Feature
-from sonusai.mixture import Predict
-from sonusai.mixture import Truth
+from sonusai.datatypes import AudioT
+from sonusai.datatypes import Feature
+from sonusai.datatypes import Predict
+from sonusai.datatypes import Truth
 
 
 def signal_handler(_sig, _frame):

sonusai/doc/doc.py

@@ -1,4 +1,4 @@
-from sonusai.mixture import get_default_config
+from ..mixture.config import get_default_config
 
 
 def doc_seed() -> str:
@@ -124,7 +124,7 @@ in this dataset.
 
 
 def get_truth_functions() -> str:
-    from sonusai.mixture import truth_functions
+    from ..mixture import truth_functions
 
     functions = [function for function in dir(truth_functions) if not function.startswith("__")]
     text = "\nSupported truth functions:\n\n"
@@ -471,7 +471,7 @@ Rules must specify all the following parameters:
 
 
 def doc_config() -> str:
-    from sonusai.mixture import VALID_CONFIGS
+    from ..mixture.constants import VALID_CONFIGS
 
     text = "\n"
     text += "The SonusAI database is defined using a config.yml file.\n\n"
@@ -482,7 +482,7 @@ def doc_config() -> str:
 
 
 def doc_asr_configs() -> str:
-    from sonusai.utils import get_available_engines
+    from ..utils.asr import get_available_engines
 
     default = f"\nDefault value: {get_default_config()['asr_configs']}"
     engines = get_available_engines()

sonusai/doc.py

@@ -13,16 +13,16 @@ Show SonusAI documentation.
 def main() -> None:
     from docopt import docopt
 
-    import sonusai
+    from sonusai import __version__ as sai_version
     from sonusai.utils import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     from sonusai import doc
 
     topic = args["TOPIC"]
 
-    print(f"SonusAI {sonusai.__version__} Documentation")
+    print(f"SonusAI {sai_version} Documentation")
     print("")
 
     topics = sorted([item[4:] for item in dir(doc) if item.startswith("doc_")])
@@ -42,4 +42,11 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    from sonusai import exception_handler
+    from sonusai.utils import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)

sonusai/genft.py

@@ -25,22 +25,8 @@ Outputs the following to the mixture database directory:
 
 """
 
-import signal
-
-from sonusai.mixture import GeneralizedIDs
-from sonusai.mixture import GenFTData
-
-
-def signal_handler(_sig, _frame):
-    import sys
-
-    from sonusai import logger
-
-    logger.info("Canceled due to keyboard interrupt")
-    sys.exit(1)
-
-
-signal.signal(signal.SIGINT, signal_handler)
+from sonusai.datatypes import GeneralizedIDs
+from sonusai.datatypes import GenFTData
 
 
 def genft(
@@ -48,7 +34,7 @@ def genft(
     mixids: GeneralizedIDs = "*",
     compute_truth: bool = True,
     compute_segsnr: bool = False,
-    write: bool = False,
+    cache: bool = False,
     show_progress: bool = False,
     force: bool = True,
     no_par: bool = False,
@@ -69,8 +55,8 @@ def genft(
             location=location,
             compute_truth=compute_truth,
             compute_segsnr=compute_segsnr,
+            cache=cache,
             force=force,
-            write=write,
         ),
         mixids,
         progress=progress,
@@ -82,33 +68,42 @@ def genft(
 
 
 def _genft_kernel(
-    m_id: int, location: str, compute_truth: bool, compute_segsnr: bool, force: bool, write: bool
+    m_id: int,
+    location: str,
+    compute_truth: bool,
+    compute_segsnr: bool,
+    cache: bool,
+    force: bool,
 ) -> GenFTData:
+    from functools import partial
+    from typing import Any
+
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import write_cached_data
     from sonusai.mixture import write_mixture_metadata
 
     mixdb = MixtureDatabase(location)
 
+    write_func = partial(write_cached_data, location=mixdb.location, name="mixture", index=mixdb.mixture(m_id).name)
+
     result = GenFTData()
 
-    feature, truth_f = mixdb.mixture_ft(m_id=m_id, force=force)
+    mixture = mixdb.mixture_mixture(m_id)
+
+    feature, truth_f = mixdb.mixture_ft(m_id=m_id, mixture=mixture, force=force)
     result.feature = feature
-    if write:
-        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("feature", feature)])
+    items: dict[str, Any] = {"feature": feature}
 
     if compute_truth:
         result.truth_f = truth_f
-        if write:
-            write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("truth_f", truth_f)])
+        items["truth_f"] = truth_f
 
     if compute_segsnr:
-        segsnr = mixdb.mixture_segsnr(m_id=m_id, force=force)
-        result.segsnr = segsnr
-        if write:
-            write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("segsnr", segsnr)])
+        segsnr_t = mixdb.mixture_segsnr_t(m_id)
+        result.segsnr = mixdb.mixture_segsnr(m_id=m_id, segsnr_t=segsnr_t, force=force, cache=cache)
 
-    if write:
+    if cache:
+        write_func(items=items)
         write_mixture_metadata(mixdb, m_id=m_id)
 
     return result
@@ -117,10 +112,10 @@ def _genft_kernel(
 def main() -> None:
     from docopt import docopt
 
-    import sonusai
+    from sonusai import __version__ as sai_version
     from sonusai.utils import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     import time
     from os.path import join
@@ -129,7 +124,7 @@ def main() -> None:
     from sonusai import initial_log_messages
     from sonusai import logger
     from sonusai import update_console_handler
-    from sonusai.mixture import SAMPLE_RATE
+    from sonusai.constants import SAMPLE_RATE
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import check_audio_files_exist
     from sonusai.utils import human_readable_size
@@ -166,18 +161,14 @@ def main() -> None:
 
     check_audio_files_exist(mixdb)
 
-    try:
-        genft(
-            location=location,
-            mixids=mixids,
-            compute_segsnr=compute_segsnr,
-            write=True,
-            show_progress=True,
-            no_par=no_par,
-        )
-    except Exception as e:
-        logger.debug(e)
-        raise
+    genft(
+        location=location,
+        mixids=mixids,
+        compute_segsnr=compute_segsnr,
+        cache=True,
+        show_progress=True,
+        no_par=no_par,
+    )
 
     logger.info(f"Wrote {len(mixids)} mixtures to {location}")
     logger.info("")
@@ -195,4 +186,11 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    from sonusai import exception_handler
+    from sonusai.utils import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)

sonusai/genmetrics.py

@@ -1,12 +1,13 @@
 """sonusai genmetrics
 
-usage: genmetrics [-hvusd] [-i MIXID] [-n INCLUDE] [-x EXCLUDE] LOC
+usage: genmetrics [-hvusd] [-i MIXID] [-n INCLUDE] [-p NUMPROC] [-x EXCLUDE] LOC
 
 options:
     -h, --help
     -v, --verbose                   Be verbose.
     -i MIXID, --mixid MIXID         Mixture ID(s) to generate. [default: *].
     -n INCLUDE, --include INCLUDE   Metrics to include. [default: all]
+    -p NUMPROC, --nproc NUMPROC     Number of parallel processes to use. Default single thread.
     -x EXCLUDE, --exclude EXCLUDE   Metrics to exclude. [default: none]
     -u, --update                    Update metrics (do not regenerate existing metrics).
     -s, --supported                 Show list of supported metrics.
@@ -46,28 +47,16 @@ Generate all available metrics except for mxcovl
 
 """
 
-import signal
-
-
-def signal_handler(_sig, _frame):
-    import sys
-
-    from sonusai import logger
-
-    logger.info("Canceled due to keyboard interrupt")
-    sys.exit(1)
-
-
-signal.signal(signal.SIGINT, signal_handler)
-
 
 def _process_mixture(mixid: int, location: str, metrics: list[str], update: bool = False) -> set[str]:
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import write_cached_data
+    from sonusai.mixture import write_mixture_metadata
 
     mixdb = MixtureDatabase(location)
     results = mixdb.mixture_metrics(m_id=mixid, metrics=metrics, force=not update)
-    write_cached_data(mixdb.location, "mixture", mixdb.mixture(mixid).name, list(results.items()))
+    write_cached_data(mixdb.location, "mixture", mixdb.mixture(mixid).name, results)
+    write_mixture_metadata(mixdb, mixture=mixdb.mixture(mixid))
 
     return set(results.keys())
 
@@ -75,11 +64,11 @@ def _process_mixture(mixid: int, location: str, metrics: list[str], update: bool
 def main() -> None:
     from docopt import docopt
 
-    import sonusai
+    from sonusai import __version__ as sai_version
     from sonusai.mixture import MixtureDatabase
     from sonusai.utils import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     verbose = args["--verbose"]
     mixids = args["--mixid"]
@@ -88,6 +77,7 @@ def main() -> None:
     update = args["--update"]
     show_supported = args["--supported"]
     dryrun = args["--dryrun"]
+    num_proc = args["--nproc"]
     location = args["LOC"]
 
     import fnmatch
@@ -157,11 +147,20 @@ def main() -> None:
     logger.info("")
     logger.info(f"Found {len(mixids):,} mixtures to process")
 
+    if num_proc is None or len(mixids) == 1:
+        no_par = True
+        num_proc = None
+    else:
+        no_par = False
+        num_proc = int(num_proc)  # TBD add support for 'auto'
+
     progress = track(total=len(mixids), desc="genmetrics")
     results = par_track(
         partial(_process_mixture, location=location, metrics=metrics, update=update),
         mixids,
         progress=progress,
+        num_cpus=num_proc,
+        no_par=no_par,
     )
     progress.close()
 
@@ -176,4 +175,11 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    from sonusai import exception_handler
+    from sonusai.utils import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)