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

sonusai/genft.py

@@ -15,19 +15,20 @@ Inputs:
     MIXID       A glob of mixture ID(s) to generate.
 
 Outputs the following to the mixture database directory:
-    <id>.h5:
-        dataset:    feature
-        dataset:    truth_f
-        dataset:    segsnr (optional)
-    <id>.txt
+    <id>
+        feature.pkl
+        truth_f.pkl
+        segsnr.pkl (optional)
+        metadata.txt
     genft.log
 
 """
+
 import signal
 from dataclasses import dataclass
 
-from sonusai.mixture import GenFTData
 from sonusai.mixture import GeneralizedIDs
+from sonusai.mixture import GenFTData
 from sonusai.mixture import MixtureDatabase
 
 
@@ -36,7 +37,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)
 
 
@@ -45,81 +46,87 @@ signal.signal(signal.SIGINT, signal_handler)
 
 @dataclass
 class MPGlobal:
-    mixdb: MixtureDatabase = None
-    compute_truth: bool = None
-    compute_segsnr: bool = None
-    force: bool = None
-    write: bool = None
+    mixdb: MixtureDatabase
+    compute_truth: bool
+    compute_segsnr: bool
+    force: bool
+    write: bool
 
 
-MP_GLOBAL = MPGlobal()
+MP_GLOBAL: MPGlobal
 
 
-def genft(mixdb: MixtureDatabase,
-          mixids: GeneralizedIDs = None,
-          compute_truth: bool = True,
-          compute_segsnr: bool = False,
-          write: bool = False,
-          show_progress: bool = False,
-          force: bool = True) -> list[GenFTData]:
-    from tqdm import tqdm
-
-    from sonusai.utils import pp_tqdm_imap
+def genft(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    compute_truth: bool = True,
+    compute_segsnr: bool = False,
+    write: bool = False,
+    show_progress: bool = False,
+    force: bool = True,
+) -> list[GenFTData]:
+    from sonusai.utils import par_track
+    from sonusai.utils import track
 
     mixids = mixdb.mixids_to_list(mixids)
 
-    progress = tqdm(total=len(mixids), disable=not show_progress)
-    results = pp_tqdm_imap(_genft_kernel,
-                           mixids,
-                           initializer=_genft_initializer,
-                           initargs=(mixdb.location, compute_truth, compute_segsnr, force, write),
-                           progress=progress)
+    progress = track(total=len(mixids), disable=not show_progress)
+    results = par_track(
+        _genft_kernel,
+        mixids,
+        initializer=_genft_initializer,
+        initargs=(mixdb.location, compute_truth, compute_segsnr, force, write),
+        progress=progress,
+    )
     progress.close()
 
     return results
 
 
-def _genft_initializer(location: str,
-                       compute_truth: bool,
-                       compute_segsnr: bool,
-                       force: bool,
-                       write: bool) -> None:
-    MP_GLOBAL.mixdb = MixtureDatabase(location)
-    MP_GLOBAL.compute_truth = compute_truth
-    MP_GLOBAL.compute_segsnr = compute_segsnr
-    MP_GLOBAL.force = force
-    MP_GLOBAL.write = write
+def _genft_initializer(location: str, compute_truth: bool, compute_segsnr: bool, force: bool, write: bool) -> None:
+    global MP_GLOBAL
+
+    MP_GLOBAL = MPGlobal(
+        mixdb=MixtureDatabase(location),
+        compute_truth=compute_truth,
+        compute_segsnr=compute_segsnr,
+        force=force,
+        write=write,
+    )
 
 
 def _genft_kernel(m_id: int) -> GenFTData:
-    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
     compute_truth = MP_GLOBAL.compute_truth
     compute_segsnr = MP_GLOBAL.compute_segsnr
     force = MP_GLOBAL.force
     write = MP_GLOBAL.write
 
+    result = GenFTData()
+
     feature, truth_f = mixdb.mixture_ft(m_id=m_id, force=force)
-    write_data = [('feature', feature)]
+    write_data = [("feature", feature)]
+    result.feature = feature
 
     if compute_truth:
-        write_data.append(('truth_f', truth_f))
-    else:
-        truth_f = None
+        write_data.append(("truth_f", truth_f))
+        result.truth_f = truth_f
 
     if compute_segsnr:
         segsnr = mixdb.mixture_segsnr(m_id=m_id, force=force)
-        write_data.append(('segsnr', segsnr))
-    else:
-        segsnr = None
+        write_data.append(("segsnr", segsnr))
+        result.segsnr = segsnr
 
     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 GenFTData(feature=feature, truth_f=truth_f, segsnr=segsnr)
+    return result
 
 
 def main() -> None:
@@ -141,53 +148,62 @@ def main() -> None:
     from sonusai.utils import human_readable_size
     from sonusai.utils import seconds_to_hms
 
-    verbose = args['--verbose']
-    mixids = args['--mixid']
-    compute_segsnr = args['--segsnr']
-    location = args['LOC']
+    verbose = args["--verbose"]
+    mixids = args["--mixid"]
+    compute_segsnr = args["--segsnr"]
+    location = args["LOC"]
 
     start_time = time.monotonic()
 
-    create_file_handler(join(location, 'genft.log'))
+    create_file_handler(join(location, "genft.log"))
     update_console_handler(verbose)
-    initial_log_messages('genft')
+    initial_log_messages("genft")
 
-    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
-    total_transform_frames = total_samples // mixdb.ft_config.R
+    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')
+    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(mixdb=mixdb,
-          mixids=mixids,
-          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)}')
+    try:
+        genft(
+            mixdb=mixdb,
+            mixids=mixids,
+            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'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)}')
+        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)}')
+        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('')
+    logger.info(f"Completed in {seconds_to_hms(seconds=end_time - start_time)}")
+    logger.info("")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

sonusai/genmetrics.py

@@ -32,6 +32,7 @@ Generate all available metrics except for mxwer.faster:
 > sonusai genmetrics -x"mxwer.faster" mixdb_loc
 
 """
