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

sonusai/metrics/calc_speech.py

@@ -1,18 +1,23 @@
 import numpy as np
 
-from sonusai.mixture.constants import SAMPLE_RATE
-from sonusai.mixture.datatypes import SpeechMetrics
-
+from ..constants import SAMPLE_RATE
+from ..datatypes import SpeechMetrics
 from .calc_pesq import calc_pesq
 
 
-def calc_speech(hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int = SAMPLE_RATE) -> SpeechMetrics:
-    """Calculate speech metrics pesq, c_sig, c_bak, and c_ovl.
+def calc_speech(
+    hypothesis: np.ndarray,
+    reference: np.ndarray,
+    pesq: float | None = None,
+    sample_rate: int = SAMPLE_RATE,
+) -> SpeechMetrics:
+    """Calculate speech metrics c_sig, c_bak, and c_ovl.
 
     These are all related and thus included in one function. Reference: matlab script "compute_metrics.m".
 
     :param hypothesis: estimated audio
     :param reference: reference audio
+    :param pesq: pesq
     :param sample_rate: sample rate of audio
     :return: SpeechMetrics named tuple
     """
@@ -36,18 +41,21 @@ def calc_speech(hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int
     seg_snr = np.mean(segsnr_dist)
 
     # PESQ
-    _pesq = calc_pesq(hypothesis=hypothesis, reference=reference, sample_rate=sample_rate)
+    if pesq is None:
+        pesq = calc_pesq(hypothesis=hypothesis, reference=reference, sample_rate=sample_rate)
 
     # Now compute the composite measures
-    csig = float(np.clip(3.093 - 1.029 * llr_mean + 0.603 * _pesq - 0.009 * wss_dist, 1, 5))
-    cbak = float(np.clip(1.634 + 0.478 * _pesq - 0.007 * wss_dist + 0.063 * seg_snr, 1, 5))
-    covl = float(np.clip(1.594 + 0.805 * _pesq - 0.512 * llr_mean - 0.007 * wss_dist, 1, 5))
+    csig = float(np.clip(3.093 - 1.029 * llr_mean + 0.603 * pesq - 0.009 * wss_dist, 1, 5))
+    cbak = float(np.clip(1.634 + 0.478 * pesq - 0.007 * wss_dist + 0.063 * seg_snr, 1, 5))
+    covl = float(np.clip(1.594 + 0.805 * pesq - 0.512 * llr_mean - 0.007 * wss_dist, 1, 5))
 
-    return SpeechMetrics(_pesq, csig, cbak, covl)
+    return SpeechMetrics(csig, cbak, covl)
 
 
 def _calc_weighted_spectral_slope_measure(
-    hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int = SAMPLE_RATE
+    hypothesis: np.ndarray,
+    reference: np.ndarray,
+    sample_rate: int = SAMPLE_RATE,
 ) -> np.ndarray:
     from scipy.fftpack import fft
 
@@ -250,7 +258,9 @@ def _calc_weighted_spectral_slope_measure(
 
 
 def _calc_log_likelihood_ratio_measure(
-    hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int = SAMPLE_RATE
+    hypothesis: np.ndarray,
+    reference: np.ndarray,
+    sample_rate: int = SAMPLE_RATE,
 ) -> np.ndarray:
     from scipy.linalg import toeplitz
 
@@ -296,7 +306,9 @@ def _calc_log_likelihood_ratio_measure(
 
 
 def _calc_snr(
-    hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int = SAMPLE_RATE
+    hypothesis: np.ndarray,
+    reference: np.ndarray,
+    sample_rate: int = SAMPLE_RATE,
 ) -> tuple[float, np.ndarray]:
     # The lengths of the reference and hypothesis must be the same.
     reference_length = len(reference)

sonusai/metrics/class_summary.py

@@ -2,10 +2,10 @@
 import numpy as np
 import pandas as pd
 
-from sonusai.mixture import GeneralizedIDs
-from sonusai.mixture import MixtureDatabase
-from sonusai.mixture import Predict
-from sonusai.mixture import Truth
+from ..datatypes import GeneralizedIDs
+from ..datatypes import Predict
+from ..datatypes import Truth
+from ..mixture.mixdb import MixtureDatabase
 
 
 def class_summary(
@@ -31,7 +31,7 @@ def class_summary(
       macro avg     0.85    0.83    0.84    0.05    0.96            3768
       micro-avgwo
     """
