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

sonusai/onnx_predict.py

@@ -41,11 +41,12 @@ TBD not sure below make sense, need to continue ??
 3. Classification
 
 Outputs the following to opredict-<TIMESTAMP> directory:
-    <id>.h5
-        dataset:    predict
+    <id>
+        predict.pkl
     onnx_predict.log
 
 """
+
 import signal
 
 
@@ -54,7 +55,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)
 
 
@@ -69,12 +70,12 @@ def main() -> None:
 
     args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
 
-    verbose = args['--verbose']
-    wav = args['--write-wav']
-    mixids = args['--mixid']
-    include = args['--include']
-    model_path = args['MODEL']
-    data_paths = args['DATA']
+    verbose = args["--verbose"]
+    wav = args["--write-wav"]
+    mixids = args["--mixid"]
+    include = args["--include"]
+    model_path = args["MODEL"]
+    data_paths = args["DATA"]
 
     from os import makedirs
     from os.path import abspath
@@ -103,8 +104,8 @@ def main() -> None:
     from sonusai.utils import write_audio
 
     mixdb_path = None
-    mixdb = None
-    p_mixids = None
+    mixdb: MixtureDatabase | None = None
+    p_mixids: list[int] = []
     entries: list[PathInfo] = []
 
     if len(data_paths) == 1 and isdir(data_paths[0]):
@@ -113,96 +114,98 @@ def main() -> None:
         mixdb_path = data_paths[0]
     else:
         # search all data paths for .wav, .flac (or whatever is specified in include)
-        in_basename = ''
+        in_basename = ""
 
-    output_dir = create_ts_name('opredict-' + in_basename)
+    output_dir = create_ts_name("opredict-" + in_basename)
     makedirs(output_dir, exist_ok=True)
 
     # Setup logging file
-    create_file_handler(join(output_dir, 'onnx-predict.log'))
+    create_file_handler(join(output_dir, "onnx-predict.log"))
     update_console_handler(verbose)
-    initial_log_messages('onnx_predict')
+    initial_log_messages("onnx_predict")
 
     providers = ort.get_available_providers()
-    logger.info(f'Loaded ONNX Runtime, available providers: {providers}.')
+    logger.info(f"Loaded ONNX Runtime, available providers: {providers}.")
 
     session, options, model_root, hparams, sess_inputs, sess_outputs = load_ort_session(model_path)
     if hparams is None:
-        logger.error(f'Error: ONNX model does not have required SonusAI hyperparameters, cannot proceed.')
+        logger.error("Error: ONNX model does not have required SonusAI hyperparameters, cannot proceed.")
         raise SystemExit(1)
     if len(sess_inputs) != 1:
-        logger.error(f'Error: ONNX model does not have 1 input, but {len(sess_inputs)}. Exit due to unknown input.')
+        logger.error(f"Error: ONNX model does not have 1 input, but {len(sess_inputs)}. Exit due to unknown input.")
 
     in0name = sess_inputs[0].name
     in0type = sess_inputs[0].type
     out_names = [n.name for n in session.get_outputs()]
 
-    logger.info(f'Read and compiled ONNX model from {model_path}.')
+    logger.info(f"Read and compiled ONNX model from {model_path}.")
 
     if mixdb_path is not None:
         # Assume it's a single path to SonusAI mixdb subdir
-        logger.debug(f'Attempting to load mixture database from {mixdb_path}')
+        logger.debug(f"Attempting to load mixture database from {mixdb_path}")
         mixdb = MixtureDatabase(mixdb_path)
-        logger.info(f'SonusAI mixdb: found {mixdb.num_mixtures} mixtures with {mixdb.num_classes} classes')
+        logger.info(f"SonusAI mixdb: found {mixdb.num_mixtures} mixtures with {mixdb.num_classes} classes")
         p_mixids = mixdb.mixids_to_list(mixids)
         if len(p_mixids) != mixdb.num_mixtures:
-            logger.info(f'Processing a subset of {p_mixids} from available mixtures.')
+            logger.info(f"Processing a subset of {p_mixids} from available mixtures.")
     else:
         for p in data_paths:
-            location = join(realpath(abspath(p)), '**', include)
-            logger.debug(f'Processing {location}')
+            location = join(realpath(abspath(p)), "**", include)
+            logger.debug(f"Processing {location}")
             for file in braced_iglob(pathname=location, recursive=True):
                 name = file
                 entries.append(PathInfo(abs_path=file, audio_filepath=name))
-        logger.info(f'{len(data_paths)} data paths specified, found {len(entries)} audio files.')
+        logger.info(f"{len(data_paths)} data paths specified, found {len(entries)} audio files.")
 
-    if in0type.find('float16') != -1:
+    if in0type.find("float16") != -1:
         model_is_fp16 = True
-        logger.info(f'Detected input of float16, converting all feature inputs to that type.')
+        logger.info("Detected input of float16, converting all feature inputs to that type.")
     else:
         model_is_fp16 = False
 
-    if mixdb_path is not None and hparams['batch_size'] == 1:
+    if mixdb is not None and hparams["batch_size"] == 1:
         # mixdb input
         # Assume (of course) that mixdb feature, etc. is what model expects
-        if hparams['feature'] != mixdb.feature:
-            logger.warning(f'Mixture feature does not match model feature, this inference run may fail.')
+        if hparams["feature"] != mixdb.feature:
+            logger.warning("Mixture feature does not match model feature, this inference run may fail.")
         # no choice, can't use hparams.feature since it's different from the mixdb
         feature_mode = mixdb.feature
 
         for mixid in p_mixids:
             # frames x stride x feature_params
             feature, _ = mixdb.mixture_ft(mixid)
-            if hparams['timesteps'] == 0:
+            if hparams["timesteps"] == 0:
                 # no timestep dimension, reshape will handle
                 timesteps = 0
             else:
                 # fit frames into timestep dimension (TSE mode)
                 timesteps = feature.shape[0]
 
-            feature, _ = reshape_inputs(feature=feature,
-                                        batch_size=1,
-                                        timesteps=timesteps,
-                                        flatten=hparams['flatten'],
-                                        add1ch=hparams['add1ch'])
+            feature, _ = reshape_inputs(
+                feature=feature,
+                batch_size=1,
+                timesteps=timesteps,
+                flatten=hparams["flatten"],
+                add1ch=hparams["add1ch"],
+            )
             if model_is_fp16:
-                feature = np.float16(feature)  # type: ignore
+                feature = np.float16(feature)  # type: ignore[assignment]
             # run inference, ort session wants i.e. batch x timesteps x feat_params, outputs numpy BxTxFP or BxFP
             predict = session.run(out_names, {in0name: feature})[0]
             # predict, _ = reshape_outputs(predict=predict[0], timesteps=frames)  # frames x feat_params
             output_fname = join(output_dir, mixdb.mixtures[mixid].name)
-            with h5py.File(output_fname, 'a') as f:
-                if 'predict' in f:
-                    del f['predict']
-                f.create_dataset('predict', data=predict)
+            with h5py.File(output_fname, "a") as f:
+                if "predict" in f:
+                    del f["predict"]
+                f.create_dataset("predict", data=predict)
             if wav:
                 # note only makes sense if model is predicting audio, i.e., timestep dimension exists
                 # predict_audio wants [frames, channels, feature_parameters] equivalent to timesteps, batch, bins
                 predict = np.transpose(predict, [1, 0, 2])
                 predict_audio = get_audio_from_feature(feature=predict, feature_mode=feature_mode)
-                owav_name = splitext(output_fname)[0] + '_predict.wav'
+                owav_name = splitext(output_fname)[0] + "_predict.wav"
                 write_audio(owav_name, predict_audio)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

sonusai/queries/__init__.py

@@ -1,6 +1,8 @@
 # SonusAI query utilities
+# ruff: noqa: F401
+
 from .queries import get_mixids_from_noise
 from .queries import get_mixids_from_snr
 from .queries import get_mixids_from_target
 from .queries import get_mixids_from_truth_function
-from .queries import get_mixids_from_truth_index
+from .queries import get_mixids_from_class_indices

sonusai/queries/queries.py

@@ -1,14 +1,16 @@
+from collections.abc import Callable
 from typing import Any
-from typing import Callable
 
 from sonusai.mixture.datatypes import GeneralizedIDs
 from sonusai.mixture.mixdb import MixtureDatabase
 
 
-def get_mixids_from_mixture_field_predicate(mixdb: MixtureDatabase,
-                                            field: str,
-                                            mixids: GeneralizedIDs = None,
-                                            predicate: Callable[[Any], bool] = None) -> dict[int, list[int]]:
+def get_mixids_from_mixture_field_predicate(
+    mixdb: MixtureDatabase,
+    field: str,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[int, list[int]]:
     """
     Generate mixture IDs based on mixture field and predicate
     Return a dictionary where:
