sonusai 1.0.16__cp311-abi3-macosx_10_12_x86_64.whl

sonusai/main.py

@@ -0,0 +1,134 @@
+"""sonusai
+
+usage: sonusai [--version] [--help] <command> [<args>...]
+
+The sonusai commands are:
+    <This information is automatically generated.>
+
+Aaware Sound and Voice Machine Learning Framework. See 'sonusai help <command>'
+for more information on a specific command.
+
+"""
+
+import sys
+from importlib import import_module
+from pkgutil import iter_modules
+
+from docopt import docopt
+
+from sonusai import BASEDIR
+from sonusai import __version__ as sai_version
+from sonusai import commands_list
+from sonusai import logger
+from sonusai.utils.docstring import add_commands_to_docstring
+from sonusai.utils.docstring import trim_docstring
+
+
+def discover_plugins():
+    plugins = {}
+    plugin_docstrings = []
+    for _, name, _ in iter_modules():
+        if name.startswith("sonusai_") and not name.startswith("sonusai_asr_"):
+            module = import_module(name)
+            plugins[name] = {
+                "commands": commands_list(module.commands_doc),
+                "basedir": module.BASEDIR,
+            }
+            plugin_docstrings.append(module.commands_doc)
+    return plugins, plugin_docstrings
+
+
+def execute_command_direct(command: str, argv: list[str], basedir: str) -> None:
+    """Execute a command by importing and running it directly."""
+    try:
+        # Add the command directory to the Python path temporarily
+        if basedir not in sys.path:
+            sys.path.insert(0, basedir)
+
+        # Import the command module
+        command_module = import_module(command)
+
+        # Set up sys.argv as the command module expects it
+        original_argv = sys.argv
+        sys.argv = [command, *argv]
+
+        try:
+            # Execute the main function if it exists
+            if hasattr(command_module, "main"):
+                command_module.main()
+            else:
+                logger.error(f"Command module {command} has no main() function")
+                sys.exit(1)
+        finally:
+            # Restore original sys.argv
+            sys.argv = original_argv
+
+    except ImportError as err:
+        logger.error(f"Failed to import command module {command}: {err}")
+        sys.exit(1)
+    except Exception as err:
+        logger.error(f"Error executing command {command}: {err}")
+        sys.exit(1)
+
+
+def handle_help_command_direct(argv: list[str], base_commands: list[str], plugins: dict) -> None:
+    """Handle the help command by executing modules directly."""
+    if not argv:
+        # Show the main help by re-running with -h
+        sys.argv = ["sonusai", "-h"]
+        main()
+        return
+
+    help_target = argv[0]
+
+    if help_target in base_commands:
+        execute_command_direct(help_target, ["-h"], BASEDIR)
+    else:
+        for data in plugins.values():
+            if help_target in data["commands"]:
+                execute_command_direct(help_target, ["-h"], data["basedir"])
+                return
+
+        logger.error(f"{help_target} is not a SonusAI command. See 'sonusai help'.")
+        sys.exit(1)
+
+
+def main() -> None:
+    plugins, plugin_docstrings = discover_plugins()
+    updated_docstring = add_commands_to_docstring(__doc__, plugin_docstrings)
+    args = docopt(
+        trim_docstring(updated_docstring),
+        version=sai_version,
+        options_first=True,
+    )
+
+    command = args["<command>"]
+    argv = args["<args>"]
+    base_commands = commands_list()
+
+    if command == "help":
+        handle_help_command_direct(argv, base_commands, plugins)
+        return
+
+    if command in base_commands:
+        execute_command_direct(command, argv, BASEDIR)
+        return
+
+    for data in plugins.values():
+        if command in data["commands"]:
+            execute_command_direct(command, argv, data["basedir"])
+            return
+
+    logger.error(f"{command} is not a SonusAI command. See 'sonusai help'.")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    from sonusai import exception_handler
+    from sonusai.utils.keyboard_interrupt import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)

sonusai/metrics/__init__.py

