sonusai 0.19.6__py3-none-any.whl → 0.19.9__py3-none-any.whl

sonusai/__init__.py

@@ -7,7 +7,7 @@ from rich.traceback import install
 
 install(show_locals=True)
 
-__version__ = metadata.version(__package__)
+__version__ = metadata.version(__package__)  # pyright: ignore [reportArgumentType]
 BASEDIR = dirname(__file__)
 
 commands_doc = """

sonusai/aawscd_probwrite.py

@@ -89,7 +89,7 @@ def on_message(_client, _userdata, message):
         FRAME_COUNT += 1
 
         global PROGRESS
-        PROGRESS.update()
+        PROGRESS.update()  # pyright: ignore [reportOptionalMemberAccess]
 
         if FRAME_COUNT == FRAMES:
             global DONE

sonusai/calc_metric_spenh.py

@@ -582,7 +582,7 @@ def _process_mixture(mixid: int) -> tuple[pd.DataFrame, pd.DataFrame]:
     asr_mx = None
     asr_tge = None
     asr_engines = list(mixdb.asr_configs.keys())
-    if len(asr_engines) > 0 and mixdb.mixture(mixid).snr >= -96:  # noise only, ignore/reset target asr
+    if len(asr_engines) > 0 and not mixdb.mixture(mixid).is_noise_only:  # noise only, ignore/reset target asr
         wer_mx = float(mixdb.mixture_metrics(mixid, [f"mxwer.{asr_engines[0]}"])[0]) * 100
         asr_tt = MP_GLOBAL.mixdb.mixture_speech_metadata(mixid, "text")[0]  # ignore mixup
         if asr_tt is None:

sonusai/genft.py

@@ -1,12 +1,13 @@
 """sonusai genft
 
-usage: genft [-hvs] [-i MIXID] LOC
+usage: genft [-hvsn] [-i MIXID] LOC
 
 options:
     -h, --help
-    -v, --verbose                   Be verbose.
-    -i MIXID, --mixid MIXID         Mixture ID(s) to generate. [default: *].
-    -s, --segsnr                    Save segsnr. [default: False].
+    -v, --verbose               Be verbose.
+    -i MIXID, --mixid MIXID     Mixture ID(s) to generate. [default: *].
+    -s, --segsnr                Save segsnr. [default: False].
+    -n, --nopar                 Do not run in parallel. [default: False].
 
 Generate SonusAI feature/truth data from a SonusAI mixture database.
 
@@ -50,6 +51,7 @@ def genft(
     write: bool = False,
     show_progress: bool = False,
     force: bool = True,
+    no_par: bool = False,
 ) -> list[GenFTData]:
     from functools import partial
 
@@ -62,18 +64,26 @@ def genft(
 
     progress = track(total=len(mixids), disable=not show_progress)
     results = par_track(
-        partial(_genft_kernel, location=location, compute_truth=compute_truth, compute_segsnr=compute_segsnr, force=force, write=write),
+        partial(
+            _genft_kernel,
+            location=location,
+            compute_truth=compute_truth,
+            compute_segsnr=compute_segsnr,
+            force=force,
+            write=write,
+        ),
         mixids,
         progress=progress,
+        no_par=no_par,
     )
     progress.close()
 
     return results
 
 
-def _genft_kernel(m_id: int, location: str, compute_truth: bool, compute_segsnr: bool, force: bool, write: bool) -> GenFTData:
-    from typing import Any
-
+def _genft_kernel(
+    m_id: int, location: str, compute_truth: bool, compute_segsnr: bool, force: bool, write: bool
+) -> GenFTData:
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import write_cached_data
     from sonusai.mixture import write_mixture_metadata
@@ -83,21 +93,23 @@ def _genft_kernel(m_id: int, location: str, compute_truth: bool, compute_segsnr:
     result = GenFTData()
 
     feature, truth_f = mixdb.mixture_ft(m_id=m_id, force=force)
-    write_data: list[tuple[str, Any]] = [("feature", feature)]
     result.feature = feature
+    if write:
+        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("feature", feature)])
 
     if compute_truth:
-        write_data.append(("truth_f", truth_f))
         result.truth_f = truth_f
+        if write:
+            write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("truth_f", truth_f)])
 
     if compute_segsnr:
         segsnr = mixdb.mixture_segsnr(m_id=m_id, force=force)
