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

sonusai/metrics/calc_audio_stats.py

@@ -1,3 +1,8 @@
+from pathlib import Path
+
+import numpy as np
+from sox import Transformer as SoxTransformer
+
 from sonusai.mixture.datatypes import AudioStatsMetrics
 from sonusai.mixture.datatypes import AudioT
 
@@ -12,7 +17,6 @@ def _convert_str_with_factors_to_int(x: str) -> int:
 
 def calc_audio_stats(audio: AudioT, win_len: float | None = None) -> AudioStatsMetrics:
     from sonusai.mixture import SAMPLE_RATE
-    from sonusai.mixture import Transformer
 
     args = ["stats"]
     if win_len is not None:
@@ -53,3 +57,275 @@ def calc_audio_stats(audio: AudioT, win_len: float | None = None) -> AudioStatsM
         fl=float(stats["Flat factor"]),
         pkc=_convert_str_with_factors_to_int(stats["Pk count"]),
     )
+
+
+class Transformer(SoxTransformer):
+    """Override certain sox.Transformer methods"""
+
+    def build(  # pyright: ignore [reportIncompatibleMethodOverride]
+        self,
+        input_filepath: str | Path | None = None,
+        output_filepath: str | Path | None = None,
+        input_array: np.ndarray | None = None,
+        sample_rate_in: float | None = None,
+        extra_args: list[str] | None = None,
+        return_output: bool = False,
+    ) -> tuple[bool, str | None, str | None]:
+        """Given an input file or array, creates an output_file on disk by
+        executing the current set of commands. This function returns True on
+        success. If return_output is True, this function returns a triple of
+        (status, out, err), giving the success state, along with stdout and
+        stderr returned by sox.
+
+        Parameters
+        ----------
+        input_filepath : str or None
+            Either path to input audio file or None for array input.
+        output_filepath : str
+            Path to desired output file. If a file already exists at
+            the given path, the file will be overwritten.
+            If '-n', no file is created.
+        input_array : np.ndarray or None
+            An np.ndarray of an waveform with shape (n_samples, n_channels).
+            sample_rate_in must also be provided.
+            If None, input_filepath must be specified.
+        sample_rate_in : int
+            Sample rate of input_array.
+            This argument is ignored if input_array is None.
+        extra_args : list or None, default=None
+            If a list is given, these additional arguments are passed to SoX
+            at the end of the list of effects.
+            Don't use this argument unless you know exactly what you're doing!
+        return_output : bool, default=False
+            If True, returns the status and information sent to stderr and
+            stdout as a tuple (status, stdout, stderr).
+            If output_filepath is None, return_output=True by default.
+            If False, returns True on success.
+
+        Returns
+        -------
+        status : bool
+            True on success.
+        out : str (optional)
+            This is not returned unless return_output is True.
+            When returned, captures the stdout produced by sox.
+        err : str (optional)
+            This is not returned unless return_output is True.
+            When returned, captures the stderr produced by sox.
+
+        Examples
+        --------
+        > import numpy as np
+        > import sox
+        > tfm = sox.Transformer()
+        > sample_rate = 44100
+        > y = np.sin(2 * np.pi * 440.0 * np.arange(sample_rate * 1.0) / sample_rate)
+
+        file in, file out - basic usage
+
+        > status = tfm.build('path/to/input.wav', 'path/to/output.mp3')
+
+        file in, file out - equivalent usage
+
+        > status = tfm.build(
+                input_filepath='path/to/input.wav',
+                output_filepath='path/to/output.mp3'
+            )
+
+        array in, file out
+
+        > status = tfm.build(
+                input_array=y, sample_rate_in=sample_rate,
+                output_filepath='path/to/output.mp3'
+            )
+
+        """
+        from sox import file_info
+        from sox.core import SoxError
+        from sox.core import sox
+        from sox.log import logger
+
+        input_format, input_filepath = self._parse_inputs(input_filepath, input_array, sample_rate_in)
+
+        if output_filepath is None:
+            raise ValueError("output_filepath is not specified!")
+
+        # set output parameters
+        if input_filepath == output_filepath:
+            raise ValueError("input_filepath must be different from output_filepath.")
+        file_info.validate_output_file(output_filepath)
+
+        args = []
+        args.extend(self.globals)
+        args.extend(self._input_format_args(input_format))
+        args.append(input_filepath)
+        args.extend(self._output_format_args(self.output_format))
+        args.append(output_filepath)
+        args.extend(self.effects)
+
+        if extra_args is not None:
+            if not isinstance(extra_args, list):
+                raise ValueError("extra_args must be a list.")
+            args.extend(extra_args)
+
+        status, out, err = sox(args, input_array, True)
+        if status != 0:
+            raise SoxError(f"Stdout: {out}\nStderr: {err}")
+
+        logger.info("Created %s with effects: %s", output_filepath, " ".join(self.effects_log))
+
+        if return_output:
+            return status, out, err  # pyright: ignore [reportReturnType]
+
+        return True, None, None
+
+    def build_array(  # pyright: ignore [reportIncompatibleMethodOverride]
+        self,
+        input_filepath: str | Path | None = None,
+        input_array: np.ndarray | None = None,
+        sample_rate_in: int | None = None,
+        extra_args: list[str] | None = None,
+    ) -> np.ndarray:
+        """Given an input file or array, returns the output as a numpy array
+        by executing the current set of commands. By default, the array will
+        have the same sample rate as the input file unless otherwise specified
+        using set_output_format. Functions such as channels and convert
+        will be ignored!
+
+        The SonusAI override does not generate a warning for rate transforms.
+
+        Parameters
+        ----------
+        input_filepath : str, Path or None
+            Either path to input audio file or None.
+        input_array : np.ndarray or None
+            A np.ndarray of a waveform with shape (n_samples, n_channels).
+            If this argument is passed, sample_rate_in must also be provided.
+            If None, input_filepath must be specified.
+        sample_rate_in : int
+            Sample rate of input_array.
+            This argument is ignored if input_array is None.
+        extra_args : list or None, default=None
+            If a list is given, these additional arguments are passed to SoX
+            at the end of the list of effects.
+            Don't use this argument unless you know exactly what you're doing!
+
+        Returns
+        -------
+        output_array : np.ndarray
+            Output audio as a numpy array
+
+        Examples
+        --------
+
+        > import numpy as np
+        > import sox
+        > tfm = sox.Transformer()
+        > sample_rate = 44100
+        > y = np.sin(2 * np.pi * 440.0 * np.arange(sample_rate * 1.0) / sample_rate)
+
+        file in, array out
+
+        > output_array = tfm.build(input_filepath='path/to/input.wav')
+
+        array in, array out
+
+        > output_array = tfm.build(input_array=y, sample_rate_in=sample_rate)
+
+        specifying the output sample rate
+
+        > tfm.set_output_format(rate=8000)
+        > output_array = tfm.build(input_array=y, sample_rate_in=sample_rate)
+
+        if an effect changes the number of channels, you must explicitly
+        specify the number of output channels
+
+        > tfm.remix(remix_dictionary={1: [1], 2: [1], 3: [1]})
+        > tfm.set_output_format(channels=3)
+        > output_array = tfm.build(input_array=y, sample_rate_in=sample_rate)
+
+
+        """
+        from sox.core import SoxError
+        from sox.core import sox
+        from sox.log import logger
+        from sox.transform import ENCODINGS_MAPPING
+
+        input_format, input_filepath = self._parse_inputs(input_filepath, input_array, sample_rate_in)
+
+        # check if any of the below commands are part of the effects chain
+        ignored_commands = ["channels", "convert"]
+        if set(ignored_commands) & set(self.effects_log):
+            logger.warning(
+                "When outputting to an array, channels and convert "
+                + "effects may be ignored. Use set_output_format() to "
+                + "specify output formats."
+            )
+
+        output_filepath = "-"
+
+        if input_format.get("file_type") is None:
+            encoding_out = np.int16
+        else:
+            encoding_out = next(k for k, v in ENCODINGS_MAPPING.items() if input_format["file_type"] == v)
+
+        n_bits = np.dtype(encoding_out).itemsize * 8
+
+        output_format = {
+            "file_type": "raw",
+            "rate": sample_rate_in,
+            "bits": n_bits,
+            "channels": input_format["channels"],
+            "encoding": None,
+            "comments": None,
+            "append_comments": True,
+        }
+
+        if self.output_format.get("rate") is not None:
+            output_format["rate"] = self.output_format["rate"]
+
+        if self.output_format.get("channels") is not None:
+            output_format["channels"] = self.output_format["channels"]
+
+        if self.output_format.get("bits") is not None:
+            n_bits = self.output_format["bits"]
+            output_format["bits"] = n_bits
+
+        match n_bits:
+            case 8:
+                encoding_out = np.int8  # type: ignore[assignment]
+            case 16:
+                encoding_out = np.int16
+            case 32:
+                encoding_out = np.float32  # type: ignore[assignment]
+            case 64:
+                encoding_out = np.float64  # type: ignore[assignment]
+            case _:
+                raise ValueError(f"invalid n_bits {n_bits}")
+
+        args = []
+        args.extend(self.globals)
+        args.extend(self._input_format_args(input_format))
+        args.append(input_filepath)
+        args.extend(self._output_format_args(output_format))
+        args.append(output_filepath)
+        args.extend(self.effects)
+
+        if extra_args is not None:
+            if not isinstance(extra_args, list):
+                raise ValueError("extra_args must be a list.")
+            args.extend(extra_args)
+
+        status, out, err = sox(args, input_array, False)
+        if status != 0:
+            raise SoxError(f"Stdout: {out}\nStderr: {err}")
+
+        out = np.frombuffer(out, dtype=encoding_out)  # pyright: ignore [reportArgumentType, reportCallIssue]
+        if output_format["channels"] > 1:
+            out = out.reshape(
+                (output_format["channels"], int(len(out) / output_format["channels"])),
+                order="F",
+            ).T
+        logger.info("Created array with effects: %s", " ".join(self.effects_log))
+
+        return out

