sonusai 1.0.16__cp311-abi3-macosx_10_12_x86_64.whl

sonusai/genft.py

@@ -0,0 +1,196 @@
+"""sonusai genft
+
+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].
+    -n, --nopar                 Do not run in parallel. [default: False].
+
+Generate SonusAI feature/truth data from a SonusAI mixture database.
+
+Inputs:
+    LOC         A SonusAI mixture database directory.
+    MIXID       A glob of mixture ID(s) to generate.
+
+Outputs the following to the mixture database directory:
+    <id>
+        feature.pkl
+        truth_f.pkl
+        segsnr.pkl (optional)
+        metadata.txt
+    genft.log
+
+"""
+
+from sonusai.datatypes import GeneralizedIDs
+from sonusai.datatypes import GenFTData
+
+
+def genft(
+    location: str,
+    mixids: GeneralizedIDs = "*",
+    compute_truth: bool = True,
+    compute_segsnr: bool = False,
+    cache: bool = False,
+    show_progress: bool = False,
+    force: bool = True,
+    no_par: bool = False,
+) -> list[GenFTData]:
+    from functools import partial
+
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.utils.parallel import par_track
+    from sonusai.utils.parallel import track
+
+    mixdb = MixtureDatabase(location)
+    mixids = mixdb.mixids_to_list(mixids)
+
+    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,
+            cache=cache,
+            force=force,
+        ),
+        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,
+    cache: bool,
+    force: bool,
+) -> GenFTData:
+    from functools import partial
+    from typing import Any
+
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture.data_io import write_cached_data
+    from sonusai.mixture.helpers import write_mixture_metadata
+
+    mixdb = MixtureDatabase(location)
+
+    write_func = partial(write_cached_data, location=mixdb.location, name="mixture", index=mixdb.mixture(m_id).name)
+
+    result = GenFTData()
+
+    mixture = mixdb.mixture_mixture(m_id)
+
+    feature, truth_f = mixdb.mixture_ft(m_id=m_id, mixture=mixture, force=force)
+    result.feature = feature
+    items: dict[str, Any] = {"feature": feature}
+
+    if compute_truth:
+        result.truth_f = truth_f
+        items["truth_f"] = truth_f
+
+    if compute_segsnr:
+        segsnr_t = mixdb.mixture_segsnr_t(m_id)
+        result.segsnr = mixdb.mixture_segsnr(m_id=m_id, segsnr_t=segsnr_t, force=force, cache=cache)
+
+    if cache:
+        write_func(items=items)
+        write_mixture_metadata(mixdb, m_id=m_id)
+
+    return result
+
+
+def main() -> None:
+    from docopt import docopt
+
+    from sonusai import __version__ as sai_version
+    from sonusai.utils.docstring import trim_docstring
+
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
+
+    import time
+    from os.path import join
+
+    from sonusai import create_file_handler
+    from sonusai import initial_log_messages
+    from sonusai import logger
+    from sonusai import update_console_handler
+    from sonusai.constants import SAMPLE_RATE
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture.helpers import check_audio_files_exist
+    from sonusai.utils.human_readable_size import human_readable_size
+    from sonusai.utils.seconds_to_hms import seconds_to_hms
+
+    verbose = args["--verbose"]
+    mixids = args["--mixid"]
+    compute_segsnr = args["--segsnr"]
+    no_par = args["--nopar"]
+    location = args["LOC"]
+
+    start_time = time.monotonic()
+
+    create_file_handler(join(location, "genft.log"), verbose)
+    update_console_handler(verbose)
+    initial_log_messages("genft")
+
+    logger.info(f"Load mixture database from {location}")
+    mixdb = MixtureDatabase(location)
+    mixids = mixdb.mixids_to_list(mixids)
+
+    total_samples = mixdb.total_samples(mixids)
+    duration = total_samples / SAMPLE_RATE
+    total_transform_frames = total_samples // mixdb.ft_config.overlap
+    total_feature_frames = total_samples // mixdb.feature_step_samples
+
+    logger.info("")
+    logger.info(f"Found {len(mixids):,} mixtures to process")
+    logger.info(
+        f"{total_samples:,} samples, "
+        f"{total_transform_frames:,} transform frames, "
+        f"{total_feature_frames:,} feature frames"
+    )
+
+    check_audio_files_exist(mixdb)
+
+    genft(
+        location=location,
+        mixids=mixids,
+        compute_segsnr=compute_segsnr,
+        cache=True,
+        show_progress=True,
+        no_par=no_par,
+    )
+
+    logger.info(f"Wrote {len(mixids)} mixtures to {location}")
+    logger.info("")
+    logger.info(f"Duration: {seconds_to_hms(seconds=duration)}")
+    logger.info(
+        f"feature:  {human_readable_size(total_feature_frames * mixdb.fg_stride * mixdb.feature_parameters * 4, 1)}"
+    )
+    logger.info(f"truth_f:  {human_readable_size(total_feature_frames * mixdb.num_classes * 4, 1)}")
+    if compute_segsnr:
+        logger.info(f"segsnr:   {human_readable_size(total_transform_frames * 4, 1)}")
+
+    end_time = time.monotonic()
+    logger.info(f"Completed in {seconds_to_hms(seconds=end_time - start_time)}")
+    logger.info("")
+
+
+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/genmetrics.py