@@ -0,0 +1,43 @@
+# SonusAI metrics utilities for model training and validation
+
+from .calc_audio_stats import calc_audio_stats
+from .calc_class_weights import calc_class_weights_from_mixdb
+from .calc_class_weights import calc_class_weights_from_truth
+from .calc_optimal_thresholds import calc_optimal_thresholds
+from .calc_pcm import calc_pcm
+from .calc_pesq import calc_pesq
+from .calc_phase_distance import calc_phase_distance
+from .calc_sa_sdr import calc_sa_sdr
+from .calc_sample_weights import calc_sample_weights
+from .calc_segsnr_f import calc_segsnr_f
+from .calc_segsnr_f import calc_segsnr_f_bin
+from .calc_speech import calc_speech
+from .calc_wer import calc_wer
+from .calc_wsdr import calc_wsdr
+from .calculate_metrics import calculate_metrics
+from .class_summary import class_summary
+from .confusion_matrix_summary import confusion_matrix_summary
+from .one_hot import one_hot
+from .snr_summary import snr_summary
+
+__all__ = [
+    "calc_audio_stats",
+    "calc_class_weights_from_mixdb",
+    "calc_class_weights_from_truth",
+    "calc_optimal_thresholds",
+    "calc_pcm",
+    "calc_pesq",
+    "calc_phase_distance",
+    "calc_sa_sdr",
+    "calc_sample_weights",
+    "calc_segsnr_f",
+    "calc_segsnr_f_bin",
+    "calc_speech",
+    "calc_wer",
+    "calc_wsdr",
+    "calculate_metrics",
+    "class_summary",
+    "confusion_matrix_summary",
+    "one_hot",
+    "snr_summary",
+]

sonusai/metrics/calc_audio_stats.py

@@ -0,0 +1,42 @@
+from ..datatypes import AudioStatsMetrics
+from ..datatypes import AudioT
+
+
+def _convert_str_with_factors_to_int(x: str) -> int:
+    if "k" in x:
+        return int(1000 * float(x.replace("k", "")))
+    if "M" in x:
+        return int(1000000 * float(x.replace("M", "")))
+    return int(x)
+
+
+def calc_audio_stats(audio: AudioT, win_len: float | None = None) -> AudioStatsMetrics:
+    from ..mixture.sox_effects import sox_stats
+
+    out = sox_stats(audio, win_len)
+
+    if out is None:
+        raise SystemError("Call to sox failed")
+
+    stats = {}
+    lines = out.split("\n")
+    for line in lines:
+        split_line = line.split()
+        if len(split_line) == 0:
+            continue
+        value = split_line[-1]
+        key = " ".join(split_line[:-1])
+        stats[key] = value
+
+    return AudioStatsMetrics(
+        dco=float(stats["DC offset"]),
+        min=float(stats["Min level"]),
+        max=float(stats["Max level"]),
+        pkdb=float(stats["Pk lev dB"]),
+        lrms=float(stats["RMS lev dB"]),
+        pkr=float(stats["RMS Pk dB"]),
+        tr=float(stats["RMS Tr dB"]),
+        cr=float(stats["Crest factor"]),
+        fl=float(stats["Flat factor"]),
+        pkc=_convert_str_with_factors_to_int(stats["Pk count"]),
+    )

sonusai/metrics/calc_class_weights.py