sonusai/metrics_summary.py

@@ -137,7 +137,9 @@ def main() -> None:
         print(f"Could not open SonusAI mixture database in {location}, exiting ...")
         return
 
-    metrics_present = mixdb.cached_metrics()
+    # Only check first and last mixture in order to save time
+    metrics_present = mixdb.cached_metrics([0, mixdb.num_mixtures - 1])
+
     num_metrics_present = len(metrics_present)
     if num_metrics_present < 1:
         print(f"mixdb reports no pre-generated metrics are present. Nothing to summarize in {location}, exiting ...")
@@ -150,7 +152,7 @@ def main() -> None:
         create_file_handler(join(location, "metrics_summary.log"))
         update_console_handler(verbose)
         initial_log_messages("metrics_summary")
-        logger.info(f"Logging summary of SonusAI mixture db at {location}")
+        logger.info(f"Logging summary of SonusAI mixture database at {location}")
     else:
         update_console_handler(verbose)
 
@@ -164,7 +166,7 @@ def main() -> None:
         fsuffix = f"_s{len(mixids)}t{mixdb.num_mixtures}"
     else:
         logger.info(
-            f"Summarizing SonusAI mixture db with {mixdb.num_mixtures} mixtures "
+            f"Summarizing SonusAI mixture database with {mixdb.num_mixtures} mixtures "
             f"and {num_metrics_present} pre-generated metrics ..."
         )
         fsuffix = ""

