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

sonusai/genmix.py

@@ -1,17 +1,16 @@
 """sonusai genmix
 
-usage: genmix [-hvgtsn] [-i MIXID] LOC
+usage: genmix [-hvtsn] [-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].
     -n, --nopar                 Do not run in parallel. [default: False].
 
-Generate SonusAI mixture data from a SonusAI mixture database.
+Create mixture data from a SonusAI mixture database.
 
 Inputs:
     LOC         A SonusAI mixture database directory.
@@ -19,41 +18,26 @@ Inputs:
 
 Outputs the following to the mixture database directory:
     <id>
-        mixture.pkl
-        targets.pkl
+        metadata.txt
+        sources.pkl
+        source.pkl
         noise.pkl
-        target.pkl (optional)
+        mixture.pkl
         truth_t.pkl (optional)
         segsnr_t.pkl (optional)
-        metadata.txt
     genmix.log
 """
 
-import signal
-
-from sonusai.mixture import GeneralizedIDs
-from sonusai.mixture import GenMixData
-
-
-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)
+from sonusai.datatypes import GeneralizedIDs
+from sonusai.datatypes import GenMixData
 
 
 def genmix(
     location: str,
     mixids: GeneralizedIDs = "*",
-    save_target: bool = False,
     compute_truth: bool = False,
     compute_segsnr: bool = False,
-    write: bool = False,
+    cache: bool = False,
     show_progress: bool = False,
     force: bool = True,
     no_par: bool = False,
@@ -71,11 +55,10 @@ def genmix(
         partial(
             _genmix_kernel,
             location=location,
-            save_target=save_target,
             compute_truth=compute_truth,
             compute_segsnr=compute_segsnr,
+            cache=cache,
             force=force,
-            write=write,
         ),
         mixids,
         progress=progress,
@@ -89,56 +72,43 @@ def genmix(
 def _genmix_kernel(
     m_id: int,
     location: str,
-    save_target: bool,
     compute_truth: bool,
     compute_segsnr: bool,
+    cache: bool,
     force: bool,
-    write: bool,
 ) -> GenMixData:
     from sonusai.mixture import MixtureDatabase
-    from sonusai.mixture import write_cached_data
     from sonusai.mixture import write_mixture_metadata
 
     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)
-    result.noise = noise
-    if write:
-        write_cached_data(
-            mixdb.location,
-            "mixture",
-            mixdb.mixture(m_id).name,
-            [
-                ("targets", targets),
-                ("noise", noise),
-            ],
-        )
+    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:
-        truth_t = mixdb.mixture_truth_t(m_id=m_id, targets=targets, noise=noise, force=force)
-        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)
-    result.target = target
-    if save_target and write:
-        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("target", target)])
+        result.truth_t = mixdb.mixture_truth_t(m_id=m_id, force=force, cache=cache)
 
     if compute_segsnr:
-        segsnr_t = mixdb.mixture_segsnr_t(m_id=m_id, targets=targets, target=target, noise=noise, force=force)
-        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)
-    result.mixture = mixture
-    if write:
-        write_cached_data(mixdb.location, "mixture", mixdb.mixture(m_id).name, [("mixture", mixture)])
+        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
@@ -147,10 +117,10 @@ def _genmix_kernel(
 def main() -> None:
     from docopt import docopt
 
-    import sonusai
+    from sonusai import __version__ as sai_version
     from sonusai.utils import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     import time
     from os.path import join
@@ -159,7 +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.constants import SAMPLE_RATE
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import check_audio_files_exist
     from sonusai.utils import human_readable_size
@@ -168,7 +138,6 @@ def main() -> None:
     verbose = args["--verbose"]
     location = args["LOC"]
     mixids = args["--mixid"]
-    save_target = args["--target"]
     compute_truth = args["--truth"]
     compute_segsnr = args["--segsnr"]
     no_par = args["--nopar"]
@@ -192,20 +161,16 @@ def main() -> None:
 
     check_audio_files_exist(mixdb)
 
-    try:
-        genmix(
-            location=location,
-            mixids=mixids,
-            save_target=save_target,
-            compute_truth=compute_truth,
-            compute_segsnr=compute_segsnr,
-            write=True,
-            show_progress=True,
-            no_par=no_par,
-        )
-    except Exception as e:
-        logger.debug(e)
-        raise
+    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("")
@@ -224,4 +189,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)