sonusai 0.15.6__py3-none-any.whl → 0.15.9__py3-none-any.whl

sonusai/mixture/helpers.py

@@ -1,5 +1,9 @@
 from typing import Any
 
+from pyaaware import ForwardTransform
+from pyaaware import InverseTransform
+
+from sonusai.mixture import EnergyT
 from sonusai.mixture.datatypes import AudioF
 from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import AudiosT
@@ -78,7 +82,7 @@ def get_feature_generator_info(fg_config: FeatureGeneratorConfig) -> FeatureGene
         decimation=fg.decimation,
         stride=fg.stride,
         step=fg.step,
-        num_bands=fg.num_bands,
+        feature_parameters=fg.feature_parameters,
         ft_config=TransformConfig(N=fg.ftransform_N,
                                   R=fg.ftransform_R,
                                   bin_start=fg.bin_start,
@@ -327,15 +331,14 @@ def get_ft(mixdb: MixtureDatabase, mixture: Mixture, mixture_audio: AudioT, trut
     import numpy as np
     from pyaaware import FeatureGenerator
 
-    from .spectral_mask import apply_spectral_mask
     from .truth import truth_reduction
 
-    mixture_f = get_mixture_f(mixdb=mixdb, mixture_audio=mixture_audio)
+    mixture_f = get_mixture_f(mixdb=mixdb, mixture=mixture, mixture_audio=mixture_audio)
 
     transform_frames = mixdb.mixture_transform_frames(mixture.samples)
     feature_frames = mixdb.mixture_feature_frames(mixture.samples)
 
-    feature = np.empty((feature_frames, mixdb.fg_stride, mixdb.fg_num_bands), dtype=np.float32)
+    feature = np.empty((feature_frames, mixdb.fg_stride, mixdb.feature_parameters), dtype=np.float32)
     truth_f = np.empty((feature_frames, mixdb.num_classes), dtype=np.complex64)
 
     fg = FeatureGenerator(**asdict(mixdb.fg_config))
@@ -350,11 +353,6 @@ def get_ft(mixdb: MixtureDatabase, mixture: Mixture, mixture_audio: AudioT, trut
             truth_f[feature_frame] = fg.truth()
             feature_frame += 1
 
-    if mixture.spectral_mask_id is not None:
-        feature = apply_spectral_mask(feature=feature,
-                                      spectral_mask=mixdb.spectral_mask(mixture.spectral_mask_id),
-                                      seed=mixture.spectral_mask_seed)
-
     if np.isreal(truth_f).all():
         return feature, truth_f.real
 
@@ -444,14 +442,35 @@ def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: AudiosT)
     return np.sum(targets_ir, axis=0)
 
 
-def get_mixture_f(mixdb: MixtureDatabase, mixture_audio: AudioT) -> AudioF:
+def get_mixture_f(mixdb: MixtureDatabase, mixture: Mixture, mixture_audio: AudioT) -> AudioF:
     """Get the mixture transform for the given mixture
 
     :param mixdb: Mixture database
+    :param mixture: Mixture record
     :param mixture_audio: Mixture audio data for the given mixid
     :return: Mixture transform data
     """
-    return forward_transform(mixture_audio, mixdb.ft_config)
+    from .spectral_mask import apply_spectral_mask
+
+    mixture_f = forward_transform(mixture_audio, mixdb.ft_config)
+
+    if mixture.spectral_mask_id is not None:
+        mixture_f = apply_spectral_mask(audio_f=mixture_f,
+                                        spectral_mask=mixdb.spectral_mask(mixture.spectral_mask_id),
+                                        seed=mixture.spectral_mask_seed)
+
+    return mixture_f
+
+
+def get_transform_from_audio(audio: AudioT, transform: ForwardTransform) -> tuple[AudioF, EnergyT]:
+    """Apply forward transform to input audio data to generate transform data
+
+    :param audio: Time domain data [samples]
+    :param transform: ForwardTransform object
+    :return: Frequency domain data [frames, bins], Energy [frames]
+    """
+    f, e = transform.execute_all(audio)
+    return f.transpose(), e
 
 
 def forward_transform(audio: AudioT, config: TransformConfig) -> AudioF:
@@ -465,17 +484,30 @@ def forward_transform(audio: AudioT, config: TransformConfig) -> AudioF:
     """
     from pyaaware import AawareForwardTransform
 
-    from .audio import calculate_transform_from_audio
-
-    audio_f, _ = calculate_transform_from_audio(audio=audio,
-                                                transform=AawareForwardTransform(N=config.N,
-                                                                                 R=config.R,
-                                                                                 bin_start=config.bin_start,
-                                                                                 bin_end=config.bin_end,
-                                                                                 ttype=config.ttype))
+    audio_f, _ = get_transform_from_audio(audio=audio,
+                                          transform=AawareForwardTransform(N=config.N,
+                                                                           R=config.R,
+                                                                           bin_start=config.bin_start,
+                                                                           bin_end=config.bin_end,
+                                                                           ttype=config.ttype))
     return audio_f
 
 
+def get_audio_from_transform(data: AudioF, transform: InverseTransform, trim: bool = True) -> tuple[AudioT, EnergyT]:
+    """Apply inverse transform to input transform data to generate audio data
+
+    :param data: Frequency domain data [frames, bins]
+    :param transform: InverseTransform object
+    :param trim: Removes starting samples so output waveform will be time-aligned with input waveform to the transform
+    :return: Time domain data [samples], Energy [frames]
+    """
+    t, e = transform.execute_all(data.transpose())
+    if trim:
+        t = t[transform.N - transform.R:]
+
+    return t, e
+
+
 def inverse_transform(transform: AudioF, config: TransformConfig, trim: bool = True) -> AudioT:
     """Transform frequency domain data into time domain using the inverse transform config from the feature
 
@@ -490,16 +522,14 @@ def inverse_transform(transform: AudioF, config: TransformConfig, trim: bool = T
     import numpy as np
     from pyaaware import AawareInverseTransform
 
-    from .audio import calculate_audio_from_transform
-
-    audio, _ = calculate_audio_from_transform(data=transform,
-                                              transform=AawareInverseTransform(N=config.N,
-                                                                               R=config.R,
-                                                                               bin_start=config.bin_start,
-                                                                               bin_end=config.bin_end,
-                                                                               ttype=config.ttype,
-                                                                               gain=np.float32(1)),
-                                              trim=trim)
+    audio, _ = get_audio_from_transform(data=transform,
+                                        transform=AawareInverseTransform(N=config.N,
+                                                                         R=config.R,
+                                                                         bin_start=config.bin_start,
+                                                                         bin_end=config.bin_end,
+                                                                         ttype=config.ttype,
+                                                                         gain=np.float32(1)),
+                                        trim=trim)
     return audio
 
 
@@ -534,7 +564,7 @@ def augmented_target_samples(target_files: TargetFiles,
     it = list(product(*[target_ids, target_augmentation_ids]))
     return sum([estimate_augmented_length_from_length(
         length=target_files[fi].samples,
-        tempo=target_augmentations[ai].tempo,
+        tempo=float(target_augmentations[ai].tempo),
         frame_length=feature_step_samples) for fi, ai, in it])
 
 

sonusai/mixture/log_duration_and_sizes.py

@@ -1,7 +1,7 @@
 def log_duration_and_sizes(total_duration: float,
                            num_classes: int,
                            feature_step_samples: int,
-                           num_bands: int,
+                           feature_parameters: int,
                            stride: int,
                            desc: str) -> None:
     from sonusai import logger
@@ -14,7 +14,7 @@ def log_duration_and_sizes(total_duration: float,
     total_samples = int(total_duration * SAMPLE_RATE)
     mixture_bytes = total_samples * SAMPLE_BYTES
     truth_t_bytes = total_samples * num_classes * FLOAT_BYTES
-    feature_bytes = total_samples / feature_step_samples * stride * num_bands * FLOAT_BYTES
+    feature_bytes = total_samples / feature_step_samples * stride * feature_parameters * FLOAT_BYTES
     truth_f_bytes = total_samples / feature_step_samples * num_classes * FLOAT_BYTES
 
     logger.info('')

sonusai/mixture/mixdb.py

@@ -248,8 +248,8 @@ class MixtureDatabase:
         return self.fg_info.step
 
     @cached_property
-    def fg_num_bands(self) -> int:
-        return self.fg_info.num_bands
+    def feature_parameters(self) -> int:
+        return self.fg_info.feature_parameters
 
     @cached_property
     def ft_config(self) -> TransformConfig:
@@ -809,11 +809,20 @@ class MixtureDatabase:
         :return: Mixture transform data
         """
         from .helpers import forward_transform
+        from .spectral_mask import apply_spectral_mask
 
         if force or mixture is None:
             mixture = self.mixture_mixture(m_id, targets, target, noise, force)
 
-        return forward_transform(mixture, self.ft_config)
+        mixture_f = forward_transform(mixture, self.ft_config)
+
+        m = self.mixture(m_id)
+        if m.spectral_mask_id is not None:
+            mixture_f = apply_spectral_mask(audio_f=mixture_f,
+                                            spectral_mask=self.spectral_mask(int(m.spectral_mask_id)),
+                                            seed=m.spectral_mask_seed)
+
+        return mixture_f
 
     def mixture_truth_t(self,
                         m_id: int,
@@ -938,7 +947,6 @@ class MixtureDatabase:
         import numpy as np
         from pyaaware import FeatureGenerator
 
-        from .spectral_mask import apply_spectral_mask
         from .truth import truth_reduction
 
         if not force:
@@ -964,7 +972,7 @@ class MixtureDatabase:
         if truth_t is None:
             truth_t = np.zeros((m.samples, self.num_classes), dtype=np.float32)
 
-        feature = np.empty((feature_frames, self.fg_stride, self.fg_num_bands), dtype=np.float32)
+        feature = np.empty((feature_frames, self.fg_stride, self.feature_parameters), dtype=np.float32)
         truth_f = np.empty((feature_frames, self.num_classes), dtype=np.complex64)
 
         fg = FeatureGenerator(**asdict(self.fg_config))
@@ -979,11 +987,6 @@ class MixtureDatabase:
                 truth_f[feature_frame] = fg.truth()
                 feature_frame += 1
 
-        if m.spectral_mask_id is not None:
-            feature = apply_spectral_mask(feature=feature,
-                                          spectral_mask=self.spectral_mask(int(m.spectral_mask_id)),
-                                          seed=m.spectral_mask_seed)
-
         if np.isreal(truth_f).all():
             return feature, truth_f.real
 

sonusai/mixture/spectral_mask.py

@@ -1,23 +1,23 @@
-from sonusai.mixture.datatypes import Feature
+from sonusai.mixture.datatypes import AudioF
 from sonusai.mixture.datatypes import SpectralMask
 
 
-def apply_spectral_mask(feature: Feature, spectral_mask: SpectralMask, seed: int = None) -> Feature:
+def apply_spectral_mask(audio_f: AudioF, spectral_mask: SpectralMask, seed: int = None) -> AudioF:
     """Apply frequency and time masking
 
     Implementation of SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition
 
     Ref: https://arxiv.org/pdf/1904.08779.pdf
 
-    f_width consecutive bands [f_start, f_start + f_width) are masked, where f_width is chosen from a uniform
-    distribution from 0 to the f_max_width, and f_start is chosen from [0, bands - f_width).
+    f_width consecutive bins [f_start, f_start + f_width) are masked, where f_width is chosen from a uniform
+    distribution from 0 to the f_max_width, and f_start is chosen from [0, bins - f_width).
 
     t_width consecutive frames [t_start, t_start + t_width) are masked, where t_width is chosen from a uniform
     distribution from 0 to the t_max_width, and t_start is chosen from [0, frames - t_width).
 
     A time mask cannot be wider than t_max_percent times the number of frames.
 
-    :param feature: Numpy array of feature data [frames, strides, bands]
+    :param audio_f: Numpy array of transform audio data [frames, bins]
     :param spectral_mask: Spectral mask parameters
     :param seed: Random number seed
     :return: Augmented feature
@@ -26,28 +26,28 @@ def apply_spectral_mask(feature: Feature, spectral_mask: SpectralMask, seed: int
 
     from sonusai import SonusAIError
 
-    if feature.ndim != 3:
-        raise SonusAIError('feature input must have three dimensions [frames, strides, bands]')
+    if audio_f.ndim != 2:
+        raise SonusAIError('feature input must have three dimensions [frames, bins]')
 
-    frames, strides, bands = feature.shape
+    frames, bins = audio_f.shape
 
     f_max_width = spectral_mask.f_max_width
-    if f_max_width not in range(0, bands + 1):
-        f_max_width = bands
+    if f_max_width not in range(0, bins + 1):
+        f_max_width = bins
 
     rng = np.random.default_rng(seed)
 
     # apply f_num frequency masks to the feature
     for _ in range(spectral_mask.f_num):
         f_width = int(rng.uniform(0, f_max_width))
-        f_start = rng.integers(0, bands - f_width, endpoint=True)
-        feature[:, :, f_start:f_start + f_width] = 0
+        f_start = rng.integers(0, bins - f_width, endpoint=True)
+        audio_f[:, f_start:f_start + f_width] = 0
 
     # apply t_num time masks to the feature
     t_upper_bound = int(spectral_mask.t_max_percent / 100 * frames)
     for _ in range(spectral_mask.t_num):
         t_width = min(int(rng.uniform(0, spectral_mask.t_max_width)), t_upper_bound)
         t_start = rng.integers(0, frames - t_width, endpoint=True)
-        feature[t_start:t_start + t_width, :, :] = 0
+        audio_f[t_start:t_start + t_width, :] = 0
 
-    return feature
+    return audio_f

sonusai/mixture/truth_functions/data.py

@@ -23,6 +23,8 @@ class Data:
                               num_classes=config.num_classes,
                               truth_mutex=config.mutex)
 
+        self.feature_parameters = fg.feature_parameters
+        self.ttype = fg.ftransform_ttype
         self.frame_size = fg.ftransform_R
 
         if len(target_audio) % self.frame_size != 0:

sonusai/mixture/truth_functions/target.py

@@ -13,12 +13,13 @@ 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 * bins] (stacked real, imag)
+Output shape: [:, num_classes]
+                    (target stacked real, imag; or real only for tdac-co)
     """
 
     from sonusai import SonusAIError
 
-    if data.config.num_classes != 2 * data.target_fft.bins:
+    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))
@@ -28,6 +29,7 @@ Output shape: [:, 2 * bins] (stacked real, imag)
                                       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)
 
@@ -43,11 +45,13 @@ using the STFT configuration defined by the feature. This
 will include a forward transform window if defined by the
 feature.
 
-Output shape: [:, 4 * bins] (target stacked real, imag; mixture stacked real, imag)
+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)
     """
     from sonusai import SonusAIError
 
-    if data.config.num_classes != 2 * data.target_fft.bins + 2 * data.mixture_fft.bins:
+    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 = _execute_fft(data.target_audio, data.target_fft, len(data.offsets))
@@ -59,6 +63,7 @@ Output shape: [:, 4 * bins] (target stacked real, imag; mixture stacked real, im
                                       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)
 
@@ -67,6 +72,7 @@ Output shape: [:, 4 * bins] (target stacked real, imag; mixture stacked real, im
                                       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)
 
@@ -125,6 +131,7 @@ def _stack_real_imag(data: AudioF,
                      frame_size: int,
                      zero_based_indices: list[int],
                      bins: int,
+                     ttype: str,
                      start: int,
                      truth: Truth) -> Truth:
     import numpy as np
@@ -134,7 +141,8 @@ def _stack_real_imag(data: AudioF,
         b = _get_bin_slice(index + start, bins)
         truth[i, b] = np.real(data)
 
-        b = _get_bin_slice(b.stop, bins)
-        truth[i, b] = np.imag(data)
+        if ttype != 'tdac-co':
+            b = _get_bin_slice(b.stop, bins)
+            truth[i, b] = np.imag(data)
 
     return truth

sonusai/onnx_predict.py

@@ -105,7 +105,7 @@ def main() -> None:
         logger.info('')
         logger.info(f'Run prediction on {input_name}')
         audio = read_audio(input_name)
-        feature = get_feature_from_audio(audio=audio, feature=model_metadata.feature)
+        feature = get_feature_from_audio(audio=audio, feature_mode=model_metadata.feature)
 
         predict = pad_and_predict(feature=feature,
                                   model_name=model_name,

sonusai/plot.py

@@ -314,7 +314,7 @@ def main() -> None:
             raise SonusAIError('Must specify MODEL when input is WAV')
 
         mixture_audio = read_audio(input_name)
-        feature = get_feature_from_audio(audio=mixture_audio, feature=model.feature)
+        feature = get_feature_from_audio(audio=mixture_audio, feature_mode=model.feature)
         fg_config = FeatureGeneratorConfig(feature_mode=model.feature,
                                            num_classes=model.output_shape[-1],
                                            truth_mutex=False)
@@ -406,11 +406,11 @@ def main() -> None:
         title = f'{input_name}'
         pdf_name = f'{base_name}-plot.pdf'
 
-    # Original size [frames, stride, num_bands]
+    # Original size [frames, stride, feature_parameters]
     # Decimate in the stride dimension
-    # Reshape to get frames*decimated_stride, num_bands
+    # Reshape to get frames*decimated_stride, feature_parameters
     if feature.ndim != 3:
-        raise SonusAIError(f'feature does not have 3 dimensions: frames, stride, num_bands')
+        raise SonusAIError(f'feature does not have 3 dimensions: frames, stride, feature_parameters')
     spectrogram = feature[:, -fg_step:, :]
     spectrogram = np.reshape(spectrogram, (spectrogram.shape[0] * spectrogram.shape[1], spectrogram.shape[2]))
 

sonusai/post_spenh_targetf.py

@@ -123,7 +123,7 @@ def _process(file: str) -> None:
     from pyaaware import AawareInverseTransform
 
     from sonusai import SonusAIError
-    from sonusai.mixture import calculate_audio_from_transform
+    from sonusai.mixture import get_audio_from_transform
     from sonusai.utils import float_to_int16
     from sonusai.utils import unstack_complex
     from sonusai.utils import write_wav
@@ -135,13 +135,13 @@ def _process(file: str) -> None:
         raise SonusAIError(f'Error reading {file}: {e}')
 
     output_name = join(MP_GLOBAL.output_dir, splitext(basename(file))[0] + '.wav')
-    audio, _ = calculate_audio_from_transform(data=predict,
-                                              transform=AawareInverseTransform(N=MP_GLOBAL.N,
-                                                                               R=MP_GLOBAL.R,
-                                                                               bin_start=MP_GLOBAL.bin_start,
-                                                                               bin_end=MP_GLOBAL.bin_end,
-                                                                               ttype=MP_GLOBAL.ttype,
-                                                                               gain=np.float32(1)))
+    audio, _ = get_audio_from_transform(data=predict,
+                                        transform=AawareInverseTransform(N=MP_GLOBAL.N,
+                                                                         R=MP_GLOBAL.R,
+                                                                         bin_start=MP_GLOBAL.bin_start,
+                                                                         bin_end=MP_GLOBAL.bin_end,
+                                                                         ttype=MP_GLOBAL.ttype,
+                                                                         gain=np.float32(1)))
     write_wav(name=output_name, audio=float_to_int16(audio))