sonusai/mixture/__init__.py

@@ -5,6 +5,7 @@ from .audio import get_duration
 from .audio import get_next_noise
 from .audio import get_num_samples
 from .audio import get_sample_rate
+from .audio import raw_read_audio
 from .audio import read_audio
 from .audio import read_ir
 from .audio import validate_input_file
@@ -53,7 +54,9 @@ from .datatypes import AudioF
 from .datatypes import AudioStatsMetrics
 from .datatypes import AudioT
 from .datatypes import Augmentation
+from .datatypes import AugmentationEffects
 from .datatypes import AugmentationRule
+from .datatypes import AugmentationRuleEffects
 from .datatypes import AugmentedTarget
 from .datatypes import ClassCount
 from .datatypes import EnergyF
@@ -111,10 +114,10 @@ from .helpers import get_transform_from_audio
 from .helpers import inverse_transform
 from .helpers import mixture_metadata
 from .helpers import write_mixture_metadata
+from .ir_delay import get_impulse_response_delay
 from .log_duration_and_sizes import log_duration_and_sizes
 from .mixdb import MixtureDatabase
 from .mixdb import db_file
-from .sox_audio import Transformer
 from .spectral_mask import apply_spectral_mask
 from .target_class_balancing import balance_targets
 from .targets import get_augmented_target_ids_by_class

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)