sonusai 0.19.10__py3-none-any.whl → 0.20.2__py3-none-any.whl

sonusai/mixture/audio.py

@@ -58,9 +58,62 @@ def get_sample_rate(name: str | Path, use_cache: bool = True) -> int:
 
 @lru_cache
 def _get_sample_rate(name: str | Path) -> int:
-    from .soundfile_audio import get_sample_rate
+    """Get sample rate from audio file using soundfile
 
-    return get_sample_rate(name)
+    :param name: File name
+    :return: Sample rate
+    """
+    import soundfile
+    from pydub import AudioSegment
+
+    from .tokenized_shell_vars import tokenized_expand
+
+    expanded_name, _ = tokenized_expand(name)
+
+    try:
+        if expanded_name.endswith(".mp3"):
+            return AudioSegment.from_mp3(expanded_name).frame_rate
+
+        if expanded_name.endswith(".m4a"):
+            return AudioSegment.from_file(expanded_name).frame_rate
+
+        return soundfile.info(expanded_name).samplerate
+    except Exception as e:
+        if name != expanded_name:
+            raise OSError(f"Error reading {name} (expanded: {expanded_name}): {e}") from e
+        else:
+            raise OSError(f"Error reading {name}: {e}") from e
+
+
+def raw_read_audio(name: str | Path) -> tuple[AudioT, int]:
+    import numpy as np
+    import soundfile
+    from pydub import AudioSegment
+
+    from .tokenized_shell_vars import tokenized_expand
+
+    expanded_name, _ = tokenized_expand(name)
+
+    try:
+        if expanded_name.endswith(".mp3"):
+            sound = AudioSegment.from_mp3(expanded_name)
+            raw = np.array(sound.get_array_of_samples()).astype(np.float32).reshape((-1, sound.channels))
+            raw = raw / 2 ** (sound.sample_width * 8 - 1)
+            sample_rate = sound.frame_rate
+        elif expanded_name.endswith(".m4a"):
+            sound = AudioSegment.from_file(expanded_name)
+            raw = np.array(sound.get_array_of_samples()).astype(np.float32).reshape((-1, sound.channels))
+            raw = raw / 2 ** (sound.sample_width * 8 - 1)
+            sample_rate = sound.frame_rate
+        else:
+            raw, sample_rate = soundfile.read(expanded_name, always_2d=True, dtype="float32")
+    except Exception as e:
+        if name != expanded_name:
+            raise OSError(f"Error reading {name} (expanded: {expanded_name}): {e}") from e
+        else:
+            raise OSError(f"Error reading {name}: {e}") from e
+
+    return np.squeeze(raw[:, 0].astype(np.float32)), sample_rate
 
 
 def read_audio(name: str | Path, use_cache: bool = True) -> AudioT:
@@ -77,28 +130,45 @@ def read_audio(name: str | Path, use_cache: bool = True) -> AudioT:
 
 @lru_cache
 def _read_audio(name: str | Path) -> AudioT:
-    from .soundfile_audio import read_audio
+    """Read audio data from a file using soundfile
+
+    :param name: File name
+    :return: Array of time domain audio data
+    """
+    import librosa
+
+    from .constants import SAMPLE_RATE
+
+    out, sample_rate = raw_read_audio(name)
+    out = librosa.resample(out, orig_sr=sample_rate, target_sr=SAMPLE_RATE, res_type="soxr_hq")
 
-    return read_audio(name)
+    return out
 
 
-def read_ir(name: str | Path, use_cache: bool = True) -> ImpulseResponseData:
+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)
-    return _read_ir.__wrapped__(name)
+        return _read_ir(name, delay)
+    return _read_ir.__wrapped__(name, delay)
 
 
 @lru_cache
-def _read_ir(name: str | Path) -> ImpulseResponseData:
-    from .soundfile_audio import read_ir
+def _read_ir(name: str | Path, delay: int) -> ImpulseResponseData:
+    """Read impulse response data using soundfile
 
-    return read_ir(name)
+    :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)
 
 
 def get_num_samples(name: str | Path, use_cache: bool = True) -> int:
@@ -120,6 +190,27 @@ def _get_num_samples(name: str | Path) -> int:
     :param name: File name
     :return: number of samples in resampled audio
     """
