sonusai 0.20.3__py3-none-any.whl → 1.0.2__py3-none-any.whl

sonusai/mixture/helpers.py

@@ -1,23 +1,20 @@
 from pyaaware import ForwardTransform
 from pyaaware import InverseTransform
 
-from sonusai.mixture.datatypes import AudioF
-from sonusai.mixture.datatypes import AudioT
-from sonusai.mixture.datatypes import Augmentation
-from sonusai.mixture.datatypes import AugmentationRule
-from sonusai.mixture.datatypes import EnergyT
-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 SpeechMetadata
-from sonusai.mixture.datatypes import Target
-from sonusai.mixture.datatypes import TargetFile
-from sonusai.mixture.datatypes import TransformConfig
-from sonusai.mixture.db_datatypes import MixtureRecord
-from sonusai.mixture.db_datatypes import TargetRecord
-from sonusai.mixture.mixdb import MixtureDatabase
+from ..datatypes import AudioF
+from ..datatypes import AudioT
+from ..datatypes import EnergyT
+from ..datatypes import FeatureGeneratorConfig
+from ..datatypes import FeatureGeneratorInfo
+from ..datatypes import GeneralizedIDs
+from ..datatypes import Mixture
+from ..datatypes import Source
+from ..datatypes import Sources
+from ..datatypes import SpeechMetadata
+from ..datatypes import TransformConfig
+from .db_datatypes import MixtureRecord
+from .db_datatypes import SourceRecord
+from .mixdb import MixtureDatabase
 
 
 def generic_ids_to_list(num_ids: int, ids: GeneralizedIDs = "*") -> list[int]:
@@ -58,17 +55,12 @@ def generic_ids_to_list(num_ids: int, ids: GeneralizedIDs = "*") -> list[int]:
     return result
 
 
-def get_feature_generator_info(
-    fg_config: FeatureGeneratorConfig,
-) -> FeatureGeneratorInfo:
-    from dataclasses import asdict
-
+def get_feature_generator_info(fg_config: FeatureGeneratorConfig) -> FeatureGeneratorInfo:
     from pyaaware import FeatureGenerator
 
-    from .datatypes import FeatureGeneratorInfo
-    from .datatypes import TransformConfig
+    from ..datatypes import TransformConfig
 