-        write_data.append(("segsnr", segsnr))
         result.segsnr = segsnr
+        if write:
+            write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("segsnr", segsnr)])
 
     if write:
-        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, write_data)
-        write_mixture_metadata(mixdb, mixdb.mixture(m_id))
+        write_mixture_metadata(mixdb, m_id)
 
     return result
 
@@ -117,6 +129,7 @@ def main() -> None:
     from sonusai import initial_log_messages
     from sonusai import logger
     from sonusai import update_console_handler
+    from sonusai.mixture import SAMPLE_RATE
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import check_audio_files_exist
     from sonusai.utils import human_readable_size
@@ -125,6 +138,7 @@ def main() -> None:
     verbose = args["--verbose"]
     mixids = args["--mixid"]
     compute_segsnr = args["--segsnr"]
+    no_par = args["--nopar"]
     location = args["LOC"]
 
     start_time = time.monotonic()
@@ -138,7 +152,7 @@ def main() -> None:
     mixids = mixdb.mixids_to_list(mixids)
 
     total_samples = mixdb.total_samples(mixids)
-    duration = total_samples / sonusai.mixture.SAMPLE_RATE
+    duration = total_samples / SAMPLE_RATE
     total_transform_frames = total_samples // mixdb.ft_config.overlap
     total_feature_frames = total_samples // mixdb.feature_step_samples
 
@@ -159,6 +173,7 @@ def main() -> None:
             compute_segsnr=compute_segsnr,
             write=True,
             show_progress=True,
+            no_par=no_par,
         )
     except Exception as e:
         logger.debug(e)

sonusai/genmetrics.py

@@ -1,6 +1,6 @@
 """sonusai genmetrics
 
-usage: genmetrics [-hvs] [-i MIXID] [-n INCLUDE] [-x EXCLUDE] LOC
+usage: genmetrics [-hvsd] [-i MIXID] [-n INCLUDE] [-x EXCLUDE] LOC
 
 options:
     -h, --help
@@ -9,27 +9,39 @@ options:
     -n INCLUDE, --include INCLUDE   Metrics to include. [default: all]
     -x EXCLUDE, --exclude EXCLUDE   Metrics to exclude. [default: none]
     -s, --supported                 Show list of supported metrics.
+    -d, --dryrun                    Show list of metrics that will be generated and exit.
 
 Calculate speech enhancement metrics of SonusAI mixture data in LOC.
 
 Inputs:
     LOC         A SonusAI mixture database directory.
     MIXID       A glob of mixture ID(s) to generate.
-    INCLUDE     Comma separated list of metrics to include. Can be 'all' or
-                any of the supported metrics.
-    EXCLUDE     Comma separated list of metrics to exclude. Can be 'none' or
-                any of the supported metrics.
+    INCLUDE     Comma separated list of metrics to include. Can be "all" or
+                any of the supported metrics or glob(s).
+    EXCLUDE     Comma separated list of metrics to exclude. Can be "none" or
+                any of the supported metrics or glob(s)
+
+Note: The default include of "all" excludes the generation of ASR metrics,
+i.e., "*asr*,*wer*". However, if include is manually specified to something other than "all",
+then this behavior is overridden.
+
+Similarly, the default exclude of "none" excludes the generation of ASR metrics,
+i.e., "*asr*,*wer*". However, if exclude is manually specified to something other than "none",
+then this behavior is also overridden.
 
 Examples:
 
 Generate all available mxwer metrics (as determined by mixdb asr_configs parameter):
-> sonusai genmetrics -n"mxwer" mixdb_loc
+> sonusai genmetrics -n"mxwer*" mixdb_loc
 
 Generate only mxwer.faster metrics:
 > sonusai genmetrics -n"mxwer.faster" mixdb_loc
 
-Generate all available metrics except for mxwer.faster:
-> sonusai genmetrics -x"mxwer.faster" mixdb_loc
+Generate only faster metrics:
+> sonusai genmetrics -n"*faster" mixdb_loc
+
+Generate all available metrics except for mxcovl
+> sonusai genmetrics -x"mxcovl" mixdb_loc
 
 """
 
@@ -71,11 +83,13 @@ def main() -> None:
 
     verbose = args["--verbose"]
     mixids = args["--mixid"]