-    from .soundfile_audio import get_num_samples
+    import math
+
+    import soundfile
+    from pydub import AudioSegment
 
-    return get_num_samples(name)
+    from .constants import SAMPLE_RATE
+    from .tokenized_shell_vars import tokenized_expand
+
+    expanded_name, _ = tokenized_expand(name)
+
+    if expanded_name.endswith(".mp3"):
+        sound = AudioSegment.from_mp3(expanded_name)
+        samples = sound.frame_count()
+        sample_rate = sound.frame_rate
+    elif expanded_name.endswith(".m4a"):
+        sound = AudioSegment.from_file(expanded_name)
+        samples = sound.frame_count()
+        sample_rate = sound.frame_rate
+    else:
+        info = soundfile.info(name)
+        samples = info.frames
+        sample_rate = info.samplerate
+
+    return math.ceil(SAMPLE_RATE * samples / sample_rate)

sonusai/mixture/augmentation.py

@@ -1,8 +1,10 @@
 from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import Augmentation
+from sonusai.mixture.datatypes import AugmentationEffects
 from sonusai.mixture.datatypes import AugmentationRule
 from sonusai.mixture.datatypes import ImpulseResponseData
 from sonusai.mixture.datatypes import OptionalNumberStr
+from sonusai.mixture.mixdb import MixtureDatabase
 
 
 def get_augmentation_rules(rules: list[dict] | dict, num_ir: int = 0) -> list[AugmentationRule]:
@@ -41,49 +43,63 @@ def _expand_rules(expanded_rules: list[dict], rule: dict) -> list[dict]:
     from .constants import VALID_AUGMENTATIONS
     from .eq_rule_is_valid import eq_rule_is_valid
 
-    for key, value in list(rule.items()):
-        if value is None:
-            del rule[key]
+    if "pre" not in rule:
+        raise ValueError("Rule must have 'pre' key")
 
-    # replace old 'eq' rule with new 'eq1' rule to allow both for backward compatibility
-    rule = {"eq1" if key == "eq" else key: value for key, value in rule.items()}
+    if "post" not in rule:
+        rule["post"] = {}
+
+    for key in rule:
+        if key not in ("pre", "post", "mixup"):
+            raise ValueError(f"Invalid augmentation key: {key}")
 
+        if key in ("pre", "post"):
+            for k, v in list(rule[key].items()):
+                if v is None:
+                    del rule[key][k]
+
+    # replace old 'eq' rule with new 'eq1' rule to allow both for backward compatibility
     for key in rule:
-        if key not in VALID_AUGMENTATIONS:
-            nice_list = "\n".join([f"  {item}" for item in VALID_AUGMENTATIONS])
-            raise ValueError(f"Invalid augmentation: {key}.\nValid augmentations are:\n{nice_list}")
-
-        if key in ["eq1", "eq2", "eq3"]:
-            if not eq_rule_is_valid(rule[key]):
-                raise ValueError(f"Invalid augmentation value for {key}: {rule[key]}")
-
-            if all(isinstance(el, list) or (isinstance(el, str) and el == "none") for el in rule[key]):
-                # Expand multiple rules
-                for value in rule[key]:
-                    expanded_rule = deepcopy(rule)
-                    if isinstance(value, str) and value == "none":
-                        expanded_rule[key] = None
-                    else:
-                        expanded_rule[key] = deepcopy(value)
-                    _expand_rules(expanded_rules, expanded_rule)
-                return expanded_rules
-
-        elif key in ["mixup"]:
-            pass
-
-        else:
-            if isinstance(rule[key], list):
-                for value in rule[key]:
-                    if isinstance(value, list):
-                        raise TypeError(f"Invalid augmentation value for {key}: {rule[key]}")
-                    expanded_rule = deepcopy(rule)
-                    expanded_rule[key] = deepcopy(value)
-                    _expand_rules(expanded_rules, expanded_rule)
-                return expanded_rules
+        rule[key] = {"eq1" if k == "eq" else k: v for k, v in rule[key].items()}
+
+    for key in ("pre", "post"):
+        for k in rule[key]:
+            if k not in VALID_AUGMENTATIONS:
+                nice_list = "\n".join([f"  {item}" for item in VALID_AUGMENTATIONS])
+                raise ValueError(f"Invalid augmentation: {k}.\nValid augmentations are:\n{nice_list}")
+
+            if k in ["eq1", "eq2", "eq3"]:
+                if not eq_rule_is_valid(rule[key][k]):
+                    raise ValueError(f"Invalid augmentation value for {k}: {rule[key][k]}")
+
+                if all(isinstance(el, list) or (isinstance(el, str) and el == "none") for el in rule[key][k]):
+                    # Expand multiple rules
+                    for value in rule[key][k]:
+                        expanded_rule = deepcopy(rule)
+                        if isinstance(value, str) and value == "none":
+                            expanded_rule[key][k] = None
+                        else:
+                            expanded_rule[key][k] = deepcopy(value)
+                        _expand_rules(expanded_rules, expanded_rule)
+                    return expanded_rules
+
             else:
