sonusai 0.19.5__py3-none-any.whl → 0.19.8__py3-none-any.whl

sonusai/genmixdb.py

@@ -1,15 +1,16 @@
 """sonusai genmixdb
 
-usage: genmixdb [-hvmfsdj] LOC
+usage: genmixdb [-hvmfsdjn] LOC
 
 options:
-   -h, --help
-   -v, --verbose    Be verbose.
-   -m, --mix        Save mixture data. [default: False].
-   -f, --ft         Save feature/truth_f data. [default: False].
-   -s, --segsnr     Save segsnr data. [default: False].
-   -d, --dryrun     Perform a dry run showing the processed config. [default: False].
-   -j, --json       Save JSON version of database. [default: False].
+    -h, --help
+    -v, --verbose   Be verbose.
+    -m, --mix       ave mixture data. [default: False].
+    -f, --ft        Save feature/truth_f data. [default: False].
+    -s, --segsnr    Save segsnr data. [default: False].
+    -d, --dryrun    Perform a dry run showing the processed config. [default: False].
+    -j, --json      Save JSON version of database. [default: False].
+    -n, --nopar     Do not run in parallel. [default: False].
 
 Create mixture database data for training and evaluation. Optionally, also create mixture audio and feature/truth data.
 
@@ -114,10 +115,6 @@ will find all .wav files in the specified directories and process them as target
 """
 
 import signal
-from dataclasses import dataclass
-
-from sonusai.mixture import Mixture
-from sonusai.mixture import MixtureDatabase
 
 
 def signal_handler(_sig, _frame):
@@ -132,17 +129,6 @@ def signal_handler(_sig, _frame):
 signal.signal(signal.SIGINT, signal_handler)
 
 
-@dataclass
-class MPGlobal:
-    mixdb: MixtureDatabase
-    save_mix: bool
-    save_ft: bool
-    save_segsnr: bool
-
-
-MP_GLOBAL: MPGlobal
-
-
 def genmixdb(
     location: str,
     save_mix: bool = False,
@@ -152,7 +138,9 @@ def genmixdb(
     show_progress: bool = False,
     test: bool = False,
     save_json: bool = False,
-) -> MixtureDatabase:
+    no_par: bool = False,
+) -> None:
+    from functools import partial
     from random import seed
 
     import yaml