@@ -18,6 +20,7 @@ def get_mixids_from_mixture_field_predicate(mixdb: MixtureDatabase,
     mixid_out = mixdb.mixids_to_list(mixids)
 
     if predicate is None:
+
         def predicate(_: Any) -> bool:
             return True
 
@@ -30,7 +33,7 @@ def get_mixids_from_mixture_field_predicate(mixdb: MixtureDatabase,
                     criteria_set.add(v)
         elif predicate(value):
             criteria_set.add(value)
-    criteria = sorted(list(criteria_set))
+    criteria = sorted(criteria_set)
 
     result: dict[int, list[int]] = {}
     for criterion in criteria:
@@ -47,22 +50,27 @@ def get_mixids_from_mixture_field_predicate(mixdb: MixtureDatabase,
     return result
 
 
-def get_mixids_from_truth_settings_field_predicate(mixdb: MixtureDatabase,
-                                                   field: str,
-                                                   mixids: GeneralizedIDs = None,
-                                                   predicate: Callable[[Any], bool] = None) -> dict[int, list[int]]:
+def get_mixids_from_truth_configs_field_predicate(
+    mixdb: MixtureDatabase,
+    field: str,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[int, list[int]]:
     """
-    Generate mixture IDs based on target truth_settings field and predicate
+    Generate mixture IDs based on target truth_configs field and predicate
     Return a dictionary where:
         - keys are the matching field values
         - values are lists of the mixids that match the criteria
     """
+    from sonusai.mixture import REQUIRED_TRUTH_CONFIGS
+
     mixid_out = mixdb.mixids_to_list(mixids)
 
     # Get all field values
-    values = get_all_truth_settings_values_from_field(mixdb, field)
+    values = get_all_truth_configs_values_from_field(mixdb, field)
 
     if predicate is None:
+
         def predicate(_: Any) -> bool:
             return True
 
@@ -75,10 +83,14 @@ def get_mixids_from_truth_settings_field_predicate(mixdb: MixtureDatabase,
         indices = []
         for t_id in mixdb.target_file_ids:
             target = mixdb.target_file(t_id)
-            for truth_setting in target.truth_settings:
-                if value in getattr(truth_setting, field):
-                    indices.append(t_id)
-        indices = sorted(list(set(indices)))
+            for truth_config in target.truth_configs.values():
+                if field in REQUIRED_TRUTH_CONFIGS:
+                    if value in getattr(truth_config, field):
+                        indices.append(t_id)
+                else:
+                    if value in getattr(truth_config.config, field):
+                        indices.append(t_id)
+        indices = sorted(set(indices))
 
         mixids = []
         for index in indices:
@@ -86,61 +98,66 @@ def get_mixids_from_truth_settings_field_predicate(mixdb: MixtureDatabase,
                 if index in [target.file_id for target in mixdb.mixture(m_id).targets]:
                     mixids.append(m_id)
 
-        mixids = sorted(list(set(mixids)))
+        mixids = sorted(set(mixids))
         if mixids:
             result[value] = mixids
 
     return result
 
 
-def get_all_truth_settings_values_from_field(mixdb: MixtureDatabase, field: str) -> list:
+def get_all_truth_configs_values_from_field(mixdb: MixtureDatabase, field: str) -> list:
     """
-    Generate a list of all values corresponding to the given field in truth_settings
+    Generate a list of all values corresponding to the given field in truth_configs
     """
+    from sonusai.mixture import REQUIRED_TRUTH_CONFIGS
+
     result = []
     for target in mixdb.target_files:
-        for truth_setting in target.truth_settings:
-            value = getattr(truth_setting, field)
+        for truth_config in target.truth_configs.values():
+            if field in REQUIRED_TRUTH_CONFIGS:
+                value = getattr(truth_config, field)
+            else:
+                value = getattr(truth_config.config, field, None)
             if isinstance(value, str):
                 value = [value]
             result.extend(value)
 
-    return sorted(list(set(result)))
+    return sorted(set(result))
 
 
-def get_mixids_from_noise(mixdb: MixtureDatabase,
-                          mixids: GeneralizedIDs = None,
-                          predicate: Callable[[Any], bool] = None) -> dict[int, list[int]]:
+def get_mixids_from_noise(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[int, list[int]]:
     """
     Generate mixids based on noise index predicate
     Return a dictionary where:
         - keys are the noise indices
         - values are lists of the mixids that match the noise index
     """
-    return get_mixids_from_mixture_field_predicate(mixdb=mixdb,
-                                                   mixids=mixids,
-                                                   field='noise_id',
-                                                   predicate=predicate)
+    return get_mixids_from_mixture_field_predicate(mixdb=mixdb, mixids=mixids, field="noise_id", predicate=predicate)
 
 
-def get_mixids_from_target(mixdb: MixtureDatabase,
-                           mixids: GeneralizedIDs = None,
-                           predicate: Callable[[Any], bool] = None) -> dict[int, list[int]]:
+def get_mixids_from_target(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[int, list[int]]:
     """
     Generate mixids based on a target index predicate
     Return a dictionary where:
         - keys are the target indices
         - values are lists of the mixids that match the target index
     """
-    return get_mixids_from_mixture_field_predicate(mixdb=mixdb,
-                                                   mixids=mixids,
-                                                   field='target_ids',
-                                                   predicate=predicate)
+    return get_mixids_from_mixture_field_predicate(mixdb=mixdb, mixids=mixids, field="target_ids", predicate=predicate)
 
 
-def get_mixids_from_snr(mixdb: MixtureDatabase,
-                        mixids: GeneralizedIDs = None,
-                        predicate: Callable[[Any], bool] = None) -> dict[float, list[int]]:
+def get_mixids_from_snr(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[float, list[int]]:
     """
     Generate mixids based on an SNR predicate
     Return a dictionary where:
@@ -155,46 +172,70 @@ def get_mixids_from_snr(mixdb: MixtureDatabase,
     snrs = [float(snr) for snr in mixdb.all_snrs if not snr.is_random]
 
     if predicate is None:
+
         def predicate(_: Any) -> bool:
             return True
 
     # Get only the SNRs of interest (filter on predicate)
     snrs = [snr for snr in snrs if predicate(snr)]
 
-    result = {}
+    result: dict[float, list[int]] = {}
     for snr in snrs:
         # Get a list of mixids for each SNR
-        result[snr] = sorted(
-            [i for i, mixture in enumerate(mixdb.mixtures) if mixture.snr == snr and i in mixid_out])
+        result[snr] = sorted([i for i, mixture in enumerate(mixdb.mixtures) if mixture.snr == snr and i in mixid_out])
 
     return result
 
 
-def get_mixids_from_truth_index(mixdb: MixtureDatabase,
-                                mixids: GeneralizedIDs = None,
-                                predicate: Callable[[Any], bool] = None) -> dict[int, list[int]]:
+def get_mixids_from_class_indices(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[int, list[int]]:
     """
-    Generate mixids based on a truth index predicate
+    Generate mixids based on a class index predicate
     Return a dictionary where:
-        - keys are the truth indices
-        - values are lists of the mixids that match the truth index
+        - keys are the class indices
+        - values are lists of the mixids that match the class index
     """
-    return get_mixids_from_truth_settings_field_predicate(mixdb=mixdb,
-                                                          mixids=mixids,
-                                                          field='index',
-                                                          predicate=predicate)
+    mixid_out = mixdb.mixids_to_list(mixids)
+
+    if predicate is None:
+
+        def predicate(_: Any) -> bool:
+            return True
+
+    criteria_set = set()
+    for m_id in mixid_out:
+        class_indices = mixdb.mixture_class_indices(m_id)
+        for class_index in class_indices:
+            if predicate(class_index):
+                criteria_set.add(class_index)
+    criteria = sorted(criteria_set)
+
+    result: dict[int, list[int]] = {}
+    for criterion in criteria:
+        result[criterion] = []
+        for m_id in mixid_out:
+            class_indices = mixdb.mixture_class_indices(m_id)
+            for class_index in class_indices:
+                if class_index == criterion:
+                    result[criterion].append(m_id)
+
+    return result
 
 
-def get_mixids_from_truth_function(mixdb: MixtureDatabase,
-                                   mixids: GeneralizedIDs = None,
-                                   predicate: Callable[[Any], bool] = None) -> dict[int, list[int]]:
+def get_mixids_from_truth_function(
+    mixdb: MixtureDatabase,
+    mixids: GeneralizedIDs = "*",
+    predicate: Callable[[Any], bool] | None = None,
+) -> dict[int, list[int]]:
     """
     Generate mixids based on a truth function predicate
     Return a dictionary where:
         - keys are the truth functions
         - values are lists of the mixids that match the truth function
     """
-    return get_mixids_from_truth_settings_field_predicate(mixdb=mixdb,
-                                                          mixids=mixids,
-                                                          field='function',
-                                                          predicate=predicate)
+    return get_mixids_from_truth_configs_field_predicate(
+        mixdb=mixdb, mixids=mixids, field="function", predicate=predicate
+    )

sonusai/speech/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401
+
 from .textgrid import annotate_textgrid
 from .textgrid import create_textgrid
 from .types import TimeAlignedType

sonusai/speech/l2arctic.py

@@ -1,7 +1,6 @@
 import os
 import string
 from pathlib import Path
-from typing import Optional
 
 from .types import TimeAlignedType
 
@@ -9,54 +8,54 @@ from .types import TimeAlignedType
 def _get_duration(name: str) -> float:
     import soundfile
 
-    from sonusai import SonusAIError
-
     try:
         return soundfile.info(name).duration
     except Exception as e:
-        raise SonusAIError(f'Error reading {name}: {e}')
+        raise OSError(f"Error reading {name}: {e}") from e
 
 
-def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
+def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
     """Load time-aligned text data given a L2-ARCTIC audio file.
 
     :param audio: Path to the L2-ARCTIC audio file.
     :return: A TimeAlignedType object.
     """
-    file = Path(audio).parent.parent / 'transcript' / (Path(audio).stem + '.txt')
+    file = Path(audio).parent.parent / "transcript" / (Path(audio).stem + ".txt")
     if not os.path.exists(file):
         return None
 
-    with open(file, mode='r', encoding='utf-8') as f:
+    with open(file, encoding="utf-8") as f:
         line = f.read()
 
-    return TimeAlignedType(0,
-                           _get_duration(str(audio)),
-                           line.strip().lower().translate(str.maketrans('', '', string.punctuation)))
+    return TimeAlignedType(
+        0,
+        _get_duration(str(audio)),
+        line.strip().lower().translate(str.maketrans("", "", string.punctuation)),
+    )
 
 
-def load_words(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
+def load_words(audio: str | os.PathLike[str]) -> list[TimeAlignedType] | None:
     """Load time-aligned word data given a L2-ARCTIC audio file.
 
     :param audio: Path to the L2-ARCTIC audio file.
     :return: A list of TimeAlignedType objects.
     """
-    return _load_ta(audio, 'words')
+    return _load_ta(audio, "words")
 
 
-def load_phonemes(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
+def load_phonemes(audio: str | os.PathLike[str]) -> list[TimeAlignedType] | None:
     """Load time-aligned phonemes data given a L2-ARCTIC audio file.
 
     :param audio: Path to the L2-ARCTIC audio file.
     :return: A list of TimeAlignedType objects.
     """
-    return _load_ta(audio, 'phones')
+    return _load_ta(audio, "phones")
 
 
-def _load_ta(audio: str | os.PathLike[str], tier: str) -> Optional[list[TimeAlignedType]]:
+def _load_ta(audio: str | os.PathLike[str], tier: str) -> list[TimeAlignedType] | None:
     from praatio import textgrid
 
-    file = Path(audio).parent.parent / 'textgrid' / (Path(audio).stem + '.TextGrid')
+    file = Path(audio).parent.parent / "textgrid" / (Path(audio).stem + ".TextGrid")
     if not os.path.exists(file):
         return None
 
@@ -71,7 +70,9 @@ def _load_ta(audio: str | os.PathLike[str], tier: str) -> Optional[list[TimeAlig
     return entries
 
 
-def load_annotations(audio: str | os.PathLike[str]) -> Optional[dict[str, list[TimeAlignedType]]]:
+def load_annotations(
+    audio: str | os.PathLike[str],
+) -> dict[str, list[TimeAlignedType]] | None:
     """Load time-aligned annotation data given a L2-ARCTIC audio file.
 
     :param audio: Path to the L2-ARCTIC audio file.
@@ -79,7 +80,7 @@ def load_annotations(audio: str | os.PathLike[str]) -> Optional[dict[str, list[T
     """
     from praatio import textgrid
 
-    file = Path(audio).parent.parent / 'annotation' / (Path(audio).stem + '.TextGrid')
+    file = Path(audio).parent.parent / "annotation" / (Path(audio).stem + ".TextGrid")
     if not os.path.exists(file):
         return None
 
@@ -96,21 +97,21 @@ def load_annotations(audio: str | os.PathLike[str]) -> Optional[dict[str, list[T
 
 def load_speakers(input_dir: Path) -> dict:
     speakers = {}
-    with open(input_dir / 'readme-download.txt') as file:
+    with open(input_dir / "readme-download.txt") as file:
         processing = False
         for line in file:
-            if not processing and line.startswith('|---|'):
+            if not processing and line.startswith("|---|"):
                 processing = True
                 continue
 
             if processing:
-                if line.startswith('|**Total**|'):
+                if line.startswith("|**Total**|"):
                     break
                 else:
-                    fields = line.strip().split('|')
+                    fields = line.strip().split("|")
                     speaker_id = fields[1]
                     gender = fields[2]
                     dialect = fields[3]
-                    speakers[speaker_id] = {'gender': gender, 'dialect': dialect}
+                    speakers[speaker_id] = {"gender": gender, "dialect": dialect}
 
     return speakers