-                rule[key] = convert_string_to_number(rule[key])
-                if not (isinstance(rule[key], float | int) or rule[key].startswith("rand") or rule[key] == "none"):
-                    raise ValueError(f"Invalid augmentation value for {key}: {rule[key]}")
+                if isinstance(rule[key][k], list):
+                    for value in rule[key][k]:
+                        if isinstance(value, list):
+                            raise TypeError(f"Invalid augmentation value for {k}: {rule[key][k]}")
+                        expanded_rule = deepcopy(rule)
+                        expanded_rule[key][k] = deepcopy(value)
+                        _expand_rules(expanded_rules, expanded_rule)
+                    return expanded_rules
+                else:
+                    rule[key][k] = convert_string_to_number(rule[key][k])
+                    if not (
+                        isinstance(rule[key][k], float | int)
+                        or rule[key][k].startswith("rand")
+                        or rule[key][k] == "none"
+                    ):
+                        raise ValueError(f"Invalid augmentation value for {k}: {rule[key][k]}")
 
     expanded_rules.append(rule)
     return expanded_rules
@@ -116,21 +132,22 @@ def _generate_random_rule(rule: dict, num_ir: int = 0) -> dict:
     from random import randint
 
     out_rule = deepcopy(rule)
-    for key in out_rule:
-        if key == "ir" and out_rule[key] == "rand":
-            # IR is special case
-            if num_ir == 0:
-                out_rule[key] = None
+    for key in ("pre", "post"):
+        for k in out_rule[key]:
+            if k == "ir" and out_rule[key][k] == "rand":
+                # IR is special case
+                if num_ir == 0:
+                    out_rule[key][k] = None
+                else:
+                    out_rule[key][k] = randint(0, num_ir - 1)  # noqa: S311
             else:
-                out_rule[key] = randint(0, num_ir - 1)  # noqa: S311
-        else:
-            out_rule[key] = evaluate_random_rule(str(out_rule[key]))
+                out_rule[key][k] = evaluate_random_rule(str(out_rule[key][k]))
 
-        # convert EQ values from strings to numbers
-        if key in ["eq1", "eq2", "eq3"]:
-            for n in range(3):
-                if isinstance(out_rule[key][n], str):
-                    out_rule[key][n] = eval(out_rule[key][n])  # noqa: S307
+            # convert EQ values from strings to numbers
+            if k in ("eq1", "eq2", "eq3"):
+                for n in range(3):
+                    if isinstance(out_rule[key][k][n], str):
+                        out_rule[key][k][n] = eval(out_rule[key][k][n])  # noqa: S307
 
     return out_rule
 
