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

sonusai/mixture/truth_functions/sed.py

@@ -1,65 +1,91 @@
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.data import Data
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
 
 
-def sed(data: Data) -> Truth:
+def _strictly_decreasing(list_to_check: list) -> bool:
+    from itertools import pairwise
+
+    return all(x > y for x, y in pairwise(list_to_check))
+
+
+def sed_validate(config: dict) -> None:
+    if len(config) == 0:
+        raise AttributeError("sed truth function is missing config")
+
+    parameters = ["thresholds"]
+    for parameter in parameters:
+        if parameter not in config:
+            raise AttributeError(f"sed truth function is missing required '{parameter}'")
+
+    thresholds = config["thresholds"]
+    if not _strictly_decreasing(thresholds):
+        raise ValueError(f"sed truth function 'thresholds' are not strictly decreasing: {thresholds}")
+
+
+def sed_parameters(config: TruthFunctionConfig) -> int:
+    return config.num_classes
+
+
+def sed(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
     """Sound energy detection truth generation function
 
-Calculates sound energy detection truth using simple 3 threshold
-hysteresis algorithm. SED outputs 3 possible probabilities of
-sound presence: 1.0 present, 0.5 (transition/uncertain), 0 not
-present. The output values will be assigned to the truth output
-at the index specified in the truth_settings: index.
+    Calculates sound energy detection truth using simple 3 threshold
+    hysteresis algorithm. SED outputs 3 possible probabilities of
+    sound presence: 1.0 present, 0.5 (transition/uncertain), 0 not
+    present. The output values will be assigned to the truth output
+    at the index specified in the config.
+
+    Output shape: [:, num_classes]
+
+    index       Truth index <int> or list(<int>)
 
-Output shape: [:, num_classes]
+    index indicates which truth fields should be set.
+    0 indicates none, 1 is first element in truth output vector, 2 2nd element, etc.
 
-For multilabel classification applications, num_classes should be
-set to the number of sounds/classes to be detected.
+                Examples:
+                  index = 5       truth in class 5, truth(4, 1)
+                  index = [1, 5]  truth in classes 1 and 5, truth([0, 4], 1)
 
-For single-label classification, where truth_mutex=1, num_classes
-should be set to the number of sounds/classes to be detected + 1 for
-the other class.
+                In mutually-exclusive mode, a frame is expected to only
+                belong to one class and thus all probabilities must sum to
+                1. This is effectively truth for a classifier with multichannel
+                softmax output.
+
+                For multi-label classification each class is an individual
+                probability for that class and any given frame can be
+                assigned to multiple classes/labels, i.e., the classes are
+                not mutually-exclusive. For example, a NN classifier with
+                multichannel sigmoid output. In this case, index could
+                also be a vector with multiple class indices.
     """
     import numpy as np
     import torch
     from pyaaware import SED
 
-    from sonusai import SonusAIError
-
-    if data.config.config is None:
-        raise SonusAIError('Truth function SED missing config')
+    if len(data.target_audio) % config.frame_size != 0:
+        raise ValueError(f"Number of samples in audio is not a multiple of {config.frame_size}")
 
-    parameters = ['thresholds']
-    for parameter in parameters:
-        if 'thresholds' not in data.config.config:
-            raise SonusAIError(f'Truth function SED config missing required parameter: {parameter}')
-
-    thresholds = data.config.config['thresholds']
-    if not _strictly_decreasing(thresholds):
-        raise SonusAIError(f'Truth function SED thresholds are not strictly decreasing: {thresholds}')
-
-    if len(data.target_audio) % data.frame_size != 0:
-        raise SonusAIError(f'Number of samples in audio is not a multiple of {data.frame_size}')
+    frames = config.target_fft.frames(data.target_audio)
+    parameters = sed_parameters(config)
+    if config.target_gain == 0:
+        return np.zeros((frames, parameters), dtype=np.float32)
 
     # SED wants 1-based indices
-    s = SED(thresholds=thresholds,
-            index=data.config.index,
-            frame_size=data.frame_size,
-            num_classes=data.config.num_classes,
-            mutex=data.config.mutex)
+    s = SED(
+        thresholds=config.config["thresholds"],
+        index=config.class_indices,
+        frame_size=config.frame_size,
+        num_classes=config.num_classes,
+    )
 
-    target_audio = data.target_audio / data.config.target_gain
-    energy_t = data.target_fft.execute_all(torch.from_numpy(target_audio))[1].numpy()
-    if len(energy_t) != len(data.offsets):
-        raise SonusAIError(f'Number of frames in energy_t, {len(energy_t)},'
-                           f' is not number of frames in truth, {len(data.offsets)}')
+    # Back out target gain
+    target_audio = data.target_audio / config.target_gain
 
-    for idx, offset in enumerate(data.offsets):
-        new_truth = s.execute(energy_t[idx])
-        data.truth[offset:offset + data.frame_size] = np.reshape(new_truth, (1, len(new_truth)))
+    # Compute energy
+    target_energy = config.target_fft.execute_all(torch.from_numpy(target_audio))[1].numpy()
 
-    return data.truth
+    if frames != target_energy.shape[0]:
+        raise ValueError("Incorrect frames calculation in sed truth function")
 
-
-def _strictly_decreasing(list_to_check: list) -> bool:
-    return all(x > y for x, y in zip(list_to_check, list_to_check[1:]))
+    return s.execute_all(target_energy)

sonusai/mixture/truth_functions/target.py

@@ -1,146 +1,109 @@
-from sonusai import ForwardTransform
-
 from sonusai.mixture.datatypes import AudioF
-from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import Truth
-from sonusai.mixture.truth_functions.data import Data
-
-
-def target_f(data: Data) -> Truth:
-    """Frequency domain target truth function
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionConfig
+from sonusai.mixture.truth_functions.datatypes import TruthFunctionData
 
-Calculates the true transform of the target using the STFT
-configuration defined by the feature. This will include a
-forward transform window if defined by the feature.
 
-Output shape: [:, num_classes]
-                    (target stacked real, imag; or real only for tdac-co)
-    """
-    from sonusai import SonusAIError
+def target_f_validate(_config: dict) -> None:
+    pass
 
-    if data.config.num_classes != data.feature_parameters:
-        raise SonusAIError(f'Invalid num_classes for target_f truth: {data.config.num_classes}')
 
-    target_freq = _execute_fft(data.target_audio, data.target_fft, len(data.offsets))
-    for idx, offset in enumerate(data.offsets):
-        data.truth = _stack_real_imag(data=target_freq[idx],
-                                      offset=offset,
-                                      frame_size=data.frame_size,
-                                      zero_based_indices=data.zero_based_indices,
-                                      bins=data.target_fft.bins,
-                                      ttype=data.ttype,
-                                      start=0,
-                                      truth=data.truth)
+def target_f_parameters(config: TruthFunctionConfig) -> int:
+    if config.ttype == "tdac-co":
+        return config.target_fft.bins
 
-    return data.truth
+    return config.target_fft.bins * 2
 
 
-# TODO: Need Data to include mixture audio to do this correctly
-def target_mixture_f(data: Data) -> Truth:
-    """Frequency domain target and mixture truth function
+def target_f(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
+    """Frequency domain target truth function
 
-Calculates the true transform of the target and the mixture
-using the STFT configuration defined by the feature. This
-will include a forward transform window if defined by the
-feature.
+    Calculates the true transform of the target using the STFT
+    configuration defined by the feature. This will include a
+    forward transform window if defined by the feature.
 
-Output shape: [:, 2 * num_classes]
-                    (target stacked real, imag; or real only for tdac-co)
-                    (mixture stacked real, imag; or real only for tdac-co)
+    Output shape: [:, 2 * bins] (target stacked real, imag) or
+                  [:, bins] (target real only for tdac-co)
     """
-    from sonusai import SonusAIError
+    import torch
 
-    if data.config.num_classes != 2 * data.feature_parameters:
-        raise SonusAIError(f'Invalid num_classes for target_mixture_f truth: {data.config.num_classes}')
+    target_freq = config.target_fft.execute_all(torch.from_numpy(data.target_audio))[0].numpy()
+    return _stack_real_imag(target_freq, config.ttype)
 
-    target_freq = _execute_fft(data.target_audio, data.target_fft, len(data.offsets))
-    mixture_freq = _execute_fft(data.mixture_audio, data.mixture_fft, len(data.offsets))
 
-    for idx, offset in enumerate(data.offsets):
-        data.truth = _stack_real_imag(data=target_freq[idx],
-                                      offset=offset,
-                                      frame_size=data.frame_size,
-                                      zero_based_indices=data.zero_based_indices,
-                                      bins=data.target_fft.bins,
-                                      ttype=data.ttype,
-                                      start=0,
-                                      truth=data.truth)
+def target_mixture_f_validate(_config: dict) -> None:
+    pass
 
-        data.truth = _stack_real_imag(data=mixture_freq[idx],
-                                      offset=offset,
-                                      frame_size=data.frame_size,
-                                      zero_based_indices=data.zero_based_indices,
-                                      bins=data.target_fft.bins,
-                                      ttype=data.ttype,
-                                      start=data.target_fft.bins * 2,
-                                      truth=data.truth)
 
-    return data.truth
+def target_mixture_f_parameters(config: TruthFunctionConfig) -> int:
+    if config.ttype == "tdac-co":
+        return config.target_fft.bins * 2
 
+    return config.target_fft.bins * 4
 
-def target_swin_f(data: Data) -> Truth:
-    """Frequency domain target with synthesis window truth function
 
-Calculates the true transform of the target using the STFT
-configuration defined by the feature. This will include a
-forward transform window if defined by the feature and also
-the inverse transform (or synthesis) window.
+def target_mixture_f(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
+    """Frequency domain target and mixture truth function
+
+    Calculates the true transform of the target and the mixture
+    using the STFT configuration defined by the feature. This
+    will include a forward transform window if defined by the
+    feature.
 
-Output shape: [:, 2 * bins] (stacked real, imag)
+    Output shape: [:, 4 * bins] (target stacked real, imag; mixture stacked real, imag) or
+                  [:, 2 * bins] (target real; mixture real for tdac-co)
     """
     import numpy as np
+    import torch
 
-    from sonusai import SonusAIError
+    target_freq = config.target_fft.execute_all(torch.from_numpy(data.target_audio))[0].numpy()
+    mixture_freq = config.mixture_fft.execute_all(torch.from_numpy(data.mixture_audio))[0].numpy()
 
-    if data.config.num_classes != 2 * data.target_fft.bins:
-        raise SonusAIError(f'Invalid num_classes for target_swin_f truth: {data.config.num_classes}')
+    frames, bins = target_freq.shape
+    truth = np.empty((frames, bins * 4), dtype=np.float32)
+    truth[:, : bins * 2] = _stack_real_imag(target_freq, config.ttype)
+    truth[:, bins * 2 :] = _stack_real_imag(mixture_freq, config.ttype)
+    return truth
 
-    for idx, offset in enumerate(data.offsets):
-        target_freq, _ = data.target_fft.execute(
-            np.multiply(data.target_audio[offset:offset + data.frame_size], data.swin))
 
-        indices = slice(offset, offset + data.frame_size)
-        for index in data.zero_based_indices:
-            bins = _get_bin_slice(index, data.target_fft.bins)
-            data.truth[indices, bins] = np.real(target_freq[idx])
+def target_swin_f_validate(_config: dict) -> None:
+    pass
 
-            bins = _get_bin_slice(bins.stop, data.target_fft.bins)
-            data.truth[indices, bins] = np.imag(target_freq[idx])
 
-    return data.truth
+def target_swin_f_parameters(config: TruthFunctionConfig) -> int:
+    return config.target_fft.bins * 2
 
 
-def _execute_fft(audio: AudioT, transform: ForwardTransform, expected_frames: int) -> AudioF:
-    import torch
-    from sonusai import SonusAIError
+def target_swin_f(data: TruthFunctionData, config: TruthFunctionConfig) -> Truth:
+    """Frequency domain target with synthesis window truth function
+
+    Calculates the true transform of the target using the STFT
+    configuration defined by the feature. This will include a
+    forward transform window if defined by the feature and also
+    the inverse transform (or synthesis) window.
 
-    freq = transform.execute_all(torch.from_numpy(audio))[0].numpy()
-    if len(freq) != expected_frames:
-        raise SonusAIError(f'Number of frames, {len(freq)}, is not number of frames expected, {expected_frames}')
-    return freq
+    Output shape: [:, 2 * bins] (stacked real, imag)
+    """
+    import numpy as np
 
+    from sonusai.utils import stack_complex
 
-def _get_bin_slice(start: int, length: int) -> slice:
-    return slice(start, start + length)
+    truth = np.empty((len(data.target_audio) // config.frame_size, config.target_fft.bins * 2), dtype=np.float32)
+    for idx, offset in enumerate(range(0, len(data.target_audio), config.frame_size)):
+        target_freq = config.target_fft.execute(
+            np.multiply(data.target_audio[offset : offset + config.frame_size], config.swin)
+        )[0]
+        truth[idx] = stack_complex(target_freq)
 
+    return truth
 
-def _stack_real_imag(data: AudioF,
-                     offset: int,
-                     frame_size: int,
-                     zero_based_indices: list[int],
-                     bins: int,
-                     ttype: str,
-                     start: int,
-                     truth: Truth) -> Truth:
+
+def _stack_real_imag(data: AudioF, ttype: str) -> Truth:
     import numpy as np
 
-    i = _get_bin_slice(offset, frame_size)
-    for index in zero_based_indices:
-        b = _get_bin_slice(index + start, bins)
-        truth[i, b] = np.real(data)
+    from sonusai.utils import stack_complex
 
-        if ttype != 'tdac-co':
-            b = _get_bin_slice(b.stop, bins)
-            truth[i, b] = np.imag(data)
+    if ttype == "tdac-co":
+        return np.real(data)
 
-    return truth
+    return stack_complex(data)

sonusai/mkwav.py

@@ -16,18 +16,16 @@ Inputs:
     MIXID       A glob of mixture ID(s) to generate.
 
 Outputs the following to the mixture database directory:
-    <id>_mixture.wav:   mixture
-    <id>_target.wav:    target (optional)
-    <id>_noise.wav:     noise (optional)
-    <id>.txt
+    <id>
+        mixture.wav:   mixture
+        target.wav:    target (optional)
+        noise.wav:     noise (optional)
+        metadata.txt
     mkwav.log
 
 """
-import signal
-from dataclasses import dataclass
 
-from sonusai.mixture import AudioT
-from sonusai.mixture import MixtureDatabase
+import signal
 
 
 def signal_handler(_sig, _frame):
@@ -35,79 +33,33 @@ 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)
 
 
 signal.signal(signal.SIGINT, signal_handler)
 
 
-@dataclass
-class MPGlobal:
-    mixdb: MixtureDatabase = None
-    write_target: bool = None
-    write_noise: bool = None
-
-
-MP_GLOBAL = MPGlobal()
-
-
-def mkwav(location: str, mixid: int) -> tuple[AudioT, AudioT, AudioT]:
-    import numpy as np
-
-    from sonusai.genmix import genmix
-
-    data = genmix(location=location, mixids=mixid, force=False)
-
-    return data[0].mixture, np.sum(data[0].targets, axis=0), data[0].noise
-
-
-def _process_mixture(mixid: int) -> None:
-    from os.path import exists
+def _process_mixture(m_id: int, location: str, write_target: bool, write_noise: bool) -> None:
     from os.path import join
-    from os.path import splitext
-
-    import h5py
-    import numpy as np
 
-    from sonusai.mixture import mixture_metadata
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.mixture import write_mixture_metadata
     from sonusai.utils import float_to_int16
     from sonusai.utils import write_audio
 
-    mixture_filename = join(MP_GLOBAL.mixdb.location, MP_GLOBAL.mixdb.mixtures[mixid].name)
-    mixture_basename = splitext(mixture_filename)[0]
-
-    target = None
-    noise = None
-
-    need_data = True
-    if exists(mixture_filename + '.h5'):
-        with h5py.File(mixture_filename, 'r') as f:
-            if 'mixture' in f:
-                need_data = False
-            if MP_GLOBAL.write_target and 'targets' not in f:
-                need_data = True
-            if MP_GLOBAL.write_noise and 'noise' not in f:
-                need_data = True
-
-    if need_data:
-        mixture, target, noise = mkwav(location=MP_GLOBAL.mixdb.location, mixid=mixid)
-    else:
-        with h5py.File(mixture_filename, 'r') as f:
-            mixture = np.array(f['mixture'])
-            if MP_GLOBAL.write_target:
-                target = np.sum(np.array(f['targets']), axis=0)
-            if MP_GLOBAL.write_noise:
-                noise = np.array(f['noise'])
-
-    write_audio(name=mixture_basename + '_mixture.wav', audio=float_to_int16(mixture))
-    if MP_GLOBAL.write_target:
-        write_audio(name=mixture_basename + '_target.wav', audio=float_to_int16(target))
-    if MP_GLOBAL.write_noise:
-        write_audio(name=mixture_basename + '_noise.wav', audio=float_to_int16(noise))
-
-    with open(file=mixture_basename + '.txt', mode='w') as f:
-        f.write(mixture_metadata(MP_GLOBAL.mixdb, MP_GLOBAL.mixdb.mixture(mixid)))
+    mixdb = MixtureDatabase(location)
+
+    mixture = mixdb.mixture(m_id)
+    location = join(mixdb.location, mixture.name)
+
+    write_audio(name=join(location, "mixture.wav"), audio=float_to_int16(mixdb.mixture_mixture(m_id)))
+    if write_target:
+        write_audio(name=join(location, "target.wav"), audio=float_to_int16(mixdb.mixture_target(m_id)))
+    if write_noise:
+        write_audio(name=join(location, "noise.wav"), audio=float_to_int16(mixdb.mixture_noise(m_id)))
+
+    write_mixture_metadata(mixdb, mixture)
 
 
 def main() -> None:
@@ -118,63 +70,68 @@ def main() -> None:
 
     args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
 
-    verbose = args['--verbose']
-    mixid = args['--mixid']
-    MP_GLOBAL.write_target = args['--target']
-    MP_GLOBAL.write_noise = args['--noise']
-    location = args['LOC']
+    verbose = args["--verbose"]
+    mixid = args["--mixid"]
+    write_target = args["--target"]
+    write_noise = args["--noise"]
+    location = args["LOC"]
 
     import time
+    from functools import partial
     from os.path import join
 
-    from tqdm import tqdm
-
     import sonusai
     from sonusai import create_file_handler
     from sonusai import initial_log_messages
     from sonusai import logger
     from sonusai import update_console_handler
+    from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import check_audio_files_exist
-    from sonusai.utils import pp_tqdm_imap
     from sonusai.utils import human_readable_size
+    from sonusai.utils import par_track
     from sonusai.utils import seconds_to_hms
+    from sonusai.utils import track
 
     start_time = time.monotonic()
 
-    create_file_handler(join(location, 'mkwav.log'))
+    create_file_handler(join(location, "mkwav.log"))
     update_console_handler(verbose)
-    initial_log_messages('mkwav')
+    initial_log_messages("mkwav")
 
-    logger.info(f'Load mixture database from {location}')
-    MP_GLOBAL.mixdb = MixtureDatabase(location)
-    mixid = MP_GLOBAL.mixdb.mixids_to_list(mixid)
+    logger.info(f"Load mixture database from {location}")
+    mixdb = MixtureDatabase(location)
+    mixid = mixdb.mixids_to_list(mixid)
 
-    total_samples = MP_GLOBAL.mixdb.total_samples(mixid)
+    total_samples = mixdb.total_samples(mixid)
     duration = total_samples / sonusai.mixture.SAMPLE_RATE
 
-    logger.info('')
-    logger.info(f'Found {len(mixid):,} mixtures to process')
-    logger.info(f'{total_samples:,} samples')
+    logger.info("")
+    logger.info(f"Found {len(mixid):,} mixtures to process")
+    logger.info(f"{total_samples:,} samples")
 
-    check_audio_files_exist(MP_GLOBAL.mixdb)
+    check_audio_files_exist(mixdb)
 
-    progress = tqdm(total=len(mixid))
-    pp_tqdm_imap(_process_mixture, mixid, progress=progress)
+    progress = track(total=len(mixid))
+    par_track(
+        partial(_process_mixture, location=location, write_target=write_target, write_noise=write_noise),
+        mixid,
+        progress=progress,
+    )
     progress.close()
 
-    logger.info(f'Wrote {len(mixid)} 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 MP_GLOBAL.write_target:
-        logger.info(f'target:   {human_readable_size(total_samples * 2, 1)}')
-    if MP_GLOBAL.write_noise:
-        logger.info(f'noise:    {human_readable_size(total_samples * 2, 1)}')
+    logger.info(f"Wrote {len(mixid)} 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 write_target:
+        logger.info(f"target:   {human_readable_size(total_samples * 2, 1)}")
+    if write_noise:
+        logger.info(f"noise:    {human_readable_size(total_samples * 2, 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()