sonusai 0.20.2__py3-none-any.whl → 1.0.1__py3-none-any.whl

sonusai/metrics_summary.py

@@ -16,23 +16,9 @@ Inputs:
 
 """
 
-import signal
-
 import numpy as np
 import pandas as pd
 
-
-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)
-
 DB_99 = np.power(10, 99 / 10)
 DB_N99 = np.power(10, -99 / 10)
 
@@ -49,8 +35,8 @@ def _process_mixture(
 ) -> tuple[pd.DataFrame, pd.DataFrame]:
     from os.path import basename
 
+    from sonusai.constants import SAMPLE_RATE
     from sonusai.metrics import calc_wer
-    from sonusai.mixture import SAMPLE_RATE
     from sonusai.mixture import MixtureDatabase
 
     mixdb = MixtureDatabase(location)
@@ -61,11 +47,11 @@ def _process_mixture(
     duration = samples / SAMPLE_RATE
     tf_frames = mixdb.mixture_transform_frames(m_id)
     feat_frames = mixdb.mixture_feature_frames(m_id)
-    mxsnr = mixdb.mixture(m_id).snr
-    ti = mixdb.mixture(m_id).targets[0].file_id
+    mxsnr = mixdb.mixture(m_id).noise.snr
+    ti = mixdb.mixture(m_id).sources["primary"].file_id
     ni = mixdb.mixture(m_id).noise.file_id
-    t0file = basename(mixdb.target_file(ti).name)
-    nfile = basename(mixdb.noise_file(ni).name)
+    t0file = basename(mixdb.source_file(ti).name)
+    nfile = basename(mixdb.source_file(ni).name)
 
     all_metrics = mixdb.mixture_metrics(m_id, all_metric_names)
 
@@ -104,10 +90,10 @@ def _process_mixture(
 def main() -> None:
     from docopt import docopt
 
-    from sonusai import __version__ as sonusai_ver
-    from sonusai.utils import trim_docstring
+    from . import __version__ as sai_version
+    from .utils.docstring import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai_ver, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     verbose = args["--verbose"]
     wrlist = args["--write-list"]
@@ -121,24 +107,22 @@ def main() -> None:
 
     import psutil
 
-    from sonusai import create_file_handler
-    from sonusai import initial_log_messages
-    from sonusai import logger
-    from sonusai import update_console_handler
-    from sonusai.mixture import MixtureDatabase
-    from sonusai.utils import create_timestamp
-    from sonusai.utils import par_track
-    from sonusai.utils import track
+    from . import create_file_handler
+    from . import initial_log_messages
+    from . import logger
+    from . import update_console_handler
+    from .mixture.mixdb import MixtureDatabase
+    from .utils.create_timestamp import create_timestamp
+    from .utils.parallel import par_track
+    from .utils.parallel import track
 
-    try:
-        mixdb = MixtureDatabase(location)
-        print(f"Found SonusAI mixture database with {mixdb.num_mixtures} mixtures.")
-    except:
-        print(f"Could not open SonusAI mixture database in {location}, exiting ...")
-        return
+    mixdb = MixtureDatabase(location)
+    print(f"Found SonusAI mixture database with {mixdb.num_mixtures} mixtures.")
 
     # Only check first and last mixture in order to save time
-    metrics_present = mixdb.cached_metrics([0, mixdb.num_mixtures - 1])
+    metrics_present = mixdb.cached_metrics([0, mixdb.num_mixtures - 1])  # return pre-generated metrics in mixdb tree
+    if "mxsnr" in metrics_present:
+        metrics_present.remove("mxsnr")
 
     num_metrics_present = len(metrics_present)
     if num_metrics_present < 1:
@@ -188,8 +172,8 @@ def main() -> None:
             if len(metval) > 1:
                 logger.warning(f"Mixid {mixids[0]} metric {metric} has a list with more than 1 element, using first.")
             metval = metval[0]  # remove any list
-        if isinstance(metval, float):
-            logger.debug("Metric is scalar float, entering in summary table.")
+        if isinstance(metval, float | int):
+            logger.debug(f"Metric is scalar {type(metval)}, entering in summary table.")
             scalar_metric_names.append(metric)
         elif isinstance(metval, str):
             logger.debug("Metric is string, will summarize with word count.")
@@ -205,7 +189,7 @@ def main() -> None:
                 else:
                     logger.warning(f"Mixid {mixids[0]} metric {metric} is a vector of improper size, ignoring.")
 
-    # Setup pandas table for summarizing scalar metrics
+    # Setup pandas table for summarizing scalar metrics, always include mxsnr first
     ptab_labels = [
         "mxsnr",
         *scalar_metric_names,
@@ -276,7 +260,7 @@ def main() -> None:
         ptab1.round(2).to_csv(wlcsv_name, **table_args)
     ptab1_sorted = ptab1.sort_values(by=["mxsnr", "t0file"])
 
-    # Create metrics table except except -99 SNR
+    # Create metrics table except -99 SNR
     ptab1_nom99 = ptab1_sorted[ptab1_sorted.mxsnr != -99]
 
     # Create summary by SNR for all scalar metrics, taking mean
@@ -294,7 +278,7 @@ def main() -> None:
     nmixtot = mixdb.num_mixtures
     pd.DataFrame([["Timestamp", timestamp]]).to_csv(snrcsv_name, header=False, index=False)
     pd.DataFrame(['"Metrics avg over each SNR:"']).to_csv(snrcsv_name, **header_args)
-    mtab_snr_summary.round(2).to_csv(snrcsv_name, index=False, **table_args)
+    mtab_snr_summary.round(2).T.to_csv(snrcsv_name, index=True, header=False, mode="a", encoding="utf-8")
     pd.DataFrame(["--"]).to_csv(snrcsv_name, header=False, index=False, mode="a")
     pd.DataFrame([f'"Metrics stats over {nmix} mixtures out of {nmixtot} total:"']).to_csv(snrcsv_name, **header_args)
     ptab1.describe().round(2).T.to_csv(snrcsv_name, index=True, **table_args)
@@ -304,12 +288,14 @@ def main() -> None:
     )
     ptab1_nom99.describe().round(2).T.to_csv(snrcsv_name, index=True, **table_args)
 
-    # Write summary to .csv
+    # Write summary to text file
     snrtxt_name = str(join(location, "metric_summary_snr" + fsuffix + ".txt"))
     with open(snrtxt_name, "w") as f:
         print(f"Timestamp: {timestamp}", file=f)
         print("Metrics avg over each SNR:", file=f)
-        print(mtab_snr_summary.round(2).to_string(float_format=lambda x: f"{x:.2f}", index=False), file=f)
+        print(
+            mtab_snr_summary.round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True, header=False), file=f
+        )
         print("", file=f)
         print(f"Metrics stats over {len(mixids)} mixtures out of {mixdb.num_mixtures} total:", file=f)
         print(ptab1.describe().round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True), file=f)
@@ -319,4 +305,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/mixture/__init__.py

@@ -1,131 +1,32 @@
 # SonusAI mixture utilities
 # ruff: noqa: F401
 
-from .audio import get_duration
-from .audio import get_next_noise
-from .audio import get_num_samples
-from .audio import get_sample_rate
-from .audio import raw_read_audio
+from ..datatypes import AudioT
 from .audio import read_audio
-from .audio import read_ir
-from .audio import validate_input_file
-from .augmentation import apply_augmentation
-from .augmentation import apply_gain
-from .augmentation import apply_impulse_response
-from .augmentation import augmentation_from_rule
-from .augmentation import estimate_augmented_length_from_length
-from .augmentation import evaluate_random_rule
-from .augmentation import get_augmentation_indices_for_mixup
-from .augmentation import get_augmentation_rules
-from .augmentation import get_mixups
-from .augmentation import pad_audio_to_length
-from .class_count import get_class_count_from_mixids
-from .config import get_default_config
-from .config import get_impulse_response_files
-from .config import get_noise_files
-from .config import get_spectral_masks
-from .config import get_target_files
-from .config import get_truth_parameters
+from .config import get_ir_files
+from .config import get_source_files
 from .config import load_config
-from .config import raw_load_config
-from .config import update_config_from_file
-from .config import update_config_from_hierarchy
-from .config import validate_truth_configs
-from .constants import BIT_DEPTH
-from .constants import CHANNEL_COUNT
-from .constants import DEFAULT_CONFIG
-from .constants import DEFAULT_NOISE
-from .constants import DEFAULT_SPEECH
-from .constants import ENCODING
-from .constants import FLOAT_BYTES
-from .constants import MIXDB_VERSION
-from .constants import RAND_PATTERN
-from .constants import REQUIRED_CONFIGS
-from .constants import REQUIRED_TRUTH_CONFIGS
-from .constants import SAMPLE_BYTES
-from .constants import SAMPLE_RATE
-from .constants import VALID_AUGMENTATIONS
-from .constants import VALID_CONFIGS
-from .constants import VALID_NOISE_MIX_MODES
-from .data_io import clear_cached_data
 from .data_io import read_cached_data
 from .data_io import write_cached_data
-from .datatypes import AudioF
-from .datatypes import AudioStatsMetrics
-from .datatypes import AudioT
-from .datatypes import Augmentation
-from .datatypes import AugmentationEffects
-from .datatypes import AugmentationRule
-from .datatypes import AugmentationRuleEffects
-from .datatypes import AugmentedTarget
-from .datatypes import ClassCount
-from .datatypes import EnergyF
-from .datatypes import EnergyT
-from .datatypes import Feature
-from .datatypes import FeatureGeneratorConfig
-from .datatypes import FeatureGeneratorInfo
-from .datatypes import GeneralizedIDs
-from .datatypes import GenFTData
-from .datatypes import GenMixData
-from .datatypes import ImpulseResponseData
-from .datatypes import ImpulseResponseFile
-from .datatypes import MetricDoc
-from .datatypes import MetricDocs
-from .datatypes import Mixture
-from .datatypes import MixtureDatabaseConfig
-from .datatypes import NoiseFile
-from .datatypes import Predict
-from .datatypes import Segsnr
-from .datatypes import SnrFMetrics
-from .datatypes import SpectralMask
-from .datatypes import SpeechMetadata
-from .datatypes import SpeechMetrics
-from .datatypes import TargetFile
-from .datatypes import TransformConfig
-from .datatypes import Truth
-from .datatypes import TruthConfig
-from .datatypes import TruthConfigs
-from .datatypes import TruthDict
-from .datatypes import TruthParameter
-from .datatypes import UniversalSNR
+from .effects import get_effect_rules
 from .feature import get_audio_from_feature
 from .feature import get_feature_from_audio
 from .generation import generate_mixtures
-from .generation import get_all_snrs_from_config
 from .generation import initialize_db
 from .generation import populate_class_label_table
 from .generation import populate_class_weights_threshold_table
 from .generation import populate_impulse_response_file_table
 from .generation import populate_mixture_table
-from .generation import populate_noise_file_table
+from .generation import populate_source_file_table
 from .generation import populate_spectral_mask_table
-from .generation import populate_target_file_table
 from .generation import populate_top_table
 from .generation import populate_truth_parameters_table
 from .generation import update_mixid_width
 from .generation import update_mixture
-from .helpers import augmented_noise_samples
-from .helpers import augmented_target_samples
 from .helpers import check_audio_files_exist
 from .helpers import forward_transform
-from .helpers import frames_from_samples
-from .helpers import get_audio_from_transform
-from .helpers import get_transform_from_audio
 from .helpers import inverse_transform
-from .helpers import mixture_metadata
 from .helpers import write_mixture_metadata
-from .ir_delay import get_impulse_response_delay
 from .log_duration_and_sizes import log_duration_and_sizes
 from .mixdb import MixtureDatabase
 from .mixdb import db_file
-from .spectral_mask import apply_spectral_mask
-from .target_class_balancing import balance_targets
-from .targets import get_augmented_target_ids_by_class
-from .targets import get_augmented_target_ids_for_mixup
-from .targets import get_augmented_targets
-from .targets import get_target_augmentations_for_mixup
-from .tokenized_shell_vars import tokenized_expand
-from .tokenized_shell_vars import tokenized_replace
-from .truth import get_truth_indices_for_mixid
-from .truth import truth_function
-from .truth import truth_stride_reduction

sonusai/mixture/audio.py

@@ -1,8 +1,7 @@
 from functools import lru_cache
 from pathlib import Path
 
-from sonusai.mixture.datatypes import AudioT
-from sonusai.mixture.datatypes import ImpulseResponseData
+from ..datatypes import AudioT
 
 
 def get_next_noise(audio: AudioT, offset: int, length: int) -> AudioT:
@@ -24,7 +23,7 @@ def get_duration(audio: AudioT) -> float:
     :param audio: Time domain data [samples]
     :return: Duration of audio in seconds
     """