-    from sonusai.metrics import one_hot
+    from ..metrics.one_hot import one_hot
 
     num_classes = truth_f.shape[1]
 
@@ -58,11 +58,11 @@ def class_summary(
     else:
         row_n = [f"Class {i}" for i in range(1, num_classes + 1)]
 
-    df = pd.DataFrame(metrics[:, table_idx], columns=col_n, index=row_n) # pyright: ignore [reportArgumentType]
+    df = pd.DataFrame(metrics[:, table_idx], columns=col_n, index=row_n)  # pyright: ignore [reportArgumentType]
 
     # [miPPV, miTPR, miF1, miFPR, miACC, miAP, miAUC, TPSUM]
     avg_row_n = ["Macro-avg", "Micro-avg", "Weighted-avg"]
-    dfavg = pd.DataFrame(metavg, columns=col_n, index=avg_row_n) # pyright: ignore [reportArgumentType]
+    dfavg = pd.DataFrame(metavg, columns=col_n, index=avg_row_n)  # pyright: ignore [reportArgumentType]
 
     # dfblank = pd.DataFrame([''])
     # pd.concat([df, dfblank, dfblank, dfavg])

sonusai/metrics/confusion_matrix_summary.py

@@ -2,10 +2,10 @@
 import numpy as np
 import pandas as pd
 
-from sonusai.mixture import GeneralizedIDs
-from sonusai.mixture import MixtureDatabase
-from sonusai.mixture import Predict
-from sonusai.mixture import Truth
+from ..datatypes import GeneralizedIDs
+from ..datatypes import Predict
+from ..datatypes import Truth
+from ..mixture.mixdb import MixtureDatabase
 
 
 def confusion_matrix_summary(
@@ -30,7 +30,7 @@ def confusion_matrix_summary(
 
     Returns pandas dataframes of confusion matrix cmdf and normalized confusion matrix cmndf.
     """
-    from sonusai.metrics import one_hot
+    from ..metrics.one_hot import one_hot
 
     num_classes = truth_f.shape[1]
     # TODO: re-work for modern mixdb API

sonusai/metrics/one_hot.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from sonusai.mixture.datatypes import Predict
-from sonusai.mixture.datatypes import Truth
+from ..datatypes import Predict
+from ..datatypes import Truth
 
 
 def one_hot(
@@ -53,8 +53,8 @@ def one_hot(
     from sklearn.metrics import precision_recall_fscore_support
     from sklearn.metrics import roc_auc_score
 
-    from sonusai.utils import get_num_classes_from_predict
-    from sonusai.utils import reshape_outputs
+    from ..utils.reshape import get_num_classes_from_predict
+    from ..utils.reshape import reshape_outputs
 
     if truth.shape != predict.shape:
         raise ValueError("truth and predict are not the same shape")

sonusai/metrics/snr_summary.py

@@ -2,11 +2,11 @@
 import numpy as np
 import pandas as pd
 
-from sonusai.mixture import GeneralizedIDs
-from sonusai.mixture import MixtureDatabase
-from sonusai.mixture import Predict
-from sonusai.mixture import Segsnr
-from sonusai.mixture import Truth
+from ..datatypes import GeneralizedIDs
+from ..datatypes import Predict
+from ..datatypes import Segsnr
+from ..datatypes import Truth
+from ..mixture.mixdb import MixtureDatabase
 
 
 def snr_summary(
@@ -40,8 +40,8 @@ def snr_summary(
     """
     import warnings
 
-    from sonusai.metrics import one_hot
-    from sonusai.queries import get_mixids_from_snr
+    from ..metrics.one_hot import one_hot
+    from ..queries.queries import get_mixids_from_snr
 
     num_classes = truth_f.shape[1]
 

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,31 @@
 # 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 .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