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

sonusai/genmix.py

@@ -17,21 +17,22 @@ Inputs:
     MIXID       A glob of mixture ID(s) to generate.
 
 Outputs the following to the mixture database directory:
-    <id>.h5:
-        dataset:    mixture
-        dataset:    targets
-        dataset:    noise
-        dataset:    target (optional)
-        dataset:    truth_t (optional)
-        dataset:    segsnr_t (optional)
-    <id>.txt
+    <id>
+        mixture.pkl
+        targets.pkl
+        noise.pkl
+        target.pkl (optional)
+        truth_t.pkl (optional)
+        segsnr_t.pkl (optional)
+        metadata.txt
     genmix.log
 """
+
 import signal
 from dataclasses import dataclass
 
-from sonusai.mixture import GenMixData
 from sonusai.mixture import GeneralizedIDs
+from sonusai.mixture import GenMixData
 from sonusai.mixture import MixtureDatabase
 
 
@@ -40,7 +41,7 @@ def signal_handler(_sig, _frame):
 
     from sonusai import logger
 
-    logger.info('Canceled due to keyboard interrupt')
+    logger.info("Canceled due to keyboard interrupt")
     sys.exit(1)
 
 
@@ -49,60 +50,70 @@ signal.signal(signal.SIGINT, signal_handler)
 
 @dataclass
 class MPGlobal:
-    mixdb: MixtureDatabase = None
-    save_target: bool = None
-    compute_truth: bool = None
-    compute_segsnr: bool = None
-    force: bool = None
-    write: bool = None
-
-
-MP_GLOBAL = MPGlobal()
-
+    mixdb: MixtureDatabase
+    save_target: bool
+    compute_truth: bool
+    compute_segsnr: bool
+    force: bool
+    write: bool
+
+
+MP_GLOBAL: MPGlobal
+
+
+def genmix(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    save_target: bool = False,
+    compute_truth: bool = False,
+    compute_segsnr: bool = False,
+    write: bool = False,
+    show_progress: bool = False,
+    force: bool = True,
+) -> list[GenMixData]:
+    from sonusai.utils import par_track
+    from sonusai.utils import track
 
-def genmix(location: str,
-           mixids: GeneralizedIDs = None,
-           save_target: bool = False,
-           compute_truth: bool = False,
-           compute_segsnr: bool = False,
-           write: bool = False,
-           show_progress: bool = False,
-           force: bool = True) -> list[GenMixData]:
-    from tqdm import tqdm
-
-    from sonusai.utils import pp_tqdm_imap
-
-    mixdb = MixtureDatabase(location)
     mixids = mixdb.mixids_to_list(mixids)
-    progress = tqdm(total=len(mixids), disable=not show_progress)
-    results = pp_tqdm_imap(_genmix_kernel,
-                           mixids,
-                           initializer=_genmix_initializer,
-                           initargs=(location, save_target, compute_truth, compute_segsnr, force, write),
-                           progress=progress)
+    progress = track(total=len(mixids), disable=not show_progress)
+    results = par_track(
+        _genmix_kernel,
+        mixids,
+        initializer=_genmix_initializer,
+        initargs=(mixdb, save_target, compute_truth, compute_segsnr, force, write),
+        progress=progress,
+    )
     progress.close()
 
     return results
 
 
-def _genmix_initializer(location: str,
-                        save_target: bool,
-                        compute_truth: bool,
-                        compute_segsnr: bool,
-                        force: bool,
-                        write: bool) -> None:
-    MP_GLOBAL.mixdb = MixtureDatabase(location)
-    MP_GLOBAL.save_target = save_target
-    MP_GLOBAL.compute_truth = compute_truth
-    MP_GLOBAL.compute_segsnr = compute_segsnr
-    MP_GLOBAL.force = force
-    MP_GLOBAL.write = write
+def _genmix_initializer(
+    mixdb: MixtureDatabase,
+    save_target: bool,
+    compute_truth: bool,
+    compute_segsnr: bool,
+    force: bool,
+    write: bool,
+) -> None:
+    global MP_GLOBAL
+
+    MP_GLOBAL = MPGlobal(
+        mixdb=mixdb,
+        save_target=save_target,
+        compute_truth=compute_truth,
+        compute_segsnr=compute_segsnr,
+        force=force,
+        write=write,
+    )
 
 
 def _genmix_kernel(m_id: int) -> GenMixData:
-    from sonusai.mixture import write_mixture_data
+    from sonusai.mixture import write_cached_data
     from sonusai.mixture import write_mixture_metadata
 
+    global MP_GLOBAL
+
     mixdb = MP_GLOBAL.mixdb
     save_target = MP_GLOBAL.save_target
     compute_truth = MP_GLOBAL.compute_truth
@@ -112,44 +123,39 @@ def _genmix_kernel(m_id: int) -> GenMixData:
 
     targets = mixdb.mixture_targets(m_id=m_id, force=force)
     noise = mixdb.mixture_noise(m_id=m_id, force=force)
-    write_data = [('targets', targets), ('noise', noise)]
+    write_data = [("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))
+        write_data.append(("truth_t", truth_t))
     else:
         truth_t = None
 
     target = mixdb.mixture_target(m_id=m_id, targets=targets)
     if save_target:
-        write_data.append(('target', target))
+        write_data.append(("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))
+        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
 
-    mixture = mixdb.mixture_mixture(m_id=m_id,
-                                    targets=targets,
-                                    target=target,
-                                    noise=noise,
-                                    force=force)
-    write_data.append(('mixture', mixture))
+    mixture = mixdb.mixture_mixture(m_id=m_id, targets=targets, target=target, noise=noise, force=force)
+    write_data.append(("mixture", mixture))
 
     if write:
-        write_mixture_data(mixdb, mixdb.mixture(m_id), write_data)
+        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,
-                      noise=noise,
-                      mixture=mixture,
-                      truth_t=truth_t,
-                      segsnr_t=segsnr_t)
+    return GenMixData(
+        targets=targets,
+        target=target,
+        noise=noise,
+        mixture=mixture,
+        truth_t=truth_t,
+        segsnr_t=segsnr_t,
+    )
 
 
 def main() -> None:
@@ -171,55 +177,61 @@ def main() -> None:
     from sonusai.utils import human_readable_size
     from sonusai.utils import seconds_to_hms
 
-    verbose = args['--verbose']
-    location = args['LOC']
-    mixids = args['--mixid']
-    save_target = args['--target']
-    compute_truth = args['--truth']
-    compute_segsnr = args['--segsnr']
+    verbose = args["--verbose"]
+    location = args["LOC"]
+    mixids = args["--mixid"]
+    save_target = args["--target"]
+    compute_truth = args["--truth"]
+    compute_segsnr = args["--segsnr"]
 
     start_time = time.monotonic()
 
-    create_file_handler(join(location, 'genmix.log'))
+    create_file_handler(join(location, "genmix.log"))
     update_console_handler(verbose)
-    initial_log_messages('genmix')
+    initial_log_messages("genmix")
 
-    logger.info(f'Load mixture database from {location}')
+    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 / sonusai.mixture.SAMPLE_RATE
 
-    logger.info('')
-    logger.info(f'Found {len(mixids):,} mixtures to process')
-    logger.info(f'{total_samples:,} samples')
+    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,
-           save_target=save_target,
-           compute_truth=compute_truth,
-           compute_segsnr=compute_segsnr,
-           write=True,
-           show_progress=True)
-
-    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)}')
+    try:
+        genmix(
+            mixdb=mixdb,
+            mixids=mixids,
+            save_target=save_target,
+            compute_truth=compute_truth,
+            compute_segsnr=compute_segsnr,
+            write=True,
+            show_progress=True,
+        )
+    except Exception as e:
+        logger.debug(e)
+        raise
+
+    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)}')
+        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)}')
+        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('')
+    logger.info(f"Completed in {seconds_to_hms(seconds=end_time - start_time)}")
+    logger.info("")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()