-    from .constants import SAMPLE_RATE
+    from ..constants import SAMPLE_RATE
 
     return len(audio) / SAMPLE_RATE
 
@@ -66,7 +65,7 @@ def _get_sample_rate(name: str | Path) -> int:
     import soundfile
     from pydub import AudioSegment
 
-    from .tokenized_shell_vars import tokenized_expand
+    from ..utils.tokenized_shell_vars import tokenized_expand
 
     expanded_name, _ = tokenized_expand(name)
 
@@ -90,7 +89,7 @@ def raw_read_audio(name: str | Path) -> tuple[AudioT, int]:
     import soundfile
     from pydub import AudioSegment
 
-    from .tokenized_shell_vars import tokenized_expand
+    from ..utils.tokenized_shell_vars import tokenized_expand
 
     expanded_name, _ = tokenized_expand(name)
 
@@ -135,40 +134,12 @@ def _read_audio(name: str | Path) -> AudioT:
     :param name: File name
     :return: Array of time domain audio data
     """
-    import librosa
+    from ..constants import SAMPLE_RATE
+    from .resample import resample
 
-    from .constants import SAMPLE_RATE
-
-    out, sample_rate = raw_read_audio(name)
-    out = librosa.resample(out, orig_sr=sample_rate, target_sr=SAMPLE_RATE, res_type="soxr_hq")
-
-    return out
-
-
-def read_ir(name: str | Path, delay: int, use_cache: bool = True) -> ImpulseResponseData:
-    """Read impulse response data
-
-    :param name: File name
-    :param delay: Delay in samples
-    :param use_cache: If true, use LRU caching
-    :return: ImpulseResponseData object
-    """
-    if use_cache:
-        return _read_ir(name, delay)
-    return _read_ir.__wrapped__(name, delay)
-
-
-@lru_cache
-def _read_ir(name: str | Path, delay: int) -> ImpulseResponseData:
-    """Read impulse response data using soundfile
-
-    :param name: File name
-    :param delay: Delay in samples
-    :return: ImpulseResponseData object
-    """
     out, sample_rate = raw_read_audio(name)
 
-    return ImpulseResponseData(data=out, sample_rate=sample_rate, delay=delay)
+    return resample(out, orig_sr=sample_rate, target_sr=SAMPLE_RATE)
 
 
 def get_num_samples(name: str | Path, use_cache: bool = True) -> int:
@@ -195,8 +166,8 @@ def _get_num_samples(name: str | Path) -> int:
     import soundfile
     from pydub import AudioSegment
 
-    from .constants import SAMPLE_RATE
-    from .tokenized_shell_vars import tokenized_expand
+    from ..constants import SAMPLE_RATE
+    from ..utils.tokenized_shell_vars import tokenized_expand
 
     expanded_name, _ = tokenized_expand(name)
 
@@ -209,7 +180,7 @@ def _get_num_samples(name: str | Path) -> int:
         samples = sound.frame_count()
         sample_rate = sound.frame_rate
     else:
-        info = soundfile.info(name)
+        info = soundfile.info(expanded_name)
         samples = info.frames
         sample_rate = info.samplerate
 

sonusai/mixture/class_balancing.py

@@ -0,0 +1,103 @@
+from ..datatypes import EffectList
+from ..datatypes import EffectedFile
+from ..datatypes import File
+
+
+def balance_sources(
+    effected_sources: list[EffectedFile],
+    files: list[File],
+    effects: list[EffectList],
+    class_balancing_effect: EffectList,
+    num_classes: int,
+    num_ir: int,
+    mixups: list[int] | None = None,
+) -> tuple[list[EffectedFile], list[EffectList]]:
+    import math
+
+    from .augmentation import get_mixups
+    from .sources import get_augmented_target_ids_by_class
+
+    first_cba_id = len(effects)
+
+    if mixups is None:
+        mixups = get_mixups(effects)
+
+    for mixup in mixups:
+        if mixup == 1:
+            continue
+
+        effected_sources_indices_by_class = get_augmented_target_ids_by_class(
+            augmented_targets=effected_sources,
+            targets=files,
+            target_augmentations=effects,
+            mixup=mixup,
+            num_classes=num_classes,
+        )
+
+        largest = max([len(item) for item in effected_sources_indices_by_class])
+        largest = math.ceil(largest / mixup) * mixup
+        for es_indices in effected_sources_indices_by_class:
+            additional_effects_needed = largest - len(es_indices)
+            file_ids = sorted({effected_sources[at_index].file_id for at_index in es_indices})
+
+            tfi_idx = 0
+            for _ in range(additional_effects_needed):
+                file_id = file_ids[tfi_idx]
+                tfi_idx = (tfi_idx + 1) % len(file_ids)
+                effect_id, effects = _get_unused_balancing_effect(
+                    effected_sources=effected_sources,
+                    files=files,
+                    effects=effects,
+                    class_balancing_effect=class_balancing_effect,
+                    file_id=file_id,
+                    mixup=mixup,
+                    num_ir=num_ir,
+                    first_cbe_id=first_cba_id,
+                )
+                effected_sources.append(EffectedFile(file_id=file_id, effect_id=effect_id))
+
+    return effected_sources, effects
+
+
+def _get_unused_balancing_effect(
+    effected_sources: list[EffectedFile],
+    files: list[File],
+    effects: list[EffectList],
+    class_balancing_effect: EffectList,
+    file_id: int,
+    mixup: int,
+    num_ir: int,
+    first_cbe_id: int,
+) -> tuple[int, list[EffectList]]:
+    """Get an unused balancing augmentation for a given target file index"""
+    from dataclasses import asdict
+
+    from .augmentation import get_augmentation_rules
+
+    balancing_augmentations = [item for item in range(len(effects)) if item >= first_cbe_id]
+    used_balancing_augmentations = [
+        effected_source.effect_id
+        for effected_source in effected_sources
+        if effected_source.file_id == file_id and effected_source.effect_id in balancing_augmentations
+    ]
+
+    augmentation_indices = [
+        item
+        for item in balancing_augmentations
+        if item not in used_balancing_augmentations and effects[item].mixup == mixup
+    ]
+    if len(augmentation_indices) > 0:
+        return augmentation_indices[0], effects
+
+    class_balancing_effect = get_class_balancing_effect(file=files[file_id], default_cbe=class_balancing_effect)
+    new_effect = get_augmentation_rules(rules=asdict(class_balancing_effect), num_ir=num_ir)[0]
+    new_effect.mixup = mixup
+    effects.append(new_effect)
+    return len(effects) - 1, effects
+
+
+def get_class_balancing_effect(file: File, default_cbe: EffectList) -> EffectList:
+    """Get the class balancing effect rule for the given target"""
+    if file.class_balancing_effect is not None:
+        return file.class_balancing_effect
+    return default_cbe