@@ -141,7 +158,7 @@ def _rule_has_rand(rule: dict) -> bool:
     :param rule: Rule
     :return: True if rule contains 'rand'
     """
-    return any("rand" in str(rule[key]) for key in rule)
+    return any("rand" in str(rule[key][k]) for key in rule for k in rule[key])
 
 
 def estimate_augmented_length_from_length(length: int, tempo: OptionalNumberStr = None, frame_length: int = 1) -> int:
@@ -259,67 +276,165 @@ def _parse_ir(rule: dict, num_ir: int) -> dict:
             raise ValueError(f"Invalid ir entry of {rule_in}")
         return rule_out
 
-    if "ir" not in rule:
-        return rule
+    def _process(rule_in: dict) -> dict:
+        if "ir" not in rule_in:
+            return rule_in
 
-    ir = rule["ir"]
+        ir = rule_in["ir"]
 
-    if ir is None:
-        return rule
+        if ir is None:
+            return rule_in
 
-    if isinstance(ir, str):
-        rule["ir"] = _resolve_str(ir)
-        return rule
+        if isinstance(ir, str):
+            rule_in["ir"] = _resolve_str(ir)
+            return rule_in
 
-    if isinstance(ir, list):
-        rule["ir"] = []
-        for item in ir:
-            result = _resolve_str(item)
-            if isinstance(result, str):
-                rule["ir"].append(_resolve_str(item))
-            else:
-                rule["ir"] += _resolve_str(item)
+        if isinstance(ir, list):
+            rule_in["ir"] = []
+            for item in ir:
+                result = _resolve_str(item)
+                if isinstance(result, str):
+                    rule_in["ir"].append(_resolve_str(item))
+                else:
+                    rule_in["ir"] += _resolve_str(item)
 
-        return rule
+            return rule_in
+
+        if isinstance(ir, int):
+            if ir not in range(num_ir):
+                raise ValueError(f"Invalid ir of {ir}")
+            return rule_in
 
-    if isinstance(ir, int):
-        if ir not in range(num_ir):
-            raise ValueError(f"Invalid ir of {ir}")
-        return rule
+        raise ValueError(f"Invalid ir of {ir}")
+
+    for key in rule:
+        if key in ("pre", "post"):
+            rule[key] = _process(rule[key])
 
-    raise ValueError(f"Invalid ir of {ir}")
+    return rule
 
 
-def apply_augmentation(audio: AudioT, augmentation: Augmentation, frame_length: int = 1) -> AudioT:
-    """Apply augmentations to audio data
+def apply_augmentation(
+    mixdb: MixtureDatabase,
+    audio: AudioT,
+    augmentation: AugmentationEffects,
+    frame_length: int = 1,
+) -> AudioT:
+    """Apply augmentations to audio data using torchaudio.sox_effects
 
+    :param mixdb: Mixture database
     :param audio: Audio
     :param augmentation: Augmentation
     :param frame_length: Pad resulting audio to be a multiple of this
     :return: Augmented audio
     """
-    from .torchaudio_augmentation import apply_augmentation
+    import numpy as np
+    import torch
+    import torchaudio
+
+    from .audio import read_ir
+    from .constants import SAMPLE_RATE
+
+    effects: list[list[str]] = []
+
+    # TODO: Always normalize and remove normalize from list of available augmentations
+    # Normalize to globally set level (should this be a global config parameter, or hard-coded into the script?)
+    # TODO: Support all sox effects supported by torchaudio (torchaudio.sox_effects.effect_names())
+    if augmentation.normalize is not None:
+        effects.append(["norm", str(augmentation.normalize)])
+
+    if augmentation.gain is not None:
+        effects.append(["gain", str(augmentation.gain)])
 