@@ -163,7 +151,6 @@ def genmixdb(
     from sonusai.mixture import AugmentationRule
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import balance_targets
-    from sonusai.mixture import generate_mixtures
     from sonusai.mixture import get_all_snrs_from_config
     from sonusai.mixture import get_augmentation_rules
     from sonusai.mixture import get_augmented_targets
@@ -329,7 +316,8 @@ def genmixdb(
             f"{seconds_to_hms(seconds=noise_audio_duration)}"
         )
 
-    used_noise_files, used_noise_samples, mixtures = generate_mixtures(
+    used_noise_files, used_noise_samples = populate_mixture_table(
+        location=location,
         noise_mix_mode=mixdb.noise_mix_mode,
         augmented_targets=augmented_targets,
         target_files=target_files,
@@ -342,16 +330,17 @@ def genmixdb(
         num_classes=mixdb.num_classes,
         feature_step_samples=mixdb.feature_step_samples,
         num_ir=mixdb.num_impulse_response_files,
+        test=test,
     )
 
-    num_mixtures = len(mixtures)
+    num_mixtures = len(mixdb.mixtures)
     update_mixid_width(location, num_mixtures, test)
 
     if logging:
         logger.info("")
         logger.info(f"Found {num_mixtures:,} mixtures to process")
 
-    total_duration = float(sum([mixture.samples for mixture in mixtures])) / SAMPLE_RATE
+    total_duration = float(sum([mixture.samples for mixture in mixdb.mixtures])) / SAMPLE_RATE
 
     if logging:
         log_duration_and_sizes(
@@ -375,17 +364,21 @@ def genmixdb(
     if logging:
         logger.info("Generating mixtures")
     progress = track(total=num_mixtures, disable=not show_progress)
-    mixtures = par_track(
-        _process_mixture,
-        mixtures,
+    par_track(
+        partial(
+            _process_mixture,
+            location=location,
+            save_mix=save_mix,
+            save_ft=save_ft,
+            save_segsnr=save_segsnr,
+            test=test,
+        ),
+        range(num_mixtures),
         progress=progress,
-        initializer=_initializer,
-        initargs=(location, save_mix, save_ft, save_segsnr, test),
+        no_par=no_par,
     )
     progress.close()
 
-    populate_mixture_table(location, mixtures, test)
-
     total_noise_files = len(noise_files)
 
     total_samples = mixdb.total_samples()
@@ -414,79 +407,62 @@ def genmixdb(
         mixdb = MixtureDatabase(location)
         mixdb.save()
 
-    return mixdb
-
 
-def _initializer(location: str, save_mix: bool, save_ft: bool, save_segsnr: bool, test: bool) -> None:
-    global MP_GLOBAL
-
-    MP_GLOBAL = MPGlobal(
-        mixdb=MixtureDatabase(location, test),
-        save_mix=save_mix,
-        save_ft=save_ft,
-        save_segsnr=save_segsnr,
-    )
-
-
-def _process_mixture(mixture: Mixture) -> Mixture:
-    from typing import Any
+def _process_mixture(
+    m_id: int,
+    location: str,
+    save_mix: bool,
+    save_ft: bool,
+    save_segsnr: bool,
+    test: bool,
+) -> None:
+    from functools import partial
 
-    from sonusai.mixture import get_ft
-    from sonusai.mixture import get_segsnr
-    from sonusai.mixture import get_truth
-    from sonusai.mixture import update_mixture
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture import clear_cached_data
+    from sonusai.mixture import update_mixture_table
     from sonusai.mixture import write_cached_data
     from sonusai.mixture import write_mixture_metadata
 
-    global MP_GLOBAL
+    with_data = save_mix or save_ft or save_segsnr
+
+    genmix_data = update_mixture_table(location, m_id, with_data, test)
 
-    with_data = MP_GLOBAL.save_mix or MP_GLOBAL.save_ft
-    mixdb = MP_GLOBAL.mixdb
+    mixdb = MixtureDatabase(location, test)
+    mixture = mixdb.mixture(m_id)
 
-    mixture, genmix_data = update_mixture(mixdb, mixture, with_data)
+    write = partial(write_cached_data, location=location, name="mixture", index=mixture.name)
+    clear = partial(clear_cached_data, location=location, name="mixture", index=mixture.name)
 
     if with_data:
-        write_data: list[tuple[str, Any]] = []
-
-        if MP_GLOBAL.save_mix:
-            write_data.append(("targets", genmix_data.targets))
-            write_data.append(("noise", genmix_data.noise))
-            write_data.append(("mixture", genmix_data.mixture))
-
-        if MP_GLOBAL.save_ft:
-            if genmix_data.targets is None or genmix_data.noise is None or genmix_data.mixture is None:
-                raise RuntimeError("Mixture data was not generated properly")
-            truth_t = get_truth(
-                mixdb=mixdb,
-                mixture=mixture,
-                targets_audio=genmix_data.targets,
-                noise_audio=genmix_data.noise,
-                mixture_audio=genmix_data.mixture,
-            )
-            feature, truth_f = get_ft(
-                mixdb=mixdb,
-                mixture=mixture,
-                mixture_audio=genmix_data.mixture,
-                truth_t=truth_t,
+        write(
+            items=[
+                ("targets", genmix_data.targets),
+                ("target", genmix_data.target),
+                ("noise", genmix_data.noise),
+                ("mixture", genmix_data.mixture),
+            ]
+        )
+
+        if save_ft:
+            clear(items=["feature", "truth_f"])
+            feature, truth_f = mixdb.mixture_ft(m_id)
+            write(
+                items=[
+                    ("feature", feature),
+                    ("truth_f", truth_f),
+                ]
             )
-            write_data.append(("feature", feature))
-            write_data.append(("truth_f", truth_f))
 
-            if MP_GLOBAL.save_segsnr:
-                if genmix_data.target is None:
-                    raise RuntimeError("Target data was not generated properly")
-                segsnr = get_segsnr(
-                    mixdb=mixdb,
-                    mixture=mixture,
-                    target_audio=genmix_data.target,
-                    noise=genmix_data.noise,
-                )
-                write_data.append(("segsnr", segsnr))
+        if save_segsnr:
+            clear(items=["segsnr"])
+            segsnr = mixdb.mixture_segsnr(m_id)
+            write(items=[("segsnr", segsnr)])
 
-        write_cached_data(mixdb.location, "mixture", mixture.name, write_data)
-        write_mixture_metadata(mixdb, mixture)
+        if not save_mix:
+            clear(items=["targets", "target", "noise", "mixture"])
 
-    return mixture
+        write_mixture_metadata(mixdb, m_id)
 
 
 def main() -> None:
@@ -519,6 +495,7 @@ def main() -> None:
     save_segsnr = args["--segsnr"]
     dryrun = args["--dryrun"]
     save_json = args["--json"]
+    no_par = args["--nopar"]
     location = args["LOC"]
 
     start_time = time.monotonic()
@@ -549,6 +526,7 @@ def main() -> None:
             save_segsnr=save_segsnr,
             show_progress=True,
             save_json=save_json,
+            no_par=no_par,
         )
     except Exception as e:
         logger.debug(e)

sonusai/metrics/calc_class_weights.py

@@ -54,7 +54,7 @@ def calc_class_weights_from_truth(truth: Truth, other_weight: float | None = Non
 
 def calc_class_weights_from_mixdb(
     mixdb: MixtureDatabase,
-    mixids: GeneralizedIDs | None = None,
+    mixids: GeneralizedIDs = "*",
     other_weight: float = 1,
     other_index: int = -1,
 ) -> tuple[np.ndarray, np.ndarray]:
@@ -77,8 +77,6 @@ def calc_class_weights_from_mixdb(
     from sonusai.mixture import get_class_count_from_mixids
 
     count = np.ceil(np.array(get_class_count_from_mixids(mixdb=mixdb, mixids=mixids)) / mixdb.feature_step_samples)
-    if mixdb.truth_mutex and other_weight is not None and other_weight > 0:
-        count[other_index] = count[other_index] / np.float32(other_weight)
     total_features = sum(count)
 
     weights = np.empty(mixdb.num_classes, dtype=np.float32)

sonusai/metrics/calc_optimal_thresholds.py

@@ -51,8 +51,8 @@ def calc_optimal_thresholds(
             AUC[nci] = np.NaN
             AP[nci] = np.NaN
         else:
-            AP[nci] = average_precision_score(truth_binary[:, nci], predict[:, nci], average=None)
-            AUC[nci] = roc_auc_score(truth_binary[:, nci], predict[:, nci], average=None)
+            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])

sonusai/metrics/calc_phase_distance.py

@@ -26,7 +26,7 @@ def calc_phase_distance(
     # weighted mean over all (scalar)
     reference_mag = np.abs(reference)
     ref_weight = reference_mag / (np.sum(reference_mag) + eps)  # frames x bins
-    err = np.around(np.sum(ref_weight * rh_angle_diff), 3)
+    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])

sonusai/metrics/calc_segsnr_f.py

@@ -45,7 +45,7 @@ def calc_segsnr_f_bin(target_f: AudioF, noise_f: AudioF) -> SnrFBinMetrics:
     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)
+    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:

sonusai/metrics/calc_speech.py

@@ -32,16 +32,16 @@ def calc_speech(hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int
     llr_mean = np.mean(ll_rs[:llr_len])
 
     # Segmental SNR
-    snr_dist, segsnr_dist = _calc_snr(hypothesis=hypothesis, reference=reference, sample_rate=sample_rate)
+    _, segsnr_dist = _calc_snr(hypothesis=hypothesis, reference=reference, sample_rate=sample_rate)
     seg_snr = np.mean(segsnr_dist)
 
     # PESQ
     _pesq = calc_pesq(hypothesis=hypothesis, reference=reference, sample_rate=sample_rate)
 
     # Now compute the composite measures
-    csig = np.clip(3.093 - 1.029 * llr_mean + 0.603 * _pesq - 0.009 * wss_dist, 1, 5)
-    cbak = np.clip(1.634 + 0.478 * _pesq - 0.007 * wss_dist + 0.063 * seg_snr, 1, 5)
-    covl = 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)
 
@@ -284,8 +284,8 @@ def _calc_log_likelihood_ratio_measure(
         hypothesis_frame = np.multiply(hypothesis_frame, window)
 
         # (2) Get the autocorrelation lags and LPC parameters used to compute the log likelihood ratio measure.
-        r_reference, ref_reference, a_reference = _lp_coefficients(reference_frame, p)
-        r_hypothesis, ref_hypothesis, a_hypothesis = _lp_coefficients(hypothesis_frame, p)
+        r_reference, _, a_reference = _lp_coefficients(reference_frame, p)
+        _, _, a_hypothesis = _lp_coefficients(hypothesis_frame, p)
 
         # (3) Compute the log likelihood ratio measure
         numerator = np.dot(np.matmul(a_hypothesis, toeplitz(r_reference)), a_hypothesis)

sonusai/metrics/class_summary.py

@@ -38,7 +38,7 @@ def class_summary(
     # TODO: re-work for modern mixdb API
     y_truth_f, y_predict = get_mixids_data(mixdb, mixids, truth_f, predict)  # type: ignore[name-defined]
 
-    if not mixdb.truth_mutex and num_classes > 1:
+    if num_classes > 1:
         if not isinstance(predict_thr, np.ndarray):
             if predict_thr == 0:
                 predict_thr = np.atleast_1d(0.5)
@@ -53,25 +53,16 @@ def class_summary(
     # [ACC, TPR, PPV, TNR, FPR, HITFA, F1, MCC, NT, PT, TP, FP, AP, AUC]
     table_idx = np.array([2, 1, 6, 4, 0, 12, 13, 9])
     col_n = ["PPV", "TPR", "F1", "FPR", "ACC", "AP", "AUC", "Support"]
-    if mixdb.truth_mutex:
-        if len(mixdb.class_labels) >= num_classes - 1:  # labels exist with or without Other
-            row_n = mixdb.class_labels
-            if len(mixdb.class_labels) == num_classes - 1:  # Other label does not exist, so add it
-                row_n.append("Other")
-        else:
-            row_n = [f"Class {i}" for i in range(1, num_classes)]
-            row_n.append("Other")
+    if len(mixdb.class_labels) == num_classes:
+        row_n = mixdb.class_labels
     else:
-        if len(mixdb.class_labels) == num_classes:
-            row_n = mixdb.class_labels
-        else:
-            row_n = [f"Class {i}" for i in range(1, num_classes + 1)]
+        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)
+    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)
+    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

@@ -37,7 +37,7 @@ def confusion_matrix_summary(
     ytrue, ypred = get_mixids_data(mixdb=mixdb, mixids=mixids, truth_f=truth_f, predict=predict)  # type: ignore[name-defined]
 
     # Check predict_thr array or scalar and return final scalar predict_thr value
-    if not mixdb.truth_mutex and num_classes > 1:
+    if num_classes > 1:
         if not isinstance(predict_thr, np.ndarray):
             if predict_thr == 0:
                 # multi-label predict_thr scalar 0 force to 0.5 default
@@ -61,31 +61,15 @@ def confusion_matrix_summary(
     else:
         class_names = [f"Class {i}" for i in range(1, num_classes + 1)]
 
-    class_nums = [f"{i}" for i in range(1, num_classes + 1)]
-
-    if mixdb.truth_mutex:
-        # single-label mode force to argmax mode
-        predict_thr = np.array(0, dtype=np.float32)
-        _, _, cm, cmn, _, _ = one_hot(ytrue, ypred, predict_thr, truth_thr, timesteps)
-        row_n = class_names
-        row_n[-1] = "Other"
-        # mux = pd.MultiIndex.from_product([['Single-label/mutex mode, truth thr = {}'.format(truth_thr)],
-        #                                   class_nums])
-        # mux = pd.MultiIndex.from_product([['truth thr = {}'.format(truth_thr)], class_nums])
-
-        cmdf = pd.DataFrame(cm, index=row_n, columns=class_nums, dtype=np.int32)
-        cmndf = pd.DataFrame(cmn, index=row_n, columns=class_nums, dtype=np.float32)
-
-    else:
-        _, _, cm, cmn, _, _ = one_hot(ytrue[:, class_idx], ypred[:, class_idx], predict_thr, truth_thr, timesteps)
-        cname = class_names[class_idx]
-        row_n = ["TrueN", "TrueP"]
-        col_n = ["N-" + cname, "P-" + cname]
-        cmdf = pd.DataFrame(cm, index=row_n, columns=col_n, dtype=np.int32)
-        cmndf = pd.DataFrame(cmn, index=row_n, columns=col_n, dtype=np.float32)
-        # add thresholds in 3rd row
-        pdnote = pd.DataFrame(np.atleast_2d([predict_thr, truth_thr]), index=["p/t thr:"], columns=col_n)
-        cmdf = pd.concat([cmdf, pdnote])
-        cmndf = pd.concat([cmndf, pdnote])
+    _, _, cm, cmn, _, _ = one_hot(ytrue[:, class_idx], ypred[:, class_idx], predict_thr, truth_thr, timesteps)
+    cname = class_names[class_idx]
+    row_n = ["TrueN", "TrueP"]
+    col_n = ["N-" + cname, "P-" + cname]
+    cmdf = pd.DataFrame(cm, index=row_n, columns=col_n, dtype=np.int32)  # pyright: ignore [reportArgumentType]
+    cmndf = pd.DataFrame(cmn, index=row_n, columns=col_n, dtype=np.float32)  # pyright: ignore [reportArgumentType]
+    # add thresholds in 3rd row
+    pdnote = pd.DataFrame(np.atleast_2d([predict_thr, truth_thr]), index=["p/t thr:"], columns=col_n)  # pyright: ignore [reportArgumentType, reportCallIssue]
+    cmdf = pd.concat([cmdf, pdnote])
+    cmndf = pd.concat([cmndf, pdnote])
 
     return cmdf, cmndf

sonusai/metrics/one_hot.py

@@ -185,11 +185,11 @@ def one_hot(
             AP = np.NaN
             # threshold_optpr[nci] = np.NaN
         else:
-            AP = average_precision_score(truthb[:, nci], predict[:, nci], average=None)
+            AP = average_precision_score(truthb[:, nci], predict[:, nci], average=None)  # pyright: ignore [reportArgumentType]
             if len(np.unique(truthb[:, nci])) < 2:  # if active classes not > 1 AUC must be NaN
                 AUC = np.NaN  # i.e. all ones sklearn auc will fail
             else:
-                AUC = roc_auc_score(truthb[:, nci], predict[:, nci], average=None)
+                AUC = roc_auc_score(truthb[:, nci], predict[:, nci], average=None)  # pyright: ignore [reportArgumentType]
             # # Optimal threshold from PR curve, optimizes f-score
             # precision, recall, thresholds = precision_recall_curve(truthb[:, nci], predict[:, nci])
             # fscore = (2 * precision * recall) / (precision + recall)
@@ -263,7 +263,7 @@ def one_hot(
         ]  # specific format, last 3 are unique
 
         # weighted average TBD
-        wp, wr, wf1, _ = precision_recall_fscore_support(truthb, predb, average="weighted", zero_division=0)
+        wp, wr, wf1, _ = precision_recall_fscore_support(truthb, predb, average="weighted", zero_division=0)  # pyright: ignore [reportArgumentType]
         if np.sum(truthb):
             taidx = np.sum(truthb, axis=0) > 0
             wap = average_precision_score(truthb[:, taidx], predict[:, taidx], average="weighted")

sonusai/metrics/snr_summary.py

@@ -48,7 +48,7 @@ def snr_summary(
     snr_mixids = get_mixids_from_snr(mixdb=mixdb, mixids=mixid)
 
     # Check predict_thr array or scalar and return final scalar predict_thr value
-    if not mixdb.truth_mutex and num_classes > 1:
+    if num_classes > 1:
         if not isinstance(predict_thr, np.ndarray):
             if predict_thr == 0:
                 # multi-label predict_thr scalar 0 force to 0.5 default
@@ -84,7 +84,7 @@ def snr_summary(
     for ii, snr in enumerate(snr_mixids):
         # TODO: re-work for modern mixdb API
         y_truth, y_predict = get_mixids_data(mixdb, snr_mixids[snr], truth_f, predict)  # type: ignore[name-defined]
-        _, metrics, _, _, _, mavg = one_hot(y_truth, y_predict, predict_thr, truth_thr, timesteps)
+        _, _, _, _, _, mavg = one_hot(y_truth, y_predict, predict_thr, truth_thr, timesteps)
 
         # mavg macro, micro, weighted: [PPV, TPR, F1, FPR, ACC, mAP, mAUC, TPSUM]
         macro_avg[ii, :] = mavg[0, 0:7]
@@ -104,21 +104,21 @@ def snr_summary(
 
     # SNR format: PPV, TPR, F1, FPR, ACC, AP, AUC
     col_n = ["PPV", "TPR", "F1", "FPR", "ACC", "AP", "AUC"]
-    snr_macrodf = pd.DataFrame(macro_avg, index=list(snr_mixids.keys()), columns=col_n)
+    snr_macrodf = pd.DataFrame(macro_avg, index=list(snr_mixids.keys()), columns=col_n)  # pyright: ignore [reportArgumentType]
     snr_macrodf.sort_index(ascending=False, inplace=True)
 
-    snr_microdf = pd.DataFrame(micro_avg, index=list(snr_mixids.keys()), columns=col_n)
+    snr_microdf = pd.DataFrame(micro_avg, index=list(snr_mixids.keys()), columns=col_n)  # pyright: ignore [reportArgumentType]
     snr_microdf.sort_index(ascending=False, inplace=True)
 
-    snr_wghtdf = pd.DataFrame(wghtd_avg, index=list(snr_mixids.keys()), columns=col_n)
+    snr_wghtdf = pd.DataFrame(wghtd_avg, index=list(snr_mixids.keys()), columns=col_n)  # pyright: ignore [reportArgumentType]
     snr_wghtdf.sort_index(ascending=False, inplace=True)
 
     # Add segmental SNR columns if provided
     if segsnr is not None:
         ssnrdf = pd.DataFrame(
             ssnr_stats,
-            index=list(snr_mixids.keys()),
-            columns=["SSNRavg", "SSNR80p", "SSNRmax"],
+            index=list(snr_mixids.keys()),  # pyright: ignore [reportArgumentType]
+            columns=["SSNRavg", "SSNR80p", "SSNRmax"],  # pyright: ignore [reportArgumentType]
         )
         ssnrdf.sort_index(ascending=False, inplace=True)
         snr_macrodf = pd.concat([snr_macrodf, ssnrdf], axis=1)

sonusai/mixture/__init__.py

@@ -46,19 +46,15 @@ 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 AudiosF
-from .datatypes import AudiosT
 from .datatypes import AudioStatsMetrics
 from .datatypes import AudioT
 from .datatypes import Augmentation
 from .datatypes import AugmentationRule
-from .datatypes import AugmentationRules
-from .datatypes import Augmentations
 from .datatypes import AugmentedTarget
-from .datatypes import AugmentedTargets
 from .datatypes import ClassCount
 from .datatypes import EnergyF
 from .datatypes import EnergyT
@@ -70,34 +66,27 @@ from .datatypes import GenFTData
 from .datatypes import GenMixData
 from .datatypes import ImpulseResponseData
 from .datatypes import ImpulseResponseFile
-from .datatypes import ImpulseResponseFiles
-from .datatypes import ListAudiosT
 from .datatypes import MetricDoc
 from .datatypes import MetricDocs
 from .datatypes import Mixture
 from .datatypes import MixtureDatabaseConfig
-from .datatypes import Mixtures
 from .datatypes import NoiseFile
-from .datatypes import NoiseFiles
 from .datatypes import Predict
 from .datatypes import Segsnr
 from .datatypes import SnrFMetrics
 from .datatypes import SpectralMask
-from .datatypes import SpectralMasks
 from .datatypes import SpeechMetadata
 from .datatypes import SpeechMetrics
 from .datatypes import TargetFile
-from .datatypes import TargetFiles
 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 TruthParameters
 from .datatypes import UniversalSNR
 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
@@ -110,17 +99,14 @@ 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 .generation import update_mixture_table
 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_ft
-from .helpers import get_segsnr
 from .helpers import get_transform_from_audio
-from .helpers import get_truth
 from .helpers import inverse_transform
 from .helpers import mixture_metadata
 from .helpers import write_mixture_metadata

sonusai/mixture/augmentation.py

@@ -1,12 +1,11 @@
 from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import Augmentation
 from sonusai.mixture.datatypes import AugmentationRule
-from sonusai.mixture.datatypes import AugmentationRules
 from sonusai.mixture.datatypes import ImpulseResponseData
 from sonusai.mixture.datatypes import OptionalNumberStr
 
 
-def get_augmentation_rules(rules: list[dict] | dict, num_ir: int = 0) -> AugmentationRules:
+def get_augmentation_rules(rules: list[dict] | dict, num_ir: int = 0) -> list[AugmentationRule]:
     """Generate augmentation rules from list of input rules
 
     :param rules: Dictionary of augmentation config rule[s]
@@ -25,7 +24,7 @@ def get_augmentation_rules(rules: list[dict] | dict, num_ir: int = 0) -> Augment
         rule = _parse_ir(rule, num_ir)
         processed_rules = _expand_rules(expanded_rules=processed_rules, rule=rule)
 
-    return [dataclass_from_dict(AugmentationRule, processed_rule) for processed_rule in processed_rules]
+    return [dataclass_from_dict(AugmentationRule, processed_rule) for processed_rule in processed_rules]  # pyright: ignore [reportReturnType]
 
 
 def _expand_rules(expanded_rules: list[dict], rule: dict) -> list[dict]:
@@ -163,7 +162,7 @@ def estimate_augmented_length_from_length(length: int, tempo: OptionalNumberStr
     return length
 
 
-def get_mixups(augmentations: AugmentationRules) -> list[int]:
+def get_mixups(augmentations: list[AugmentationRule]) -> list[int]:
     """Get a list of mixup values used
 
     :param augmentations: List of augmentations
@@ -172,7 +171,7 @@ def get_mixups(augmentations: AugmentationRules) -> list[int]:
     return sorted({augmentation.mixup for augmentation in augmentations})
 
 
-def get_augmentation_indices_for_mixup(augmentations: AugmentationRules, mixup: int) -> list[int]:
+def get_augmentation_indices_for_mixup(augmentations: list[AugmentationRule], mixup: int) -> list[int]:
     """Get a list of augmentation indices for a given mixup value
 
     :param augmentations: List of augmentations
@@ -327,4 +326,4 @@ def augmentation_from_rule(rule: AugmentationRule, num_ir: int) -> Augmentation:
     if _rule_has_rand(processed_rule):
         processed_rule = _generate_random_rule(processed_rule, num_ir)
 
-    return dataclass_from_dict(Augmentation, processed_rule)
+    return dataclass_from_dict(Augmentation, processed_rule)  # pyright: ignore [reportReturnType]

sonusai/mixture/class_count.py

@@ -3,7 +3,7 @@ from sonusai.mixture.datatypes import GeneralizedIDs
 from sonusai.mixture.mixdb import MixtureDatabase
 
 
-def get_class_count_from_mixids(mixdb: MixtureDatabase, mixids: GeneralizedIDs | None = None) -> ClassCount:
+def get_class_count_from_mixids(mixdb: MixtureDatabase, mixids: GeneralizedIDs = "*") -> ClassCount:
     """Sums the class counts for given mixids"""
     total_class_count = [0] * mixdb.num_classes
     m_ids = mixdb.mixids_to_list(mixids)