-    fg = FeatureGenerator(**asdict(fg_config))
+    fg = FeatureGenerator(feature_mode=fg_config.feature_mode)
 
     return FeatureGeneratorInfo(
         decimation=fg.decimation,
@@ -99,38 +91,35 @@ def get_feature_generator_info(
     )
 
 
-def mixture_all_speech_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> list[dict[str, SpeechMetadata]]:
+def mixture_all_speech_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> dict[str, dict[str, SpeechMetadata]]:
     """Get a list of all speech metadata for the given mixture"""
     from praatio.utilities.constants import Interval
 
-    from .datatypes import SpeechMetadata
+    from ..datatypes import SpeechMetadata
 
-    results: list[dict[str, SpeechMetadata]] = []
-    for target in mixture.targets:
+    results: dict[str, dict[str, SpeechMetadata]] = {}
+    for category, source in mixture.all_sources.items():
         data: dict[str, SpeechMetadata] = {}
         for tier in mixdb.speaker_metadata_tiers:
-            data[tier] = mixdb.speaker(mixdb.target_file(target.file_id).speaker_id, tier)
+            data[tier] = mixdb.speaker(mixdb.source_file(source.file_id).speaker_id, tier)
 
         for tier in mixdb.textgrid_metadata_tiers:
-            item = get_textgrid_tier_from_target_file(mixdb.target_file(target.file_id).name, tier)
+            item = get_textgrid_tier_from_source_file(mixdb.source_file(source.file_id).name, tier)
             if isinstance(item, list):
-                # Check for tempo augmentation and adjust Interval start and end data as needed
+                # Check for tempo effect and adjust Interval start and end data as needed
                 entries = []
                 for entry in item:
-                    if target.augmentation.pre.tempo is not None:
-                        entries.append(
-                            Interval(
-                                entry.start / target.augmentation.pre.tempo,
-                                entry.end / target.augmentation.pre.tempo,
-                                entry.label,
-                            )
+                    entries.append(
+                        Interval(
+                            entry.start / source.pre_tempo,
+                            entry.end / source.pre_tempo,
+                            entry.label,
                         )
-                    else:
-                        entries.append(entry)
+                    )
                 data[tier] = entries
             else:
                 data[tier] = item
-        results.append(data)
+        results[category] = data
 
     return results
 
@@ -151,27 +140,23 @@ def mixture_metadata(mixdb: MixtureDatabase, m_id: int | None = None, mixture: M
 
     metadata = ""
     speech_metadata = mixture_all_speech_metadata(mixdb, mixture)
-    for mi, target in enumerate(mixture.targets):
-        target_file = mixdb.target_file(target.file_id)
-        metadata += f"target {mi} name: {target_file.name}\n"
-        metadata += f"target {mi} augmentation: {target.augmentation.to_dict()}\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"
-            metadata += f"target {mi} truth '{key}' config:   {target_file.truth_configs[key].config}\n"
-        for key in speech_metadata[mi]:
-            metadata += f"target {mi} speech {key}: {speech_metadata[mi][key]}\n"
-    noise = mixdb.noise_file(mixture.noise.file_id)
-    noise_augmentation = mixture.noise.augmentation
-    metadata += f"noise name: {noise.name}\n"
-    metadata += f"noise augmentation: {noise_augmentation.to_dict()}\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"
-    metadata += f"target_snr_gain: {float(mixture.target_snr_gain)}\n"
-    metadata += f"noise_snr_gain: {float(mixture.noise_snr_gain)}\n"
+    for category, source in mixture.all_sources.items():
+        source_file = mixdb.source_file(source.file_id)
+        metadata += f"{category} name: {source_file.name}\n"
+        metadata += f"{category} effects: {source.effects.to_dict()}\n"
+        metadata += f"{category} pre_tempo: {source.pre_tempo}\n"
+        metadata += f"{category} class indices: {source_file.class_indices}\n"
+        metadata += f"{category} start: {source.start}\n"
+        metadata += f"{category} repeat: {source.repeat}\n"
+        metadata += f"{category} snr: {source.snr}\n"
+        metadata += f"{category} random_snr: {source.snr.is_random}\n"
+        metadata += f"{category} snr_gain: {source.snr_gain}\n"
+        for key in source_file.truth_configs:
+            metadata += f"{category} truth '{key}' function: {source_file.truth_configs[key].function}\n"
+            metadata += f"{category} truth '{key}' config:   {source_file.truth_configs[key].config}\n"
+        for key in speech_metadata[category]:
+            metadata += f"{category} speech {key}: {speech_metadata[category][key]}\n"
 
     return metadata
 
@@ -197,95 +182,51 @@ def write_mixture_metadata(mixdb: MixtureDatabase, m_id: int | None = None, mixt
         f.write(mixture_metadata(mixdb, m_id, mixture))
 
 
-def from_mixture(
-    mixture: Mixture,
-) -> tuple[str, int, str, int, float, bool, float, int, int, int, float]:
-    return (
-        mixture.name,
-        mixture.noise.file_id,
-        mixture.noise.augmentation.to_json(),
-        mixture.noise_offset,
-        mixture.noise_snr_gain,
-        mixture.snr.is_random,
-        mixture.snr,
-        mixture.samples,
-        mixture.spectral_mask_id,
-        mixture.spectral_mask_seed,
-        mixture.target_snr_gain,
-    )
-
-
-def to_mixture(entry: MixtureRecord, targets: list[Target]) -> Mixture:
-    import json
-
-    from sonusai.utils import dataclass_from_dict
+def from_mixture(mixture: Mixture) -> tuple[str, int, int, int]:
+    return mixture.name, mixture.samples, mixture.spectral_mask_id, mixture.spectral_mask_seed
 
-    from .datatypes import Noise
-    from .datatypes import UniversalSNR
 
+def to_mixture(entry: MixtureRecord, sources: Sources) -> Mixture:
     return Mixture(
-        targets=targets,
         name=entry.name,
-        noise=Noise(
-            file_id=entry.noise_file_id,
-            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,
+        all_sources=sources,
         spectral_mask_id=entry.spectral_mask_id,
         spectral_mask_seed=entry.spectral_mask_seed,
-        target_snr_gain=entry.target_snr_gain,
     )
 
 
-def from_target(target: Target) -> tuple[int, str]:
-    return target.file_id, target.augmentation.to_json()
+def from_source(source: Source) -> tuple[str, int, float, bool, float, float, bool, int]:
+    return (
+        source.effects.to_json(),
+        source.file_id,
+        source.pre_tempo,
+        source.repeat,
+        source.snr,
+        source.snr_gain,
+        source.snr.is_random,
+        source.start,
+    )
 
 
-def to_target(entry: TargetRecord) -> Target:
+def to_source(entry: SourceRecord) -> Source:
     import json
 
-    from sonusai.utils import dataclass_from_dict
-
-    from .datatypes import Augmentation
+    from ..datatypes import Effects
+    from ..datatypes import UniversalSNR
+    from ..utils.dataclass_from_dict import dataclass_from_dict
 
-    return Target(
+    return Source(
         file_id=entry.file_id,
-        augmentation=dataclass_from_dict(Augmentation, json.loads(entry.augmentation)),  # pyright: ignore [reportArgumentType]
+        effects=dataclass_from_dict(Effects, json.loads(entry.effects)),
+        start=entry.start,
+        repeat=entry.repeat,
+        snr=UniversalSNR(entry.snr, entry.snr_random),
+        snr_gain=entry.snr_gain,
+        pre_tempo=entry.pre_tempo,
     )
 
 
-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
-    :param mixture: Mixture record
-    :param targets_audio: List of augmented target audio data (one per target in the mixup)
-    :return: Sum of augmented target audio data
-    """
-    # Apply post-truth augmentation effects to targets and sum
-    import numpy as np
-
-    from .augmentation import apply_augmentation
-
-    targets_post = []
-    for idx, target_audio in enumerate(targets_audio):
-        target = mixture.targets[idx]
-        targets_post.append(
-            apply_augmentation(
-                mixdb=mixdb,
-                audio=target_audio,
-                augmentation=target.augmentation.post,
-                frame_length=mixdb.feature_step_samples,
-            )
-        )
-
-    # Return sum of targets
-    return np.sum(targets_post, axis=0)
-
-
 def get_transform_from_audio(audio: AudioT, transform: ForwardTransform) -> tuple[AudioF, EnergyT]:
     """Apply forward transform to input audio data to generate transform data
 
@@ -368,67 +309,24 @@ def check_audio_files_exist(mixdb: MixtureDatabase) -> None:
     """Walk through all the noise and target audio files in a mixture database ensuring that they exist"""
     from os.path import exists
 
-    from .tokenized_shell_vars import tokenized_expand
-
-    for noise in mixdb.noise_files:
-        file_name, _ = tokenized_expand(noise.name)
-        if not exists(file_name):
-            raise OSError(f"Could not find {file_name}")
-
-    for target in mixdb.target_files:
-        file_name, _ = tokenized_expand(target.name)
-        if not exists(file_name):
-            raise OSError(f"Could not find {file_name}")
-
-
-def augmented_target_samples(
-    target_files: list[TargetFile],
-    target_augmentations: list[AugmentationRule],
-    feature_step_samples: int,
-) -> int:
-    from itertools import product
-
-    from .augmentation import estimate_augmented_length_from_length
-
-    target_ids = list(range(len(target_files)))
-    target_augmentation_ids = list(range(len(target_augmentations)))
-    it = list(product(*[target_ids, target_augmentation_ids]))
-    return sum(
-        [
-            estimate_augmented_length_from_length(
-                length=target_files[fi].samples,
-                tempo=target_augmentations[ai].pre.tempo,
-                frame_length=feature_step_samples,
-            )
-            for fi, ai in it
-        ]
-    )
-
-
-def augmented_noise_samples(noise_files: list[NoiseFile], noise_augmentations: list[Augmentation]) -> int:
-    from itertools import product
-
-    noise_ids = list(range(len(noise_files)))
-    noise_augmentation_ids = list(range(len(noise_augmentations)))
-    it = list(product(*[noise_ids, noise_augmentation_ids]))
-    return sum([augmented_noise_length(noise_files[fi], noise_augmentations[ai]) for fi, ai in it])
-
-
-def augmented_noise_length(noise_file: NoiseFile, noise_augmentation: Augmentation) -> int:
-    from .augmentation import estimate_augmented_length_from_length
+    from ..utils.tokenized_shell_vars import tokenized_expand
 
-    return estimate_augmented_length_from_length(length=noise_file.samples, tempo=noise_augmentation.pre.tempo)
+    for source_files in mixdb.source_files.values():
+        for source_file in source_files:
+            file_name, _ = tokenized_expand(source_file.name)
+            if not exists(file_name):
+                raise OSError(f"Could not find {file_name}")
 
 
-def get_textgrid_tier_from_target_file(target_file: str, tier: str) -> SpeechMetadata | None:
+def get_textgrid_tier_from_source_file(source_file: str, tier: str) -> SpeechMetadata | None:
     from pathlib import Path
 
     from praatio import textgrid
     from praatio.utilities.constants import Interval
 
-    from .tokenized_shell_vars import tokenized_expand
+    from ..utils.tokenized_shell_vars import tokenized_expand
 
-    textgrid_file = Path(tokenized_expand(target_file)[0]).with_suffix(".TextGrid")
+    textgrid_file = Path(tokenized_expand(source_file)[0]).with_suffix(".TextGrid")
     if not textgrid_file.exists():
         return None
 

sonusai/mixture/ir_delay.py

@@ -1,14 +1,13 @@
 import numpy as np
 
 
-def get_impulse_response_delay(file: str) -> int:
-    from sonusai.utils import temp_seed
-
+def get_ir_delay(file: str) -> int:
+    from ..utils.rand import seed_context
     from .audio import raw_read_audio
 
     ir, sample_rate = raw_read_audio(file)
 
-    with temp_seed(42):
+    with seed_context(42):
         wgn_ref = np.random.normal(loc=0, scale=0.2, size=int(np.ceil(0.05 * sample_rate))).astype(np.float32)
 
     wgn_conv = np.convolve(ir, wgn_ref)

sonusai/mixture/ir_effects.py

@@ -0,0 +1,77 @@
+from functools import lru_cache
+from pathlib import Path
+
+from ..datatypes import AudioT
+from ..datatypes import ImpulseResponseData
+from .audio import raw_read_audio
+
+
+def apply_ir(audio: AudioT, ir: ImpulseResponseData) -> AudioT:
+    """Apply impulse response to audio data using scipy
+
+    :param audio: Audio
+    :param ir: Impulse response data
+    :return: Effected audio
+    """
+    import numpy as np
+    from librosa import resample
+    from scipy.signal import fftconvolve
+
+    from ..constants import SAMPLE_RATE
+
+    if not isinstance(audio, np.ndarray):
+        raise TypeError("audio must be a numpy array")
+
+    # Early exit if no ir or if all audio is zero
+    if ir is None or not audio.any():
+        return audio
+
+    pk_in = np.max(np.abs(audio))
+
+    # Convert audio to IR sample rate
+    audio_in = resample(audio, orig_sr=SAMPLE_RATE, target_sr=ir.sample_rate, res_type="soxr_hq")
+
+    # Apply IR
+    audio_out = fftconvolve(audio_in, ir.data, mode="full")
+
+    # Delay compensation
+    audio_out = audio_out[ir.delay :]
+
+    # Convert back to global sample rate
+    audio_out = resample(audio_out, orig_sr=ir.sample_rate, target_sr=SAMPLE_RATE, res_type="soxr_hq")
+
+    # Trim to length
+    audio_out = audio_out[: len(audio)]
+
+    # Gain compensation
+    pk_out = np.max(np.abs(audio_out))
+    pk_gain = pk_in / pk_out
+    audio_out = audio_out * pk_gain
+
+    return audio_out
+
+
+def read_ir(name: str | Path, delay: int, use_cache: bool = True) -> ImpulseResponseData:
+    """Read impulse response data
+
+    :param name: File name
+    :param delay: Delay in samples
+    :param use_cache: If true, use LRU caching
+    :return: ImpulseResponseData object
+    """
+    if use_cache:
+        return _read_ir(name, delay)
+    return _read_ir.__wrapped__(name, delay)
+
+
+@lru_cache
+def _read_ir(name: str | Path, delay: int) -> ImpulseResponseData:
+    """Read impulse response data using soundfile
+
+    :param name: File name
+    :param delay: Delay in samples
+    :return: ImpulseResponseData object
+    """
+    out, sample_rate = raw_read_audio(name)
+
+    return ImpulseResponseData(data=out, sample_rate=sample_rate, delay=delay)

sonusai/mixture/log_duration_and_sizes.py

@@ -1,29 +1,23 @@
 def log_duration_and_sizes(
     total_duration: float,
-    num_classes: int,
     feature_step_samples: int,
     feature_parameters: int,
     stride: int,
     desc: str,
 ) -> None:
-    from sonusai import logger
-    from sonusai.utils import human_readable_size
-    from sonusai.utils import seconds_to_hms
-
-    from .constants import FLOAT_BYTES
-    from .constants import SAMPLE_BYTES
-    from .constants import SAMPLE_RATE
+    from .. import logger
+    from ..constants import FLOAT_BYTES
+    from ..constants import SAMPLE_BYTES
+    from ..constants import SAMPLE_RATE
+    from ..utils.human_readable_size import human_readable_size
+    from ..utils.seconds_to_hms import seconds_to_hms
 
     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 * feature_parameters * FLOAT_BYTES
-    truth_f_bytes = total_samples / feature_step_samples * num_classes * FLOAT_BYTES
 
     logger.info("")
     logger.info(f"{desc} duration:   {seconds_to_hms(seconds=total_duration)}")
     logger.info(f"{desc} sizes:")
     logger.info(f" mixture:             {human_readable_size(mixture_bytes, 1)}")
-    logger.info(f" truth_t:             {human_readable_size(truth_t_bytes, 1)}")
     logger.info(f" feature:             {human_readable_size(feature_bytes, 1)}")
-    logger.info(f" truth_f:             {human_readable_size(truth_f_bytes, 1)}")