-    return apply_augmentation(audio, augmentation, frame_length)
+    if augmentation.pitch is not None:
+        effects.append(["pitch", str(augmentation.pitch)])
+        effects.append(["rate", str(SAMPLE_RATE)])
+
+    if augmentation.tempo is not None:
+        effects.append(["tempo", "-s", str(augmentation.tempo)])
+
+    if augmentation.eq1 is not None:
+        effects.append(["equalizer", *[str(item) for item in augmentation.eq1]])
+
+    if augmentation.eq2 is not None:
+        effects.append(["equalizer", *[str(item) for item in augmentation.eq2]])
+
+    if augmentation.eq3 is not None:
+        effects.append(["equalizer", *[str(item) for item in augmentation.eq3]])
+
+    if augmentation.lpf is not None:
+        effects.append(["lowpass", "-2", str(augmentation.lpf), "0.707"])
+
+    if effects:
+        if audio.ndim == 1:
+            audio = np.reshape(audio, (1, audio.shape[0]))
+        out = torch.tensor(audio)
+
+        try:
+            out, _ = torchaudio.sox_effects.apply_effects_tensor(out, sample_rate=SAMPLE_RATE, effects=effects)
+        except Exception as e:
+            raise RuntimeError(f"Error applying {augmentation}: {e}") from e
+
+        audio_out = np.squeeze(np.array(out))
+    else:
+        audio_out = audio
+
+    if augmentation.ir is not None:
+        audio_out = apply_impulse_response(
+            audio=audio_out,
+            ir=read_ir(
+                name=mixdb.impulse_response_file(augmentation.ir),  # pyright: ignore [reportArgumentType]
+                delay=mixdb.impulse_response_delay(augmentation.ir),  # pyright: ignore [reportArgumentType]
+                use_cache=mixdb.use_cache,
+            ),
+        )
+
+    # make sure length is multiple of frame_length
+    return pad_audio_to_frame(audio=audio_out, frame_length=frame_length)
 
 
 def apply_impulse_response(audio: AudioT, ir: ImpulseResponseData) -> AudioT:
-    """Apply impulse response to audio data
+    """Apply impulse response to audio data using scipy
 
     :param audio: Audio
     :param ir: Impulse response data
     :return: Augmented audio
     """
-    from .torchaudio_augmentation import apply_impulse_response
+    import numpy as np
+    from librosa import resample
+    from scipy.signal import fftconvolve
+
+    from .constants import SAMPLE_RATE
+
+    # Early exit if no ir or if all audio is zero
+    if ir is None or not audio.any():
+        return audio
+
+    # Convert audio to IR sample rate
+    audio_in = resample(audio, orig_sr=SAMPLE_RATE, target_sr=ir.sample_rate, res_type="soxr_hq")
+    max_in = np.max(np.abs(audio_in))
+
+    # Apply IR
+    audio_out = fftconvolve(audio_in, ir.data, mode="full")
 
-    return apply_impulse_response(audio, ir)
+    # 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)]
+    max_out = np.max(np.abs(audio_out))
+
+    compensation_gain = max_in / max_out
+
+    return audio_out * compensation_gain
 
 
 def augmentation_from_rule(rule: AugmentationRule, num_ir: int) -> Augmentation:
     from sonusai.utils import dataclass_from_dict
 
-    from .datatypes import Augmentation
-
     processed_rule = rule.to_dict()
     del processed_rule["mixup"]
     processed_rule = _generate_none_rule(processed_rule)

sonusai/mixture/config.py

@@ -529,7 +529,7 @@ def get_impulse_response_files(config: dict) -> list[ImpulseResponseFile]:
     return list(
         chain.from_iterable(
             [
-                append_impulse_response_files(entry=ImpulseResponseFile(entry["name"], entry.get("tags", [])))
+                append_impulse_response_files(entry=ImpulseResponseFile(entry["name"], entry.get("tags", []), 0))
                 for entry in config["impulse_responses"]
             ]
         )
@@ -552,6 +552,7 @@ def append_impulse_response_files(entry: ImpulseResponseFile, tokens: dict | Non
     from os.path import splitext
 
     from .audio import validate_input_file
+    from .ir_delay import get_impulse_response_delay
     from .tokenized_shell_vars import tokenized_expand
     from .tokenized_shell_vars import tokenized_replace
 
@@ -572,7 +573,7 @@ def append_impulse_response_files(entry: ImpulseResponseFile, tokens: dict | Non
             for file in listdir(name):
                 if not isabs(file):
                     file = join(dir_name, file)
-                child = ImpulseResponseFile(file, entry.tags)
+                child = ImpulseResponseFile(file, entry.tags, get_impulse_response_delay(file))
                 impulse_response_files.extend(append_impulse_response_files(entry=child, tokens=tokens))
         else:
             try:
@@ -587,7 +588,7 @@ def append_impulse_response_files(entry: ImpulseResponseFile, tokens: dict | Non
                                 tokens.update(new_tokens)
                                 if not isabs(file):
                                     file = join(dir_name, file)
-                                child = ImpulseResponseFile(file, entry.tags)
+                                child = ImpulseResponseFile(file, entry.tags, get_impulse_response_delay(file))
                                 impulse_response_files.extend(append_impulse_response_files(entry=child, tokens=tokens))
                 elif ext == ".yml":
                     try:
@@ -602,7 +603,11 @@ def append_impulse_response_files(entry: ImpulseResponseFile, tokens: dict | Non
                         raise OSError(f"Error processing {name}: {e}") from e
                 else:
                     validate_input_file(name)
-                    impulse_response_files.append(ImpulseResponseFile(tokenized_replace(name, tokens), entry.tags))
+                    impulse_response_files.append(
+                        ImpulseResponseFile(
+                            tokenized_replace(name, tokens), entry.tags, get_impulse_response_delay(name)
+                        )
+                    )
             except Exception as e:
                 raise OSError(f"Error processing {name}: {e}") from e
 

sonusai/mixture/constants.py

@@ -38,7 +38,6 @@ VALID_AUGMENTATIONS = [
     "eq3",
     "lpf",
     "ir",
-    "mixup",
 ]
 VALID_NOISE_MIX_MODES = ["exhaustive", "non-exhaustive", "non-combinatorial"]
 RAND_PATTERN = re.compile(r"rand\(([-+]?(\d+(\.\d*)?|\.\d+)),\s*([-+]?(\d+(\.\d*)?|\.\d+))\)")

sonusai/mixture/datatypes.py

@@ -75,7 +75,7 @@ EQ: TypeAlias = tuple[float | int, float | int, float | int]
 
 
 @dataclass
-class AugmentationRule(DataClassSonusAIMixin):
+class AugmentationRuleEffects(DataClassSonusAIMixin):
     normalize: OptionalNumberStr = None
     pitch: OptionalNumberStr = None
     tempo: OptionalNumberStr = None
@@ -85,11 +85,17 @@ class AugmentationRule(DataClassSonusAIMixin):
     eq3: OptionalListNumberStr = None
     lpf: OptionalNumberStr = None
     ir: OptionalNumberStr = None
+
+
+@dataclass
+class AugmentationRule(DataClassSonusAIMixin):
+    pre: AugmentationRuleEffects
+    post: AugmentationRuleEffects | None = None
     mixup: int = 1
 
 
 @dataclass
-class Augmentation(DataClassSonusAIMixin):
+class AugmentationEffects(DataClassSonusAIMixin):
     normalize: float | None = None
     pitch: float | None = None
     tempo: float | None = None
@@ -101,6 +107,12 @@ class Augmentation(DataClassSonusAIMixin):
     ir: int | None = None
 
 
+@dataclass
+class Augmentation(DataClassSonusAIMixin):
+    pre: AugmentationEffects
+    post: AugmentationEffects
+
+
 @dataclass(frozen=True)
 class UniversalSNRGenerator:
     is_random: bool
@@ -191,19 +203,16 @@ class GenFTData:
 
 @dataclass
 class ImpulseResponseData:
-    name: str
-    sample_rate: int
     data: AudioT
-
-    @property
-    def length(self) -> int:
-        return len(self.data)
+    sample_rate: int
+    delay: int
 
 
 @dataclass
 class ImpulseResponseFile:
     file: str
     tags: list[str]
+    delay: int
 
 
 @dataclass(frozen=True)
@@ -230,9 +239,9 @@ class Target(DataClassSonusAIMixin):
     def gain(self) -> float:
         # gain is used to back out the gain augmentation in order to return the target audio
         # to its normalized level when calculating truth (if needed).
-        if self.augmentation.gain is None:
+        if self.augmentation.pre.gain is None:
             return 1.0
-        return round(10 ** (self.augmentation.gain / 20), ndigits=5)
+        return round(10 ** (self.augmentation.pre.gain / 20), ndigits=5)
 
 
 Targets: TypeAlias = list[Target]