@@ -0,0 +1,90 @@
+import numpy as np
+
+from ..datatypes import GeneralizedIDs
+from ..datatypes import Truth
+from ..mixture.mixdb import MixtureDatabase
+
+
+def calc_class_weights_from_truth(truth: Truth, other_weight: float | None = None, other_index: int = -1) -> np.ndarray:
+    """Calculate class weights.
+
+    Supports non-existent classes (a problem with sklearn) where non-existent
+    classes get a weight of 0 (instead of inf).
+    Includes optional weighting of an "other" class if specified.
+
+    Reference:
+        weights = class_weight.compute_class_weight(class_weight='balanced', classes=clabels, y=labels)
+
+    Arguments:
+        truth: Truth data in one-hot format. Size can be:
+          - [frames, timesteps, num_classes]
+          - [frames, num_classes]
+        other_weight: float or `None`. Weight of the "other" class.
+            > 1 = increase weighting/importance relative to the true count
+            0 > `other_weight` < 1 = decrease weighting/importance relative
+            < 0 or `None` = disable, use true count (default = `None`)
+        other_index: int. Index of the "other" class in one-hot mode. Defaults to -1 (the last).
+
+    Returns:
+        A numpy array containing class weights.
+    """
+    frames, num_classes = truth.shape
+
+    if num_classes > 1:
+        labels = np.argmax(truth, axis=-1)  # [frames, 1 labels] from one-hot, last dim
+        count = np.bincount(labels, minlength=num_classes).astype(float)
+    else:
+        num_classes = 2
+        labels = np.array(truth >= 0.5).astype(np.int8)[:, 0]  # quantize to binary and shape (frames,) for bincount
+        count = np.bincount(labels, minlength=num_classes).astype(float)
+
+    if other_weight is not None and other_weight > 0:
+        count[other_index] = count[other_index] / np.float32(other_weight)
+
+    weights = np.empty((len(count)), dtype=np.float32)
+    for n in range(len(weights)):
+        if count[n] == 0:
+            # Avoid sklearn problem with absent classes and assign non-existent classes a weight of 0.
+            weights[n] = 0
+        else:
+            weights[n] = frames / (num_classes * count[n])
+
+    return weights
+
+
+def calc_class_weights_from_mixdb(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    other_weight: float = 1,
+    other_index: int = -1,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Calculate class weights using estimated feature counts from a mixture database.
+
+    Arguments:
+        mixdb: Mixture database.
+        mixids: Mixture ID's.
+        other_weight: float or `None`. Weight of the "other" class.
+            > 1 = increase weighting/importance relative to the true count
+            0 > `other_weight` < 1 = decrease weighting/importance relative
+            < 0 or `None` = disable, use true count
+        other_index: int. Index of the "other" class in one-hot mode. Defaults to -1 (the last).
+
+    Returns:
+        count: Count of features in each class.
+        weights: Class weights. [num_classes, 1]
+            Note: for Keras use dict(enumerate(weights))
+    """
+    from ..mixture.class_count import get_class_count_from_mixids
+
+    count = np.ceil(np.array(get_class_count_from_mixids(mixdb=mixdb, mixids=mixids)) / mixdb.feature_step_samples)
+    total_features = sum(count)
+
+    weights = np.empty(mixdb.num_classes, dtype=np.float32)
+    for n in range(len(weights)):
+        if count[n] == 0:
+            # Avoid sklearn problem with absent classes and assign non-existent classes a weight of 0.
+            weights[n] = 0
+        else:
+            weights[n] = total_features / (mixdb.num_classes * count[n])
+
+    return count, weights

sonusai/metrics/calc_optimal_thresholds.py

@@ -0,0 +1,73 @@
+import numpy as np
+
+from ..datatypes import Predict
+from ..datatypes import Truth
+
+
+def calc_optimal_thresholds(
+    truth: Truth,
+    predict: Predict,
+    timesteps: int = 0,
+    truth_thr: float = 0.5,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Calculates optimal thresholds for each class from one-hot prediction and truth data where both are
+    one-hot probabilities (or quantized decisions) with size [frames, num_classes] or [frames, timesteps, num_classes].
+
+    Returns:
+      thresholds_opt_pr     [num_classes, 1] optimal thresholds for PR-curve (F1) performance
+      thresholds_opt_roc    [num_classes, 1] optimal thresholds for ROC-curve (TPR/FPR) performance
+      AP                    [num_classes, 1]
+      AUC                   [num_classes, 1]
+
+    Optional truth_thr is the decision threshold(s) applied to truth one-hot input allowing truth to optionally be
+    continuous probabilities. Default is 0.5.
+    """
+    from sklearn.metrics import average_precision_score
+    from sklearn.metrics import precision_recall_curve
+    from sklearn.metrics import roc_auc_score
+    from sklearn.metrics import roc_curve
+
+    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")
+
+    predict, truth = reshape_outputs(predict=predict, truth=truth, timesteps=timesteps)  # type: ignore[assignment]
+    num_classes = get_num_classes_from_predict(predict=predict, timesteps=timesteps)
+
+    # Apply decision to truth input
+    truth_binary = np.array(truth >= truth_thr).astype(np.int8)
+
+    AP = np.zeros((num_classes, 1))
+    AUC = np.zeros((num_classes, 1))
+    thresholds_opt_pr = np.zeros((num_classes, 1))
+    thresholds_opt_roc = np.zeros((num_classes, 1))
+    eps = np.finfo(float).eps
+    for nci in range(num_classes):
+        # Average Precision also called area under the PR curve AUCPR and
+        # AUC ROC curve using binary-ized truth and continuous prediction probabilities
+        # sklearn returns nan if no active truth in a class but w/un-suppressible div-by-zero warning
+        if sum(truth_binary[:, nci]) == 0:  # no active truth must be NaN
+            thresholds_opt_pr[nci] = np.NaN
+            thresholds_opt_roc[nci] = np.NaN
+            AUC[nci] = np.NaN
+            AP[nci] = np.NaN
+        else:
+            AP[nci] = average_precision_score(truth_binary[:, nci], predict[:, nci], average=None)  # pyright: ignore [reportArgumentType]
+            AUC[nci] = roc_auc_score(truth_binary[:, nci], predict[:, nci], average=None)  # pyright: ignore [reportArgumentType]
+
+            # Optimal threshold from PR curve, optimizes f-score
+            precision, recall, thrpr = precision_recall_curve(truth_binary[:, nci], predict[:, nci])
+            fscore = (2 * precision * recall) / (precision + recall + eps)
+            ix = np.argmax(fscore)  # index of largest f1 score
+            thresholds_opt_pr[nci] = thrpr[ix]
+
+            # Optimal threshold from ROC curve, optimizes J-statistic (TPR-FPR) or gmean
+            fpr, tpr, thrroc = roc_curve(truth_binary[:, nci], predict[:, nci])
+            # J = tpr - fpr                  # J can result in thr > 1
+            gmeans = np.sqrt(tpr * (1 - fpr))  # gmean seems better behaved
+            ix = np.argmax(gmeans)
+            thresholds_opt_roc[nci] = thrroc[ix]
+
+    return thresholds_opt_pr, thresholds_opt_roc, AP, AUC

sonusai/metrics/calc_pcm.py

@@ -0,0 +1,45 @@
+import numpy as np
+
+
+def calc_pcm(
+    hypothesis: np.ndarray, reference: np.ndarray, with_log: bool = False
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Calculate phase constrained magnitude error
+
+    These must include a noise to make a complete mixture estimate, i.e.,
+        noise_est = mixture - sum-over-nsrc(s_est(:, nsrc, :))
+    should be one of the sources in s_true and s_est.
+
+    Calculates mean-over-srcs(mean-over-tf(| (|Sr(t, f)| + |Si(t, f)|) - (|Shr(t, f)| + |Shi(t, f)|) |))
+
+    Reference:
+        Self-attending RNN for Speech Enhancement to Improve Cross-corpus Generalization
+        Ashutosh Pandey, Student Member, IEEE and DeLiang Wang, Fellow, IEEE
+        https://doi.org/10.48550/arXiv.2105.12831
+
+    :param hypothesis: complex [frames, nsrc, bins]
+    :param reference: complex [frames, nsrc, bins]
+    :param with_log: enable log
+    :return: (error, error per bin, error per frame)
+    """
+    # LSM = 1/(T*F) * sumtf(| (|Sr(t, f)| + |Si(t, f)|) - (|Shr(t, f)| + |Shi(t, f)|) |)
+    # LPCM = 1/2 * LSM(s, sh) + 1/2 * LSM(n, nh)
+
+    # [frames, nsrc, bins]
+    hypothesis_abs = np.abs(np.real(hypothesis)) + np.abs(np.imag(hypothesis))
+    reference_abs = np.abs(np.real(reference)) + np.abs(np.imag(reference))
+    err = np.abs(reference_abs - hypothesis_abs)
+
+    # mean over frames, nsrc for value per bin
+    err_b = np.mean(np.mean(err, axis=0), axis=0)
+    # mean over bins, nsrc for value per frame
+    err_f = np.mean(np.mean(err, axis=2), axis=1)
+    # mean over bins and frames, nsrc for scalar value
+    err = np.mean(np.mean(err, axis=(0, 2)), axis=0)
+
+    if with_log:
+        err_b = np.around(20 * np.log10(err_b + np.finfo(np.float32).eps), 3)
+        err_f = np.around(20 * np.log10(err_f + np.finfo(np.float32).eps), 3)
+        err = np.around(20 * np.log10(err + np.finfo(np.float32).eps), 3)
+
+    return err, err_b, err_f

sonusai/metrics/calc_pesq.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+from ..constants import SAMPLE_RATE
+
+
+def calc_pesq(
+    hypothesis: np.ndarray,
+    reference: np.ndarray,
+    error_value: float = 0.0,
+    sample_rate: int = SAMPLE_RATE,
+) -> float:
+    """Computes the PESQ score of hypothesis vs. reference
+
+    Upon error, assigns a value of 0, or user specified value in error_value
+
+    :param hypothesis: estimated audio
+    :param reference: reference audio
+    :param error_value: value to use if error occurs
+    :param sample_rate: sample rate of audio
+    :return: value between -0.5 to 4.5
+    """
+    import warnings
+
+    from pesq import pesq
+
+    from .. import logger
+
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            score = pesq(fs=sample_rate, ref=reference, deg=hypothesis, mode="wb")
+    except Exception as e:
+        logger.debug(f"PESQ error {e}")
+        score = error_value
+
+    return score

sonusai/metrics/calc_phase_distance.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+
+def calc_phase_distance(
+    reference: np.ndarray, hypothesis: np.ndarray, eps: float = 1e-9
+) -> tuple[float, np.ndarray, np.ndarray]:
+    """Calculate weighted phase distance error (weight normalization over bins per frame)
+
+    :param reference: complex [frames, bins]
+    :param hypothesis: complex [frames, bins]
+    :param eps: epsilon value
+    :return: mean, mean per bin, mean per frame
+    """
+    ang_diff = np.angle(reference) - np.angle(hypothesis)
+    phd_mod = (ang_diff + np.pi) % (2 * np.pi) - np.pi
+    rh_angle_diff = phd_mod * 180 / np.pi  # angle diff in deg
+
+    # Use complex divide to intrinsically keep angle diff +/-180 deg, but avoid div by zero (real hyp)
+    # hyp_real = np.real(hypothesis)
+    # near_zeros = np.real(hyp_real) < eps
+    # hyp_real = hyp_real * (np.logical_not(near_zeros))
+    # hyp_real = hyp_real + (near_zeros * eps)
+    # hypothesis = hyp_real + 1j*np.imag(hypothesis)
+    # rh_angle_diff = np.angle(reference / hypothesis) * 180 / np.pi  # angle diff +/-180
+
+    # weighted mean over all (scalar)
+    reference_mag = np.abs(reference)
+    ref_weight = reference_mag / (np.sum(reference_mag) + eps)  # frames x bins
+    err = float(np.around(np.sum(ref_weight * rh_angle_diff), 3))
+
+    # weighted mean over frames (value per bin)
+    err_b = np.zeros(reference.shape[1])
+    for bi in range(reference.shape[1]):
+        ref_weight = reference_mag[:, bi] / (np.sum(reference_mag[:, bi], axis=0) + eps)
+        err_b[bi] = np.around(np.sum(ref_weight * rh_angle_diff[:, bi]), 3)
+
+    # weighted mean over bins (value per frame)
+    err_f = np.zeros(reference.shape[0])
+    for fi in range(reference.shape[0]):
+        ref_weight = reference_mag[fi, :] / (np.sum(reference_mag[fi, :]) + eps)
+        err_f[fi] = np.around(np.sum(ref_weight * rh_angle_diff[fi, :]), 3)
+
+    return err, err_b, err_f

sonusai/metrics/calc_sa_sdr.py

@@ -0,0 +1,64 @@
+import numpy as np
+
+
+def calc_sa_sdr(
+    hypothesis: np.ndarray,
+    reference: np.ndarray,
+    with_scale: bool = False,
+    with_negate: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Calculate source-aggregated SDR (signal distortion ratio) using all source inputs which are [samples, nsrc].
+
+    These should include a noise to be a complete mixture estimate, i.e.,
+        noise_est = sum-over-all-srcs(s_est(0:nsamples, :) - sum-over-non-noisesrc(s_est(0:nsamples, n))
+    should be one of the sources in reference (s_true) and hypothesis (s_est).
+
+    Calculates -10*log10(sumn(||sn||^2) / sumn(||sn - shn||^2)
+    Note: for SA method, sums are done independently on ref and error before division, vs. SDR and SI-SDR
+    where sum over n is taken after divide (before log).  This is more stable in noise-only cases and also
+    when some sources are poorly estimated.
+    TBD: add soft-max option with eps and tau params
+
+    Reference:
+        SA-SDR: A Novel Loss Function for Separation of Meeting Style Data
+        Thilo von Neumann, Keisuke Kinoshita, Christoph Boeddeker, Marc Delcroix, Reinhold Haeb-Umbach
+        https://doi.org/10.48550/arXiv.2110.15581
+
+    :param hypothesis: [samples, nsrc]
+    :param reference: [samples, nsrc]
+    :param with_scale: enable scaling (scaling is same as in SI-SDR)
+    :param with_negate: enable negation (for use as a loss function)
+    :return: (sa_sdr, opt_scale)
+    """
+    if with_scale:
+        # calc 1 x nsrc scaling factors
+        ref_energy = np.sum(reference**2, axis=0, keepdims=True)
+        # if ref_energy is zero, just set scaling to 1.0
+        with np.errstate(divide="ignore", invalid="ignore"):
+            opt_scale = np.sum(reference * hypothesis, axis=0, keepdims=True) / ref_energy
+            opt_scale[opt_scale == np.inf] = 1.0
+            opt_scale = np.nan_to_num(opt_scale, nan=1.0)
+        scaled_ref = opt_scale * reference
+    else:
+        scaled_ref = reference
+        opt_scale = np.ones((1, reference.shape[1]), dtype=float)
+
+    # multisrc sa-sdr, inputs must be [samples, nsrc]
+    err = scaled_ref - hypothesis
+
+    # -10*log10(sumk(||sk||^2) / sumk(||sk - shk||^2)
+    # sum over samples and sources
+    num = np.sum(reference**2)
+    den = np.sum(err**2)
+    if num == 0 and den == 0:
+        ratio = np.inf
+    else:
+        ratio = num / (den + np.finfo(np.float32).eps)
+
+    sa_sdr = 10 * np.log10(ratio)
+
+    if with_negate:
+        # for use as a loss function
+        sa_sdr = -sa_sdr
+
+    return sa_sdr, opt_scale

sonusai/metrics/calc_sample_weights.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+
+def calc_sample_weights(class_weights: np.ndarray, truth: np.ndarray) -> np.ndarray:
+    """Calculate sample weights from class weights and a given truth with 2D or 3D shape.
+
+    Supports one-hot encoded multi-class or binary truth/labels
+    Note returns sum of weighted truth over classes, thus should also work for multi-label ? TBD
+
+    Inputs:
+      class_weights [num_classes, 1]        weights for each class
+      truth         [frames, timesteps, num_classes] or [frames, num_classes]
+
+    Returns:
+      sample_weights [frames, timesteps, 1] or [frames, 1]
+    """
+    ts = truth.shape
+    cs = class_weights.shape
+
+    if ts[-1] == 1 and cs[0] == 2:
+        # Binary truth needs 2nd "none" truth dimension
+        truth = np.concatenate((truth, 1 - truth), axis=1)
+
+    # broadcast [num_classes, 1] over [frames, num_classes] or [frames, timesteps, num_classes]
+    return np.sum(class_weights * truth, axis=-1)

sonusai/metrics/calc_segsnr_f.py

@@ -0,0 +1,82 @@
+import numpy as np
+
+from ..datatypes import AudioF
+from ..datatypes import Segsnr
+from ..datatypes import SnrFBinMetrics
+from ..datatypes import SnrFMetrics
+
+
+def calc_segsnr_f(segsnr_f: Segsnr) -> SnrFMetrics:
+    """Calculate metrics of snr_f truth data.
+
+    Includes mean and standard deviation of the linear values (usually energy)
+    and mean and standard deviation of the dB values (10 * log10).
+    """
+    if np.count_nonzero(segsnr_f) == 0:
+        # If all entries are zeros
+        return SnrFMetrics(0, 0, -np.inf, 0)
+
+    tmp = np.ma.array(segsnr_f, mask=np.logical_not(np.isfinite(segsnr_f)))
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are infinite
+        return SnrFMetrics(np.inf, 0, np.inf, 0)
+
+    snr_mean = np.mean(tmp, axis=0)
+    snr_std = np.std(tmp, axis=0)
+
+    tmp = 10 * np.ma.log10(tmp)
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are masked, special case where all inputs are either 0 or infinite
+        snr_db_mean = -np.inf
+        snr_db_std = np.inf
+    else:
+        snr_db_mean = np.mean(tmp, axis=0)
+        snr_db_std = np.std(tmp, axis=0)
+
+    return SnrFMetrics(snr_mean, snr_std, snr_db_mean, snr_db_std)
+
+
+def calc_segsnr_f_bin(target_f: AudioF, noise_f: AudioF) -> SnrFBinMetrics:
+    """Calculate per-bin segmental SNR metrics.
+
+    Includes per-bin mean and standard deviation of the linear values
+    and mean and standard deviation of the dB values.
+    """
+    if target_f.ndim != 2 and noise_f.ndim != 2:
+        raise ValueError("target_f and noise_f must have 2 dimensions")
+
+    segsnr_f = (np.abs(target_f) ** 2) / (np.abs(noise_f) ** 2 + np.finfo(np.float32).eps)
+
+    frames, bins = segsnr_f.shape
+    if np.count_nonzero(segsnr_f) == 0:
+        # If all entries are zeros
+        return SnrFBinMetrics(np.zeros(bins), np.zeros(bins), -np.inf * np.ones(bins), np.zeros(bins))
+
+    tmp = np.ma.array(segsnr_f, mask=np.logical_not(np.isfinite(segsnr_f)))
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are infinite
+        return SnrFBinMetrics(
+            np.inf * np.ones(bins),
+            np.zeros(bins),
+            np.inf * np.ones(bins),
+            np.zeros(bins),
+        )
+
+    snr_mean = np.mean(tmp, axis=0)
+    snr_std = np.std(tmp, axis=0)
+
+    tmp = 10 * np.ma.log10(tmp)
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are masked, special case where all inputs are either 0 or infinite
+        snr_db_mean = -np.inf * np.ones(bins)
+        snr_db_std = np.inf * np.ones(bins)
+    else:
+        snr_db_mean = np.mean(tmp, axis=0)
+        snr_db_std = np.std(tmp, axis=0)
+
+    return SnrFBinMetrics(
+        np.ma.getdata(snr_mean),
+        np.ma.getdata(snr_std),
+        np.ma.getdata(snr_db_mean),
+        np.ma.getdata(snr_db_std),
+    )