@@ -0,0 +1,183 @@
+"""sonusai genmetrics
+
+usage: genmetrics [-hvusd] [-i MIXID] [-n INCLUDE] [-x EXCLUDE] [-p NUMPROC] LOC
+
+options:
+    -h, --help
+    -v, --verbose                   Be verbose.
+    -i MIXID, --mixid MIXID         Mixture ID(s) to generate. [default: *].
+    -n INCLUDE, --include INCLUDE   Metrics to include. [default: all]
+    -x EXCLUDE, --exclude EXCLUDE   Metrics to exclude. [default: none]
+    -p NUMPROC, --nproc NUMPROC     Number of parallel processes to use. Default single thread.
+    -u, --update                    Update metrics (do not regenerate existing metrics).
+    -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 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
+
+Generate only mxwer.faster metrics:
+> sonusai genmetrics -n"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
+
+"""
+
+
+def _process_mixture(mixid: int, location: str, metrics: list[str], update: bool = False) -> set[str]:
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture.data_io import write_cached_data
+    from sonusai.mixture.helpers import write_mixture_metadata
+
+    mixdb = MixtureDatabase(location)
+    results = mixdb.mixture_metrics(m_id=mixid, metrics=metrics, force=not update)
+    write_cached_data(mixdb.location, "mixture", mixdb.mixture(mixid).name, results)
+    write_mixture_metadata(mixdb, mixture=mixdb.mixture(mixid))
+
+    return set(results.keys())
+
+
+def main() -> None:
+    from docopt import docopt
+
+    from sonusai import __version__ as sai_version
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.utils.docstring import trim_docstring
+
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
+
+    verbose = args["--verbose"]
+    mixids = args["--mixid"]
+    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 != ""}
+    update = args["--update"]
+    show_supported = args["--supported"]
+    dryrun = args["--dryrun"]
+    num_proc = args["--nproc"]
+    location = args["LOC"]
+
+    import fnmatch
+    import sys
+    import time
+    from functools import partial
+    from os.path import join
+
+    from sonusai import create_file_handler
+    from sonusai import initial_log_messages
+    from sonusai import logger
+    from sonusai import update_console_handler
+    from sonusai.utils.parallel import par_track
+    from sonusai.utils.parallel import track
+    from sonusai.utils.seconds_to_hms import seconds_to_hms
+
+    start_time = time.monotonic()
+
+    # Setup logging file
+    create_file_handler(join(location, "genmetrics.log"), verbose)
+    update_console_handler(verbose)
+    initial_log_messages("genmetrics")
+
+    logger.info(f"Load mixture database from {location}")
+
+    mixdb = MixtureDatabase(location)
+    supported = mixdb.supported_metrics
+    if show_supported:
+        logger.info(f"\nSupported metrics:\n\n{supported.pretty}")
+        sys.exit(0)
+
+    # Handle default excludes
+    if "none" in excludes:
+        if "all" in includes:
+            excludes = {"*asr*", "*wer*"}
+        else:
+            excludes = set()
+
+    # 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)
+
+    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
+
+    metrics = sorted(requested)
+
+    if len(metrics) == 0:
+        logger.warning("No metrics were requested")
+        sys.exit(1)
+
+    logger.info("Generating metrics:")
+    logger.info(f"{', '.join(metrics)}")
+    if dryrun:
+        sys.exit(0)
+
+    mixids = mixdb.mixids_to_list(mixids)
+    logger.info("")
+    logger.info(f"Found {len(mixids):,} mixtures to process")
+
+    no_par = num_proc == 1 or len(mixids) == 1
+
+    if num_proc is not None:
+        num_proc = int(num_proc)
+
+    progress = track(total=len(mixids), desc="genmetrics")
+    results = par_track(
+        partial(_process_mixture, location=location, metrics=metrics, update=update),
+        mixids,
+        progress=progress,
+        num_cpus=num_proc,
+        no_par=no_par,
+    )
+    progress.close()
+
+    written_metrics = sorted(set().union(*results))
+    logger.info(f"Wrote metrics for {len(mixids)} mixtures to {location}:")
+    logger.info(f"{', '.join(written_metrics)}")
+    logger.info("")
+
+    end_time = time.monotonic()
+    logger.info(f"Completed in {seconds_to_hms(seconds=end_time - start_time)}")
+    logger.info("")
+
+
+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/genmix.py