-    includes = [x.strip() for x in args["--include"].lower().split(",")]
-    excludes = [x.strip() for x in args["--exclude"].lower().split(",")]
+    includes = {x.strip() for x in args["--include"].replace(" ", ",").lower().split(",") if x != ""}
+    excludes = {x.strip() for x in args["--exclude"].replace(" ", ",").lower().split(",") if x != ""}
     show_supported = args["--supported"]
+    dryrun = args["--dryrun"]
     location = args["LOC"]
 
+    import fnmatch
     import sys
     import time
     from functools import partial
@@ -89,9 +103,6 @@ def main() -> None:
     from sonusai.utils import seconds_to_hms
     from sonusai.utils import track
 
-    # TODO: Check config.yml for changes to asr_configs and update mixdb
-    # TODO: Support globs for metrics (includes and excludes)
-
     start_time = time.monotonic()
 
     # Setup logging file
@@ -107,38 +118,45 @@ def main() -> None:
         logger.info(f"\nSupported metrics:\n\n{supported.pretty}")
         sys.exit(0)
 
-    if includes is None or "all" in includes:
-        metrics = supported.names
-    else:
-        metrics = set(includes)
-        if "mxwer" in metrics:
-            metrics.remove("mxwer")
-            for name in mixdb.asr_configs:
-                metrics.add(f"mxwer.{name}")
-
-    diff = metrics.difference(supported.names)
-    if diff:
-        logger.error(f"Unrecognized metric: {', '.join(diff)}")
-        sys.exit(1)
+    # Handle default excludes
+    if "none" in excludes:
+        if "all" in includes:
+            excludes = {"*asr*", "*wer*"}
+        else:
+            excludes = set()
 
-    if excludes is None or "none" in excludes:
-        _excludes = set()
-    else:
-        _excludes = set(excludes)
-        if "mxwer" in _excludes:
-            _excludes.remove("mxwer")
-            for name in mixdb.asr_configs:
-                _excludes.add(f"mxwer.{name}")
-
-    diff = _excludes.difference(supported.names)
-    if diff:
-        logger.error(f"Unrecognized metric: {', '.join(diff)}")
-        sys.exit(1)
+    # Handle default includes
+    if "all" in includes:
+        includes = {"*"}
+
+    included_metrics: set[str] = set()
+    for include in includes:
+        for m in fnmatch.filter(supported.names, include):
+            included_metrics.add(m)
 
-    for exclude in _excludes:
-        metrics.discard(exclude)
+    excluded_metrics: set[str] = set()
+    for exclude in excludes:
+        for m in fnmatch.filter(supported.names, exclude):
+            excluded_metrics.add(m)
+
+    requested = included_metrics - excluded_metrics
+
+    # Check for metrics dependencies and cache dependencies even if not explicitly requested.
+    dependencies: set[str] = set()
+    for metric in requested:
+        if metric.startswith("mxwer"):
+            dependencies.add("mxasr." + metric[6:])
+            dependencies.add("tasr." + metric[6:])
+
+    metrics = sorted(requested | dependencies)
+
+    if len(metrics) == 0:
+        logger.warning("No metrics were requested")
+        sys.exit(1)
 
     logger.info(f"Generating metrics: {', '.join(metrics)}")
+    if dryrun:
+        sys.exit(0)
 
     mixids = mixdb.mixids_to_list(mixids)
     logger.info("")
@@ -146,7 +164,7 @@ def main() -> None:
 
     progress = track(total=len(mixids), desc="genmetrics")
     par_track(
-        partial(_process_mixture, location=location, metrics=list(metrics)),
+        partial(_process_mixture, location=location, metrics=metrics),
         mixids,
         progress=progress,
     )

sonusai/genmix.py