+
 import signal
 from dataclasses import dataclass
 
@@ -43,7 +44,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)
 
 
@@ -52,28 +53,34 @@ signal.signal(signal.SIGINT, signal_handler)
 
 @dataclass
 class MPGlobal:
-    mixdb: MixtureDatabase = None
-    metrics: set[str] = None
+    mixdb: MixtureDatabase
+    metrics: set[str]
 
 
-MP_GLOBAL = MPGlobal()
+MP_GLOBAL: MPGlobal
 
 
 def _initializer(location: str, metrics: set[str]) -> None:
-    MP_GLOBAL.mixdb = MixtureDatabase(location)
-    MP_GLOBAL.metrics = metrics
+    global MP_GLOBAL
+
+    MP_GLOBAL = MPGlobal(
+        mixdb=MixtureDatabase(location),
+        metrics=metrics,
+    )
 
 
 def _process_mixture(mixid: int) -> None:
-    from sonusai.mixture import write_mixture_data
+    from sonusai.mixture import write_cached_data
+
+    global MP_GLOBAL
 
     mixdb = MP_GLOBAL.mixdb
     metrics = list(MP_GLOBAL.metrics)
 
     values = mixdb.mixture_metrics(m_id=mixid, metrics=metrics, force=True)
-    write_data = list(zip(metrics, values))
+    write_data = list(zip(metrics, values, strict=False))
 
-    write_mixture_data(mixdb, mixdb.mixture(mixid), write_data)
+    write_cached_data(mixdb.location, "mixture", mixdb.mixture(mixid).name, write_data)
 
 
 def main() -> None:
@@ -84,12 +91,12 @@ def main() -> None:
 
     args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
 
-    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(',')]
-    show_supported = args['--supported']
-    location = args['LOC']
+    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(",")]
+    show_supported = args["--supported"]
+    location = args["LOC"]
 
     import sys
     import time
@@ -99,76 +106,82 @@ def main() -> None:
     from sonusai import initial_log_messages
     from sonusai import logger
     from sonusai import update_console_handler
-    from sonusai.utils import pp_tqdm_imap
+    from sonusai.utils import par_track
     from sonusai.utils import seconds_to_hms
-    from tqdm import tqdm
+    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
-    create_file_handler(join(location, 'genmetrics.log'))
+    create_file_handler(join(location, "genmetrics.log"))
     update_console_handler(verbose)
-    initial_log_messages('genmetrics')
+    initial_log_messages("genmetrics")
 
-    logger.info(f'Load mixture database from {location}')
+    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}')
+        logger.info(f"\nSupported metrics:\n\n{supported.pretty}")
         sys.exit(0)
 
-    if includes is None or 'all' in includes:
+    if includes is None or "all" in includes:
         metrics = supported.names
     else:
         metrics = set(includes)
-        if 'mxwer' in metrics:
-            metrics.remove('mxwer')
+        if "mxwer" in metrics:
+            metrics.remove("mxwer")
             for name in mixdb.asr_configs:
-                metrics.add(f'mxwer.{name}')
+                metrics.add(f"mxwer.{name}")
 
     diff = metrics.difference(supported.names)
     if diff:
-        logger.error(f'Unrecognized metric: {", ".join(diff)}')
+        logger.error(f"Unrecognized metric: {', '.join(diff)}")
         sys.exit(1)
 
-    if excludes is None or 'none' in excludes:
-        _excludes = set([])
+    if excludes is None or "none" in excludes:
+        _excludes = set()
     else:
         _excludes = set(excludes)
-        if 'mxwer' in _excludes:
-            _excludes.remove('mxwer')
+        if "mxwer" in _excludes:
+            _excludes.remove("mxwer")
             for name in mixdb.asr_configs:
-                _excludes.add(f'mxwer.{name}')
+                _excludes.add(f"mxwer.{name}")
 
     diff = _excludes.difference(supported.names)
     if diff:
-        logger.error(f'Unrecognized metric: {", ".join(diff)}')
+        logger.error(f"Unrecognized metric: {', '.join(diff)}")
         sys.exit(1)
 
     for exclude in _excludes:
         metrics.discard(exclude)
 
-    logger.info(f'Generating metrics: {", ".join(metrics)}')
+    logger.info(f"Generating metrics: {', '.join(metrics)}")
 
     mixids = mixdb.mixids_to_list(mixids)
-    logger.info('')
-    logger.info(f'Found {len(mixids):,} mixtures to process')
-
-    progress = tqdm(total=len(mixids), desc='genmetrics')
-    pp_tqdm_imap(_process_mixture, mixids,
-                 progress=progress,
-                 initializer=_initializer,
-                 initargs=(location, metrics))
+    logger.info("")
+    logger.info(f"Found {len(mixids):,} mixtures to process")
+
+    progress = track(total=len(mixids), desc="genmetrics")
+    par_track(
+        _process_mixture,
+        mixids,
+        progress=progress,
+        initializer=_initializer,
+        initargs=(location, metrics),
+    )
     progress.close()
 
-    logger.info(f'Wrote metrics for {len(mixids)} mixtures to {location}')
-    logger.info('')
+    logger.info(f"Wrote metrics for {len(mixids)} mixtures to {location}")
+    logger.info("")
 
     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()