sonusai 0.15.8__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,7 +23,7 @@ class Data:
                               num_classes=config.num_classes,
                               truth_mutex=config.mutex)
 
-        self.num_bands = fg.num_bands
+        self.feature_parameters = fg.feature_parameters
         self.ttype = fg.ftransform_ttype
         self.frame_size = fg.ftransform_R
 

sonusai/mixture/truth_functions/target.py

@@ -19,7 +19,7 @@ Output shape: [:, num_classes]
 
     from sonusai import SonusAIError
 
-    if data.config.num_classes != data.num_bands:
+    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))
@@ -51,7 +51,7 @@ Output shape: [:, 2 * num_classes]
     """
     from sonusai import SonusAIError
 
-    if data.config.num_classes != 2 * data.num_bands:
+    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))

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))
 
 

sonusai/torchl_predict.py

@@ -43,15 +43,38 @@ Outputs the following to tpredict-<TIMESTAMP> directory:
     torch_predict.log
 
 """
+from os import makedirs
+from os.path import basename
+from os.path import isdir
 from os.path import join
+from os.path import normpath
+from os.path import splitext
 from typing import Any
 
 import h5py
 import torch
+from docopt import docopt
+from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import BasePredictionWriter
+from pyaaware import FeatureGenerator
+from pyaaware import TorchInverseTransform
+from torchinfo import summary
 
+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.data_generator import TorchFromMixtureDatabase
 from sonusai.mixture import Feature
+from sonusai.mixture import MixtureDatabase
+from sonusai.mixture import get_audio_from_feature
+from sonusai.mixture import get_feature_from_audio
+from sonusai.mixture import read_audio
+from sonusai.utils import create_ts_name
+from sonusai.utils import import_keras_model
+from sonusai.utils import trim_docstring
+from sonusai.utils import write_wav
 
 
 class CustomWriter(BasePredictionWriter):
@@ -61,7 +84,7 @@ class CustomWriter(BasePredictionWriter):
 
     def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
         # this will create N (num processes) files in `output_dir` each containing
-        # the predictions of it's respective rank
+        # the predictions of its respective rank
         # torch.save(predictions, os.path.join(self.output_dir, f"predictions_{trainer.global_rank}.pt"))
 
         # optionally, you can also save `batch_indices` to get the information about the data index
@@ -119,11 +142,6 @@ def power_uncompress(real, imag):
 
 
 def main() -> None:
-    from docopt import docopt
-
-    import sonusai
-    from sonusai.utils import trim_docstring
-
     args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
 
     verbose = args['--verbose']
@@ -139,27 +157,6 @@ def main() -> None:
     wavdbg = args['--wavdbg']  # write .wav if true
     input_name = args['INPUT']
 
-    from os import makedirs
-    from os.path import basename
-    from os.path import isdir
-    from os.path import isfile
-    from os.path import join
-    from os.path import splitext
-    from os.path import normpath
-    import h5py
-    # from sonusai.utils import float_to_int16
-
-    from torchinfo import summary
-    from sonusai import create_file_handler
-    from sonusai import initial_log_messages
-    from sonusai import update_console_handler
-    from sonusai.mixture import MixtureDatabase
-    from sonusai.mixture import get_feature_from_audio
-    from sonusai.utils import import_keras_model
-    from sonusai.mixture import read_audio
-    from sonusai.utils import create_ts_name
-    from sonusai.data_generator import TorchFromMixtureDatabase
-
     if batch_size is not None:
         batch_size = int(batch_size)
     if batch_size != 1:
@@ -222,6 +219,8 @@ def main() -> None:
                 hparams['timesteps'] = timesteps
 
         logger.info(f'Building model with hparams and batch_size={batch_size}, timesteps={timesteps}')
+        # hparams['cl_per_wght'] = 0.0
+        # hparams['feature'] = 'hum00ns1'
         try:
             model = litemodule.MyHyperModel(**hparams)  # use hparams
             # litemodule.MyHyperModel.load_from_checkpoint(ckpt_name, **hparams)
@@ -303,33 +302,25 @@ def main() -> None:
                                              drop_last=False,
                                              num_workers=dlcpu)
 
-        if wavdbg:  # setup for wav write if enabled
-            # Info needed to setup inverse transform
-            from pyaaware import FeatureGenerator
-            from pyaaware import TorchInverseTransform
-            from torchaudio import save
-            # from sonusai.utils import write_wav
-
-            half = model.num_classes // 2
-            fg = FeatureGenerator(feature_mode=model.hparams.feature,
-                                  num_classes=model.num_classes,
-                                  truth_mutex=model.truth_mutex)
-            itf = TorchInverseTransform(N=fg.itransform_N,
-                                        R=fg.itransform_R,
-                                        bin_start=fg.bin_start,
-                                        bin_end=fg.bin_end,
-                                        ttype=fg.itransform_ttype)
-
-            if mixdb.target_files[0].truth_settings[0].function == 'target_f' or \
-                    mixdb.target_files[0].truth_settings[0].function == 'target_mixture_f':
-                enable_truth_wav = True
-            else:
-                enable_truth_wav = False
-
+        # Info needed to set up inverse transform
+        half = model.num_classes // 2
+        fg = FeatureGenerator(feature_mode=model.hparams.feature,
+                              num_classes=model.num_classes,
+                              truth_mutex=model.truth_mutex)
+        itf = TorchInverseTransform(N=fg.itransform_N,
+                                    R=fg.itransform_R,
+                                    bin_start=fg.bin_start,
+                                    bin_end=fg.bin_end,
+                                    ttype=fg.itransform_ttype)
+
+        enable_truth_wav = False
+        enable_mix_wav = False
+        if wavdbg:
             if mixdb.target_files[0].truth_settings[0].function == 'target_mixture_f':
                 enable_mix_wav = True