@@ -1,14 +1,15 @@
 """sonusai genmix
 
-usage: genmix [-hvgts] [-i MIXID] LOC
+usage: genmix [-hvgtsn] [-i MIXID] LOC
 
 options:
     -h, --help
-    -v, --verbose                   Be verbose.
-    -i MIXID, --mixid MIXID         Mixture ID(s) to generate. [default: *].
-    -g, --target                    Save target. [default: False].
-    -t, --truth                     Save truth_t. [default: False].
-    -s, --segsnr                    Save segsnr_t. [default: False].
+    -v, --verbose               Be verbose.
+    -i MIXID, --mixid MIXID     Mixture ID(s) to generate. [default: *].
+    -g, --target                Save target. [default: False].
+    -t, --truth                 Save truth_t. [default: False].
+    -s, --segsnr                Save segsnr_t. [default: False].
+    -n, --nopar                 Do not run in parallel. [default: False].
 
 Generate SonusAI mixture data from a SonusAI mixture database.
 
@@ -55,6 +56,7 @@ def genmix(
     write: bool = False,
     show_progress: bool = False,
     force: bool = True,
+    no_par: bool = False,
 ) -> list[GenMixData]:
     from functools import partial
 
@@ -77,6 +79,7 @@ def genmix(
         ),
         mixids,
         progress=progress,
+        no_par=no_par,
     )
     progress.close()
 
@@ -98,41 +101,47 @@ def _genmix_kernel(
 
     mixdb = MixtureDatabase(location)
 
+    result = GenMixData()
+
     targets = mixdb.mixture_targets(m_id=m_id, force=force)
+    result.targets = targets
     noise = mixdb.mixture_noise(m_id=m_id, force=force)
-    write_data = [("targets", targets), ("noise", noise)]
+    result.noise = noise
+    if write:
+        write_cached_data(
+            mixdb.location,
+            "mixture",
+            mixdb.mixture(m_id).name,
+            [
+                ("targets", targets),
+                ("noise", noise),
+            ],
+        )
 
     if compute_truth:
         truth_t = mixdb.mixture_truth_t(m_id=m_id, targets=targets, noise=noise, force=force)
-        write_data.append(("truth_t", truth_t))
-    else:
-        truth_t = None
+        result.truth_t = truth_t
+        if write:
+            write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("truth_t", truth_t)])
 
     target = mixdb.mixture_target(m_id=m_id, targets=targets)
-    if save_target:
-        write_data.append(("target", target))
+    result.target = target
+    if save_target and write:
+        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("target", target)])
 
     if compute_segsnr:
         segsnr_t = mixdb.mixture_segsnr_t(m_id=m_id, targets=targets, target=target, noise=noise, force=force)
-        write_data.append(("segsnr_t", segsnr_t))
-    else:
-        segsnr_t = None
+        result.segsnr_t = segsnr_t
+        if write:
+            write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("segsnr_t", segsnr_t)])
 
     mixture = mixdb.mixture_mixture(m_id=m_id, targets=targets, target=target, noise=noise, force=force)
-    write_data.append(("mixture", mixture))
-
+    result.mixture = mixture
     if write:
-        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, write_data)
-        write_mixture_metadata(mixdb, mixdb.mixture(m_id))
-
-    return GenMixData(
-        targets=targets,
-        target=target,
-        noise=noise,
-        mixture=mixture,
-        truth_t=truth_t,
-        segsnr_t=segsnr_t,
-    )
+        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("mixture", mixture)])
+        write_mixture_metadata(mixdb, m_id)
+
+    return result
 
 
 def main() -> None:
@@ -150,6 +159,7 @@ def main() -> None:
     from sonusai import initial_log_messages
     from sonusai import logger
     from sonusai import update_console_handler
+    from sonusai.mixture import SAMPLE_RATE
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import check_audio_files_exist
     from sonusai.utils import human_readable_size
@@ -161,6 +171,7 @@ def main() -> None:
     save_target = args["--target"]
     compute_truth = args["--truth"]
     compute_segsnr = args["--segsnr"]
+    no_par = args["--nopar"]
 
     start_time = time.monotonic()
 
@@ -173,7 +184,7 @@ def main() -> None:
     mixids = mixdb.mixids_to_list(mixids)
 
     total_samples = mixdb.total_samples(mixids)
-    duration = total_samples / sonusai.mixture.SAMPLE_RATE
+    duration = total_samples / SAMPLE_RATE
 
     logger.info("")
     logger.info(f"Found {len(mixids):,} mixtures to process")
@@ -190,6 +201,7 @@ def main() -> None:
             compute_segsnr=compute_segsnr,
             write=True,
             show_progress=True,
+            no_par=no_par,
         )
     except Exception as e:
         logger.debug(e)