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

sonusai/mixture/helpers.py

@@ -2,26 +2,19 @@ from pyaaware import ForwardTransform
 from pyaaware import InverseTransform
 
 from sonusai.mixture.datatypes import AudioF
-from sonusai.mixture.datatypes import AudiosT
 from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import Augmentation
-from sonusai.mixture.datatypes import AugmentationRules
-from sonusai.mixture.datatypes import Augmentations
+from sonusai.mixture.datatypes import AugmentationRule
 from sonusai.mixture.datatypes import EnergyT
-from sonusai.mixture.datatypes import Feature
 from sonusai.mixture.datatypes import FeatureGeneratorConfig
 from sonusai.mixture.datatypes import FeatureGeneratorInfo
 from sonusai.mixture.datatypes import GeneralizedIDs
 from sonusai.mixture.datatypes import Mixture
 from sonusai.mixture.datatypes import NoiseFile
-from sonusai.mixture.datatypes import NoiseFiles
-from sonusai.mixture.datatypes import Segsnr
 from sonusai.mixture.datatypes import SpeechMetadata
 from sonusai.mixture.datatypes import Target
-from sonusai.mixture.datatypes import TargetFiles
-from sonusai.mixture.datatypes import Targets
+from sonusai.mixture.datatypes import TargetFile
 from sonusai.mixture.datatypes import TransformConfig
-from sonusai.mixture.datatypes import TruthDict
 from sonusai.mixture.db_datatypes import MixtureRecord
 from sonusai.mixture.db_datatypes import TargetRecord
 from sonusai.mixture.mixdb import MixtureDatabase
@@ -142,13 +135,14 @@ def mixture_all_speech_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> lis
     return results
 
 
-def mixture_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> str:
+def mixture_metadata(mixdb: MixtureDatabase, m_id: int) -> str:
     """Create a string of metadata for a Mixture
 
     :param mixdb: Mixture database
-    :param mixture: Mixture record
+    :param m_id: Mixture ID
     :return: String of metadata
     """
+    mixture = mixdb.mixture(m_id)
     metadata = ""
     speech_metadata = mixture_all_speech_metadata(mixdb, mixture)
     for mi, target in enumerate(mixture.targets):
@@ -157,7 +151,7 @@ def mixture_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> str:
         metadata += f"target {mi} name: {target_file.name}\n"
         metadata += f"target {mi} augmentation: {target.augmentation.to_dict()}\n"
         metadata += f"target {mi} ir: {mixdb.impulse_response_file(target_augmentation.ir)}\n"
-        metadata += f"target {mi} target_gain: {target.gain}\n"
+        metadata += f"target {mi} target_gain: {target.gain if not mixture.is_noise_only else 0}\n"
         metadata += f"target {mi} class indices: {target_file.class_indices}\n"
         for key in target_file.truth_configs:
             metadata += f"target {mi} truth '{key}' function: {target_file.truth_configs[key].function}\n"
@@ -169,7 +163,7 @@ def mixture_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> str:
     metadata += f"noise name: {noise.name}\n"
     metadata += f"noise augmentation: {noise_augmentation.to_dict()}\n"
     metadata += f"noise ir: {mixdb.impulse_response_file(noise_augmentation.ir)}\n"
-    metadata += f"noise offset: {mixture.noise.offset}\n"
+    metadata += f"noise offset: {mixture.noise_offset}\n"
     metadata += f"snr: {mixture.snr}\n"
     metadata += f"random_snr: {mixture.snr.is_random}\n"
     metadata += f"samples: {mixture.samples}\n"
@@ -179,17 +173,17 @@ def mixture_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> str:
     return metadata
 
 
-def write_mixture_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> None:
+def write_mixture_metadata(mixdb: MixtureDatabase, m_id: int) -> None:
     """Write mixture metadata to a text file
 
     :param mixdb: Mixture database
-    :param mixture: Mixture record
+    :param m_id: Mixture ID
     """
     from os.path import join
 
-    name = join(mixdb.location, "mixture", mixture.name, "metadata.txt")
+    name = join(mixdb.location, "mixture", mixdb.mixture(m_id).name, "metadata.txt")
     with open(file=name, mode="w") as f:
-        f.write(mixture_metadata(mixdb, mixture))
+        f.write(mixture_metadata(mixdb, m_id))
 
 
 def from_mixture(
@@ -199,7 +193,7 @@ def from_mixture(
         mixture.name,
         mixture.noise.file_id,
         mixture.noise.augmentation.to_json(),
-        mixture.noise.offset,
+        mixture.noise_offset,
         mixture.noise_snr_gain,
         mixture.snr.is_random,
         mixture.snr,
@@ -210,7 +204,7 @@ def from_mixture(
     )
 
 
-def to_mixture(entry: MixtureRecord, targets: Targets) -> Mixture:
+def to_mixture(entry: MixtureRecord, targets: list[Target]) -> Mixture:
     import json
 
     from sonusai.utils import dataclass_from_dict
@@ -223,9 +217,9 @@ def to_mixture(entry: MixtureRecord, targets: Targets) -> Mixture:
         name=entry.name,
         noise=Noise(
             file_id=entry.noise_file_id,
-            augmentation=dataclass_from_dict(Augmentation, json.loads(entry.noise_augmentation)),
-            offset=entry.noise_offset,
+            augmentation=dataclass_from_dict(Augmentation, json.loads(entry.noise_augmentation)),  # pyright: ignore [reportArgumentType]
         ),
+        noise_offset=entry.noise_offset,
         noise_snr_gain=entry.noise_snr_gain,
         snr=UniversalSNR(is_random=entry.random_snr, value=entry.snr),
         samples=entry.samples,
@@ -235,8 +229,8 @@ def to_mixture(entry: MixtureRecord, targets: Targets) -> Mixture:
     )
 
 
-def from_target(target: Target) -> tuple[int, str, float]:
-    return target.file_id, target.augmentation.to_json(), target.gain
+def from_target(target: Target) -> tuple[int, str]:
+    return target.file_id, target.augmentation.to_json()
 
 
 def to_target(entry: TargetRecord) -> Target:
@@ -245,147 +239,14 @@ def to_target(entry: TargetRecord) -> Target:
     from sonusai.utils import dataclass_from_dict
 
     from .datatypes import Augmentation
-    from .datatypes import Target
 
     return Target(
         file_id=entry.file_id,
-        augmentation=dataclass_from_dict(Augmentation, json.loads(entry.augmentation)),
-        gain=entry.gain,
-    )
-
-
-def get_truth(
-    mixdb: MixtureDatabase,
-    mixture: Mixture,
-    targets_audio: AudiosT,
-    noise_audio: AudioT,
-    mixture_audio: AudioT,
-) -> TruthDict:
-    """Get the truth data for the given mixture record
-
-    :param mixdb: Mixture database
-    :param mixture: Mixture record
-    :param targets_audio: List of augmented target audio data (one per target in the mixup) for the given mixture ID
-    :param noise_audio: Augmented noise audio data for the given mixture ID
-    :param mixture_audio: Mixture audio data for the given mixture ID
-    :return: truth data
-    """
-    from .datatypes import TruthDict
-    from .truth import truth_function
-
-    if not all(len(target) == mixture.samples for target in targets_audio):
-        raise ValueError("Lengths of targets do not match length of mixture")
-
-    if len(noise_audio) != mixture.samples:
-        raise ValueError("Length of noise does not match length of mixture")
-
-    # TODO: Need to understand how to do this correctly for mixup and target_mixture_f truth
-    if len(targets_audio) != 1:
-        raise NotImplementedError("mixup is not implemented")
-
-    truth: TruthDict = {}
-    for idx in range(len(targets_audio)):
-        target_file = mixdb.target_file(mixture.targets[idx].file_id)
-        for key, value in target_file.truth_configs.items():
-            truth[key] = truth_function(
-                target_audio=targets_audio[idx],
-                noise_audio=noise_audio,
-                mixture_audio=mixture_audio,
-                config=value,
-                feature=mixdb.feature,
-                num_classes=mixdb.num_classes,
-                class_indices=target_file.class_indices,
-                target_gain=mixture.targets[idx].gain * mixture.target_snr_gain,
-            )
-
-    return truth
-
-
-def get_ft(
-    mixdb: MixtureDatabase, mixture: Mixture, mixture_audio: AudioT, truth_t: TruthDict
-) -> tuple[Feature, TruthDict]:
-    """Get the feature and truth_f data for the given mixture record
-
-    :param mixdb: Mixture database
-    :param mixture: Mixture record
-    :param mixture_audio: Mixture audio data for the given mixid
-    :param truth_t: truth_t for the given mixid
-    :return: Tuple of (feature, truth_f) data
-    """
-
-    from pyaaware import FeatureGenerator
-
-    from .truth import truth_stride_reduction
-
-    mixture_f = get_mixture_f(mixdb=mixdb, mixture=mixture, mixture_audio=mixture_audio)
-
-    fg = FeatureGenerator(mixdb.fg_config.feature_mode, mixdb.fg_config.truth_parameters)
-    feature, truth_f = fg.execute_all(mixture_f, truth_t)
-    for name in truth_f:
-        truth_f[name] = truth_stride_reduction(truth_f[name], mixdb.truth_configs[name].stride_reduction)
-
-    return feature, truth_f
-
-
-def get_segsnr(mixdb: MixtureDatabase, mixture: Mixture, target_audio: AudioT, noise: AudioT) -> Segsnr:
-    """Get the segsnr data for the given mixture record
-
-    :param mixdb: Mixture database
-    :param mixture: Mixture record
-    :param target_audio: Augmented target audio data
-    :param noise: Augmented noise audio data
-    :return: segsnr data
-    """
-    segsnr_t = get_segsnr_t(mixdb=mixdb, mixture=mixture, target_audio=target_audio, noise_audio=noise)
-    return segsnr_t[0 :: mixdb.ft_config.overlap]
-
-
-def get_segsnr_t(mixdb: MixtureDatabase, mixture: Mixture, target_audio: AudioT, noise_audio: AudioT) -> Segsnr:
-    """Get the segsnr_t data for the given mixture record
-
-    :param mixdb: Mixture database
-    :param mixture: Mixture record
-    :param target_audio: Augmented target audio data
-    :param noise_audio: Augmented noise audio data
-    :return: segsnr_t data
-    """
-    import numpy as np
-    import torch
-    from pyaaware import ForwardTransform
-
-    fft = ForwardTransform(
-        length=mixdb.ft_config.length,
-        overlap=mixdb.ft_config.overlap,
-        bin_start=mixdb.ft_config.bin_start,
-        bin_end=mixdb.ft_config.bin_end,
-        ttype=mixdb.ft_config.ttype,
+        augmentation=dataclass_from_dict(Augmentation, json.loads(entry.augmentation)),  # pyright: ignore [reportArgumentType]
     )
 
-    segsnr_t = np.empty(mixture.samples, dtype=np.float32)
-
-    target_energy = fft.execute_all(torch.from_numpy(target_audio))[1].numpy()
-    noise_energy = fft.execute_all(torch.from_numpy(noise_audio))[1].numpy()
-
-    offsets = range(0, mixture.samples, mixdb.ft_config.overlap)
-    if len(target_energy) != len(offsets):
-        raise ValueError(
-            f"Number of frames in energy, {len(target_energy)}," f" is not number of frames in mixture, {len(offsets)}"
-        )
-
-    for idx, offset in enumerate(offsets):
-        indices = slice(offset, offset + mixdb.ft_config.overlap)
-
-        if noise_energy[idx] == 0:
-            snr = np.float32(np.inf)
-        else:
-            snr = np.float32(target_energy[idx] / noise_energy[idx])
-
-        segsnr_t[indices] = snr
-
-    return segsnr_t
 
-
-def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: AudiosT) -> AudioT:
+def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: list[AudioT]) -> AudioT:
     """Get the augmented target audio data for the given mixture record
 
     :param mixdb: Mixture database
@@ -413,28 +274,6 @@ def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: AudiosT)
     return np.sum(targets_ir, axis=0)
 
 
-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
-    """
-    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
 