@@ -0,0 +1,199 @@
+"""sonusai genmix
+
+usage: genmix [-hvtsn] [-i MIXID] LOC
+
+options:
+    -h, --help
+    -v, --verbose               Be verbose.
+    -i MIXID, --mixid MIXID     Mixture ID(s) to generate. [default: *].
+    -t, --truth                 Save truth_t. [default: False].
+    -s, --segsnr                Save segsnr_t. [default: False].
+    -n, --nopar                 Do not run in parallel. [default: False].
+
+Create mixture data from a SonusAI mixture database.
+
+Inputs:
+    LOC         A SonusAI mixture database directory.
+    MIXID       A glob of mixture ID(s) to generate.
+
+Outputs the following to the mixture database directory:
+    <id>
+        metadata.txt
+        sources.pkl
+        source.pkl
+        noise.pkl
+        mixture.pkl
+        truth_t.pkl (optional)
+        segsnr_t.pkl (optional)
+    genmix.log
+"""
+
+from sonusai.datatypes import GeneralizedIDs
+from sonusai.datatypes import GenMixData
+
+
+def genmix(
+    location: str,
+    mixids: GeneralizedIDs = "*",
+    compute_truth: bool = False,
+    compute_segsnr: bool = False,
+    cache: bool = False,
+    show_progress: bool = False,
+    force: bool = True,
+    no_par: bool = False,
+) -> list[GenMixData]:
+    from functools import partial
+
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.utils.parallel import par_track
+    from sonusai.utils.parallel import track
+
+    mixdb = MixtureDatabase(location)
+    mixids = mixdb.mixids_to_list(mixids)
+    progress = track(total=len(mixids), disable=not show_progress)
+    results = par_track(
+        partial(
+            _genmix_kernel,
+            location=location,
+            compute_truth=compute_truth,
+            compute_segsnr=compute_segsnr,
+            cache=cache,
+            force=force,
+        ),
+        mixids,
+        progress=progress,
+        no_par=no_par,
+    )
+    progress.close()
+
+    return results
+
+
+def _genmix_kernel(
+    m_id: int,
+    location: str,
+    compute_truth: bool,
+    compute_segsnr: bool,
+    cache: bool,
+    force: bool,
+) -> GenMixData:
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture.helpers import write_mixture_metadata
+
+    mixdb = MixtureDatabase(location)
+
+    result = GenMixData()
+    result.sources = mixdb.mixture_sources(m_id=m_id, force=force, cache=cache)
+    result.source = mixdb.mixture_source(m_id=m_id, sources=result.sources, force=force, cache=cache)
+    result.noise = mixdb.mixture_noise(m_id=m_id, sources=result.sources, force=force, cache=cache)
+    result.mixture = mixdb.mixture_mixture(
+        m_id=m_id,
+        sources=result.sources,
+        source=result.source,
+        noise=result.noise,
+        force=force,
+        cache=cache,
+    )
+
+    if compute_truth:
+        result.truth_t = mixdb.mixture_truth_t(m_id=m_id, force=force, cache=cache)
+
+    if compute_segsnr:
+        result.segsnr_t = mixdb.mixture_segsnr_t(
+            m_id=m_id,
+            sources=result.sources,
+            source=result.source,
+            noise=result.noise,
+            force=force,
+            cache=cache,
+        )
+
+    if cache:
+        write_mixture_metadata(mixdb, m_id=m_id)
+
+    return result
+
+
+def main() -> None:
+    from docopt import docopt
+
+    from sonusai import __version__ as sai_version
+    from sonusai.utils.docstring import trim_docstring
+
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
+
+    import time
+    from os.path import join
+
+    from sonusai import create_file_handler
+    from sonusai import initial_log_messages
+    from sonusai import logger
+    from sonusai import update_console_handler
+    from sonusai.constants import SAMPLE_RATE
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture.helpers import check_audio_files_exist
+    from sonusai.utils.human_readable_size import human_readable_size
+    from sonusai.utils.seconds_to_hms import seconds_to_hms
+
+    verbose = args["--verbose"]
+    location = args["LOC"]
+    mixids = args["--mixid"]
+    compute_truth = args["--truth"]
+    compute_segsnr = args["--segsnr"]
+    no_par = args["--nopar"]
+
+    start_time = time.monotonic()
+
+    create_file_handler(join(location, "genmix.log"), verbose)
+    update_console_handler(verbose)
+    initial_log_messages("genmix")
+
+    logger.info(f"Load mixture database from {location}")
+    mixdb = MixtureDatabase(location)
+    mixids = mixdb.mixids_to_list(mixids)
+
+    total_samples = mixdb.total_samples(mixids)
+    duration = total_samples / SAMPLE_RATE
+
+    logger.info("")
+    logger.info(f"Found {len(mixids):,} mixtures to process")
+    logger.info(f"{total_samples:,} samples")
+
+    check_audio_files_exist(mixdb)
+
+    genmix(
+        location=location,
+        mixids=mixids,
+        compute_truth=compute_truth,
+        compute_segsnr=compute_segsnr,
+        cache=True,
+        show_progress=True,
+        force=True,
+        no_par=no_par,
+    )
+
+    logger.info(f"Wrote {len(mixids)} mixtures to {location}")
+    logger.info("")
+    logger.info(f"Duration: {seconds_to_hms(seconds=duration)}")
+    logger.info(f"mixture:  {human_readable_size(total_samples * 2, 1)}")
+    if compute_truth:
+        logger.info(f"truth_t:  {human_readable_size(total_samples * mixdb.num_classes * 4, 1)}")
+    logger.info(f"target:   {human_readable_size(total_samples * 2, 1)}")
+    logger.info(f"noise:    {human_readable_size(total_samples * 2, 1)}")
+    if compute_segsnr:
+        logger.info(f"segsnr:   {human_readable_size(total_samples * 4, 1)}")
+
+    end_time = time.monotonic()
+    logger.info(f"Completed in {seconds_to_hms(seconds=end_time - start_time)}")
+    logger.info("")
+
+
+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)