-            else:
-                enable_mix_wav = False
+                enable_truth_wav = True
+            elif mixdb.target_files[0].truth_settings[0].function == 'target_f':
+                enable_truth_wav = True
 
         if reset:
             logger.info(f'Running {mixdb.num_mixtures} mixtures individually with model reset ...')
@@ -351,26 +342,25 @@ def main() -> None:
                 if wavdbg:
                     owav_base = splitext(output_name)[0]
                     tmp = torch.complex(ypred[..., :half], ypred[..., half:]).permute(2, 0, 1).detach()
+                    itf.reset()
                     predwav, _ = itf.execute_all(tmp)
-                    # predwav, _ = calculate_audio_from_transform(tmp, itf, trim=True)
-                    save(owav_base + '.wav', predwav.permute([1, 0]), 16000, encoding='PCM_S', bits_per_sample=16)
+                    # predwav, _ = calculate_audio_from_transform(tmp.numpy(), itf, trim=True)
+                    write_wav(owav_base + '.wav', predwav.permute([1, 0]).numpy(), 16000)
                     if enable_truth_wav:
                         # Note this support truth type target_f and target_mixture_f
                         tmp = torch.complex(val[0][..., :half], val[0][..., half:2 * half]).permute(2, 0, 1).detach()
+                        itf.reset()
                         truthwav, _ = itf.execute_all(tmp)
-                        save(owav_base + '_truth.wav', truthwav.permute([1, 0]), 16000, encoding='PCM_S',
-                             bits_per_sample=16)
+                        write_wav(owav_base + '_truth.wav', truthwav.permute([1, 0]).numpy(), 16000)
 
                     if enable_mix_wav:
                         tmp = torch.complex(val[0][..., 2 * half:3 * half], val[0][..., 3 * half:]).permute(2, 0, 1)
+                        itf.reset()
                         mixwav, _ = itf.execute_all(tmp.detach())
-                        save(owav_base + "_mix.wav", mixwav.permute([1, 0]), 16000, encoding='PCM_S',
-                             bits_per_sample=16)
-                        # write_wav(owav_base + "_truth.wav", truthwav, 16000)
+                        write_wav(owav_base + '_mix.wav', mixwav.permute([1, 0]).numpy(), 16000)
 
         else:
             logger.info(f'Running {mixdb.num_mixtures} mixtures with model builtin prediction loop ...')
-            from lightning.pytorch import Trainer
             pred_writer = CustomWriter(output_dir=output_dir, write_interval="epoch")
             trainer = Trainer(default_root_dir=output_dir,
                               callbacks=[pred_writer],
@@ -489,32 +479,37 @@ def main() -> None:
         # logger.info(f'Saved results to {output_dir}')
         # return
 
-    if not all(isfile(file) and splitext(file)[1] == '.wav' for file in input_name):
-        logger.exception(f'Do not know how to process input from {input_name}')
-        raise SystemExit(1)
-
-    logger.info(f'Run prediction on {len(input_name):,} WAV files')
+    logger.info(f'Run prediction on {len(input_name):,} audio files')
     for file in input_name:
-        # Convert WAV to feature data
-        audio = read_audio(file)
-        feature = get_feature_from_audio(audio=audio, feature=model.feature)
+        # Convert audio to feature data
+        audio_in = read_audio(file)
+        feature = get_feature_from_audio(audio=audio_in, feature_mode=model.hparams.feature)
 
-        # feature, predict = _pad_and_predict(hypermodel=hypermodel,
-        #                                     built_model=built_model,
-        #                                     feature=feature,
-        #                                     frames_per_batch=frames_per_batch)
+        with torch.no_grad():
+            predict = model(torch.tensor(feature))
 
-        # clean = torch_istft_olsa_hanns(clean_spec_cmplx, mixdb.ift_config.N, mixdb.ift_config.R)
+        audio_out = get_audio_from_feature(feature=predict.numpy(), feature_mode=model.hparams.feature)
 
         output_name = join(output_dir, splitext(basename(file))[0] + '.h5')
         with h5py.File(output_name, 'a') as f:
+            if 'audio_in' in f:
+                del f['audio_in']
+            f.create_dataset(name='audio_in', data=audio_in)
+
             if 'feature' in f:
                 del f['feature']
             f.create_dataset(name='feature', data=feature)
 
-            # if 'predict' in f:
-            #     del f['predict']
-            # f.create_dataset(name='predict', data=predict)
+            if 'predict' in f:
+                del f['predict']
+            f.create_dataset(name='predict', data=predict)
+
+            if 'audio_out' in f:
+                del f['audio_out']
+            f.create_dataset(name='audio_out', data=audio_out)
+
+        output_name = join(output_dir, splitext(basename(file))[0] + '_predict.wav')
+        write_wav(output_name, audio_out, 16000)
 
     logger.info(f'Saved results to {output_dir}')
     del model

sonusai/utils/__init__.py

@@ -2,10 +2,14 @@
 from .asl_p56 import asl_p56
 from .asr import ASRResult
 from .asr import calc_asr
+from .audio_devices import get_default_input_device
+from .audio_devices import get_input_device_index_by_name
+from .audio_devices import get_input_devices
 from .braced_glob import braced_glob
 from .braced_glob import braced_iglob
 from .calculate_input_shape import calculate_input_shape
 from .convert_string_to_number import convert_string_to_number
+from .create_timestamp import create_timestamp
 from .create_ts_name import create_ts_name
 from .dataclass_from_dict import dataclass_from_dict
 from .db import db_to_linear