@@ -497,7 +336,6 @@ def inverse_transform(transform: AudioF, config: TransformConfig) -> AudioT:
     :param config: Transform configuration
     :return: Time domain data [samples]
     """
-    import numpy as np
     from pyaaware import InverseTransform
 
     audio, _ = get_audio_from_transform(
@@ -508,7 +346,7 @@ def inverse_transform(transform: AudioF, config: TransformConfig) -> AudioT:
             bin_start=config.bin_start,
             bin_end=config.bin_end,
             ttype=config.ttype,
-            gain=np.float32(1),
+            gain=1,
         ),
     )
     return audio
@@ -532,8 +370,8 @@ def check_audio_files_exist(mixdb: MixtureDatabase) -> None:
 
 
 def augmented_target_samples(
-    target_files: TargetFiles,
-    target_augmentations: AugmentationRules,
+    target_files: list[TargetFile],
+    target_augmentations: list[AugmentationRule],
     feature_step_samples: int,
 ) -> int:
     from itertools import product
@@ -555,7 +393,7 @@ def augmented_target_samples(
     )
 
 
-def augmented_noise_samples(noise_files: NoiseFiles, noise_augmentations: Augmentations) -> int:
+def augmented_noise_samples(noise_files: list[NoiseFile], noise_augmentations: list[Augmentation]) -> int:
     from itertools import product
 
     noise_ids = list(range(len(noise_files)))
@@ -574,6 +412,7 @@ def get_textgrid_tier_from_target_file(target_file: str, tier: str) -> SpeechMet
     from pathlib import Path
 
     from praatio import textgrid
+    from praatio.utilities.constants import Interval
 
     from .tokenized_shell_vars import tokenized_expand
 
@@ -588,10 +427,13 @@ def get_textgrid_tier_from_target_file(target_file: str, tier: str) -> SpeechMet
 
     entries = tg.getTier(tier).entries
     if len(entries) > 1:
-        return list(entries)
-    else:
+        return [entry for entry in entries if isinstance(entry, Interval)]
+
+    if len(entries) == 1:
         return entries[0].label
 
+    return None
+
 
 def frames_from_samples(samples: int, step_samples: int) -> int:
     import numpy as np