sonusai 1.0.16__cp311-abi3-macosx_10_12_x86_64.whl

sonusai/metrics_summary.py

@@ -0,0 +1,314 @@
+"""sonusai metrics_summary
+
+usage: lsdb [-vlh] [-i MIXID] [-n NCPU] LOCATION
+
+Options:
+    -h, --help
+    -v, --verbose
+    -l, --write-list            Write .csv file list of all mixture metrics
+    -i MIXID, --mixid MIXID     Mixture ID(s) to analyze. [default: *].
+    -n, --num_process NCPU      Number of parallel processes to use [default: auto]
+
+Summarize mixture metrics across a SonusAI mixture database where metrics have been generated by SonusAI genmetrics.
+
+Inputs:
+    LOCATION     A SonusAI mixture database directory with mixdb.db and pre-generated metrics from SonusAI genmetrics.
+
+"""
+
+import numpy as np
+import pandas as pd
+
+DB_99 = np.power(10, 99 / 10)
+DB_N99 = np.power(10, -99 / 10)
+
+
+def _process_mixture(
+    m_id: int,
+    location: str,
+    all_metric_names: list[str],
+    scalar_metric_names: list[str],
+    string_metric_names: list[str],
+    frame_metric_names: list[str],
+    bin_metric_names: list[str],
+    ptab_labels: list[str],
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+    from os.path import basename
+
+    from sonusai.constants import SAMPLE_RATE
+    from sonusai.metrics import calc_wer
+    from sonusai.mixture import MixtureDatabase
+
+    mixdb = MixtureDatabase(location)
+
+    # Process mixture
+    # for mixid in mixids:
+    samples = mixdb.mixture(m_id).samples
+    duration = samples / SAMPLE_RATE
+    tf_frames = mixdb.mixture_transform_frames(m_id)
+    feat_frames = mixdb.mixture_feature_frames(m_id)
+    mxsnr = mixdb.mixture(m_id).noise.snr
+    ti = mixdb.mixture(m_id).sources["primary"].file_id
+    ni = mixdb.mixture(m_id).noise.file_id
+    t0file = basename(mixdb.source_file(ti).name)
+    nfile = basename(mixdb.source_file(ni).name)
+
+    all_metrics = mixdb.mixture_metrics(m_id, all_metric_names)
+
+    # replace dict with 'primary' value (ignore mixup)
+    scalar_metrics = {
+        key: all_metrics[key]["primary"] if isinstance(all_metrics[key], dict) else all_metrics[key]
+        for key in scalar_metric_names
+    }
+    string_metrics = {
+        key: all_metrics[key]["primary"] if isinstance(all_metrics[key], dict) else all_metrics[key]
+        for key in string_metric_names
+    }
+
+    # Convert strings into word count
+    for key in string_metrics:
+        string_metrics[key] = calc_wer(string_metrics[key], string_metrics[key]).words
+
+    # Collect pandas table values note: must match given ptab_labels
+    ptab_data: list = [
+        mxsnr,
+        *scalar_metrics.values(),
+        *string_metrics.values(),
+        tf_frames,
+        duration,
+        t0file,
+        nfile,
+    ]
+
+    ptab1 = pd.DataFrame([ptab_data], columns=ptab_labels, index=[m_id])
+
+    # TODO: collect frame metrics and bin metrics
+
+    return ptab1, ptab1
+
+
+def main() -> None:
+    from docopt import docopt
+
+    from sonusai import __version__ as sai_version
+    from sonusai.utils.docstring import trim_docstring
+
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
+
+    verbose = args["--verbose"]
+    wrlist = args["--write-list"]
+    mixids = args["--mixid"]
+    location = args["LOCATION"]
+    num_proc = args["--num_process"]
+
+    from functools import partial
+    from os.path import basename
+    from os.path import join
+
+    import psutil
+
+    from sonusai import create_file_handler
+    from sonusai import initial_log_messages
+    from sonusai import logger
+    from sonusai import update_console_handler
+    from sonusai.mixture import MixtureDatabase
+    from sonusai.utils.create_timestamp import create_timestamp
+    from sonusai.utils.parallel import par_track
+    from sonusai.utils.parallel import track
+
+    mixdb = MixtureDatabase(location)
+    print(f"Found SonusAI mixture database with {mixdb.num_mixtures} mixtures.")
+
+    # Only check first and last mixture in order to save time
+    metrics_present = mixdb.cached_metrics([0, mixdb.num_mixtures - 1])  # return pre-generated metrics in mixdb tree
+    if "mxsnr" in metrics_present:
+        metrics_present.remove("mxsnr")
+
+    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 ...")
+        return
+
+    # Setup logging file
+    timestamp = create_timestamp()  # string good for embedding into filenames
+    mixdb_fname = basename(location)
+    if verbose:
+        create_file_handler(join(location, "metrics_summary.log"), verbose)
+        update_console_handler(verbose)
+        initial_log_messages("metrics_summary")
+        logger.info(f"Logging summary of SonusAI mixture database at {location}")
+    else:
+        update_console_handler(verbose)
+
+    logger.info("")
+    mixids = mixdb.mixids_to_list(mixids)
+    if len(mixids) < mixdb.num_mixtures:
+        logger.info(
+            f"Processing a subset of {len(mixids)} out of total mixdb mixtures of {mixdb.num_mixtures}, "
+            f"summary results will not include entire dataset."
+        )
+        fsuffix = f"_s{len(mixids)}t{mixdb.num_mixtures}"
+    else:
+        logger.info(
+            f"Summarizing SonusAI mixture database with {mixdb.num_mixtures} mixtures "
+            f"and {num_metrics_present} pre-generated metrics ..."
+        )
+        fsuffix = ""
+
+    metric_sup = mixdb.supported_metrics
+    ft_bins = mixdb.ft_config.bin_end - mixdb.ft_config.bin_start + 1  # bins of forward transform
+    # Pre-process first mixid to gather metrics into 4 types: scalar, str (scalar word cnt), frame-array, bin-array
+    # Collect list of indices for each
+    scalar_metric_names: list[str] = []
+    string_metric_names: list[str] = []
+    frame_metric_names: list[str] = []
+    bin_metric_names: list[str] = []
+    all_metrics = mixdb.mixture_metrics(mixids[0], metrics_present)
+    tf_frames = mixdb.mixture_transform_frames(mixids[0])
+    for metric in metrics_present:
+        metval = all_metrics[metric]  # get metric value
+        logger.debug(f"First mixid {mixids[0]} metric {metric} = {metval}")
+        if isinstance(metval, dict):
+            logger.warning(f"Mixid {mixids[0]} metric {metric} is a dict, using 'primary'.")
+            metval = metval["primary"]  # remove any dict
+        if isinstance(metval, float | int):
+            logger.debug(f"Metric is scalar {type(metval)}, entering in summary table.")
+            scalar_metric_names.append(metric)
+        elif isinstance(metval, str):
+            logger.debug("Metric is string, will summarize with word count.")
+            string_metric_names.append(metric)
+        elif isinstance(metval, np.ndarray):
+            if metval.ndim == 1:
+                if metval.size == tf_frames:
+                    logger.debug("Metric is frames vector.")
+                    frame_metric_names.append(metric)
+                elif metval.size == ft_bins:
+                    logger.debug("Metric is bins vector.")
+                    bin_metric_names.append(metric)
+                else:
+                    logger.warning(f"Mixid {mixids[0]} metric {metric} is a vector of improper size, ignoring.")
+
+    # Setup pandas table for summarizing scalar metrics, always include mxsnr first
+    ptab_labels = [
+        "mxsnr",
+        *scalar_metric_names,
+        *string_metric_names,
+        "fcnt",
+        "duration",
+        "t0file",
+        "nfile",
+    ]
+
+    num_cpu = psutil.cpu_count()
+    cpu_percent = psutil.cpu_percent(interval=1)
+    logger.info("")
+    logger.info(f"#CPUs: {num_cpu}, current CPU utilization: {cpu_percent}%")
+    logger.info(f"Memory utilization: {psutil.virtual_memory().percent}%")
+    if num_proc == "auto":
+        use_cpu = int(num_cpu * (0.9 - cpu_percent / 100))  # default use 80% of available cpus
+    elif num_proc == "None":
+        use_cpu = None
+    else:
+        use_cpu = min(max(int(num_proc), 1), num_cpu)
+
+    logger.info(f"Summarizing metrics for {len(mixids)} mixtures using {use_cpu} parallel processes")
+
+    # progress = tqdm(total=len(mixids), desc='calc_metric_spenh', mininterval=1)
+    progress = track(total=len(mixids))
+    if use_cpu is None:
+        no_par = True
+        num_cpus = None
+    else:
+        no_par = False
+        num_cpus = use_cpu
+
+    all_metrics_tables = par_track(
+        partial(
+            _process_mixture,
+            location=location,
+            all_metric_names=metrics_present,
+            scalar_metric_names=scalar_metric_names,
+            string_metric_names=string_metric_names,
+            frame_metric_names=frame_metric_names,
+            bin_metric_names=bin_metric_names,
+            ptab_labels=ptab_labels,
+        ),
+        mixids,
+        progress=progress,
+        num_cpus=num_cpus,
+        no_par=no_par,
+    )
+    progress.close()
+
+    # Done with mixtures, write out summary metrics
+    header_args = {
+        "mode": "a",
+        "encoding": "utf-8",
+        "index": False,
+        "header": False,
+    }
+    table_args = {
+        "mode": "a",
+        "encoding": "utf-8",
+    }
+    ptab1 = pd.concat([item[0] for item in all_metrics_tables])
+    if wrlist:
+        wlcsv_name = str(join(location, "metric_summary_list" + fsuffix + ".csv"))
+        pd.DataFrame([["Timestamp", timestamp]]).to_csv(wlcsv_name, header=False, index=False)
+        pd.DataFrame([f"Metric list for {mixdb_fname}:"]).to_csv(wlcsv_name, mode="a", header=False, index=False)
+        ptab1.round(2).to_csv(wlcsv_name, **table_args)
+    ptab1_sorted = ptab1.sort_values(by=["mxsnr", "t0file"])
+
+    # Create metrics table except -99 SNR
+    ptab1_nom99 = ptab1_sorted[ptab1_sorted.mxsnr != -99]
+
+    # Create summary by SNR for all scalar metrics, taking mean
+    mtab_snr_summary = None
+    for snri in range(0, len(mixdb.snrs)):
+        tmp = ptab1_sorted.query("mxsnr==" + str(mixdb.snrs[snri])).mean(numeric_only=True).to_frame().T
+        # avoid nan when subset of mixids specified (i.e. no mixtures exist for an SNR)
+        if ~np.isnan(tmp.iloc[0].to_numpy()[0]).any():
+            mtab_snr_summary = pd.concat([mtab_snr_summary, tmp])
+    mtab_snr_summary = mtab_snr_summary.sort_values(by=["mxsnr"], ascending=False)
+
+    # Write summary to .csv
+    snrcsv_name = str(join(location, "metric_summary_snr" + fsuffix + ".csv"))
+    nmix = len(mixids)
+    nmixtot = mixdb.num_mixtures
+    pd.DataFrame([["Timestamp", timestamp]]).to_csv(snrcsv_name, header=False, index=False)
+    pd.DataFrame(['"Metrics avg over each SNR:"']).to_csv(snrcsv_name, **header_args)
+    mtab_snr_summary.round(2).T.to_csv(snrcsv_name, index=True, header=False, mode="a", encoding="utf-8")
+    pd.DataFrame(["--"]).to_csv(snrcsv_name, header=False, index=False, mode="a")
+    pd.DataFrame([f'"Metrics stats over {nmix} mixtures out of {nmixtot} total:"']).to_csv(snrcsv_name, **header_args)
+    ptab1.describe().round(2).T.to_csv(snrcsv_name, index=True, **table_args)
+    pd.DataFrame(["--"]).to_csv(snrcsv_name, header=False, index=False, mode="a")
+    pd.DataFrame([f'"Metrics stats over {len(ptab1_nom99)} non -99db mixtures out of {nmixtot} total:"']).to_csv(
+        snrcsv_name, **header_args
+    )
+    ptab1_nom99.describe().round(2).T.to_csv(snrcsv_name, index=True, **table_args)
+
+    # Write summary to text file
+    snrtxt_name = str(join(location, "metric_summary_snr" + fsuffix + ".txt"))
+    with open(snrtxt_name, "w") as f:
+        print(f"Timestamp: {timestamp}", file=f)
+        print("Metrics avg over each SNR:", file=f)
+        print(
+            mtab_snr_summary.round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True, header=False), file=f
+        )
+        print("", file=f)
+        print(f"Metrics stats over {len(mixids)} mixtures out of {mixdb.num_mixtures} total:", file=f)
+        print(ptab1.describe().round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True), file=f)
+        print("", file=f)
+        print(f"Metrics stats over {len(ptab1_nom99)} non -99db mixtures out of {mixdb.num_mixtures} total:", file=f)
+        print(ptab1_nom99.describe().round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True), file=f)
+
+
+if __name__ == "__main__":
+    from sonusai import exception_handler
+    from sonusai.utils.keyboard_interrupt import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)

sonusai/mixture/__init__.py

@@ -0,0 +1,15 @@
+# SonusAI mixture utilities
+
+from .feature import get_audio_from_feature
+from .feature import get_feature_from_audio
+from .helpers import forward_transform
+from .helpers import inverse_transform
+from .mixdb import MixtureDatabase
+
+__all__ = [
+    "MixtureDatabase",
+    "forward_transform",
+    "get_audio_from_feature",
+    "get_feature_from_audio",
+    "inverse_transform",
+]

sonusai/mixture/audio.py

@@ -0,0 +1,187 @@
+from functools import lru_cache
+from pathlib import Path
+
+from ..datatypes import AudioT
+
+
+def get_next_noise(audio: AudioT, offset: int, length: int) -> AudioT:
+    """Get the next sequence of noise data from noise audio
+
+    :param audio: Overall noise audio (entire file's worth of data)
+    :param offset: Starting sample
+    :param length: Number of samples to get
+    :return: Sequence of noise audio data
+    """
+    import numpy as np
+
+    return np.take(audio, range(offset, offset + length), mode="wrap")
+
+
+def get_duration(audio: AudioT) -> float:
+    """Get duration of audio in seconds
+
+    :param audio: Time domain data [samples]
+    :return: Duration of audio in seconds
+    """
+    from ..constants import SAMPLE_RATE
+
+    return len(audio) / SAMPLE_RATE
+
+
+def validate_input_file(input_filepath: str | Path) -> None:
+    from os.path import exists
+    from os.path import splitext
+
+    from soundfile import available_formats
+
+    if not exists(input_filepath):
+        raise OSError(f"input_filepath {input_filepath} does not exist.")
+
+    ext = splitext(input_filepath)[1][1:].lower()
+    read_formats = [item.lower() for item in available_formats()]
+    if ext not in read_formats:
+        raise OSError(f"This installation cannot process .{ext} files")
+
+
+def get_sample_rate(name: str | Path, use_cache: bool = True) -> int:
+    """Get sample rate from audio file
+
+    :param name: File name
+    :param use_cache: If true, use LRU caching
+    :return: Sample rate
+    """
+    if use_cache:
+        return _get_sample_rate(name)
+    return _get_sample_rate.__wrapped__(name)
+
+
+@lru_cache
+def _get_sample_rate(name: str | Path) -> int:
+    """Get sample rate from audio file using soundfile
+
+    :param name: File name
+    :return: Sample rate
+    """
+    import soundfile
+    from pydub import AudioSegment
+
+    from ..utils.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 ..utils.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:
+    """Read audio data from a file
+
+    :param name: File name
+    :param use_cache: If true, use LRU caching
+    :return: Array of time domain audio data
+    """
+    if use_cache:
+        return _read_audio(name)
+    return _read_audio.__wrapped__(name)
+
+
+@lru_cache
+def _read_audio(name: str | Path) -> AudioT:
+    """Read audio data from a file using soundfile
+
+    :param name: File name
+    :return: Array of time domain audio data
+    """
+    from ..constants import SAMPLE_RATE
+    from .resample import resample
+
+    out, sample_rate = raw_read_audio(name)
+
+    return resample(out, orig_sr=sample_rate, target_sr=SAMPLE_RATE)
+
+
+def get_num_samples(name: str | Path, use_cache: bool = True) -> int:
+    """Get the number of samples resampled to the SonusAI sample rate in the given file
+
+    :param name: File name
+    :param use_cache: If true, use LRU caching
+    :return: number of samples in resampled audio
+    """
+    if use_cache:
+        return _get_num_samples(name)
+    return _get_num_samples.__wrapped__(name)
+
+
+@lru_cache
+def _get_num_samples(name: str | Path) -> int:
+    """Get the number of samples resampled to the SonusAI sample rate in the given file
+
+    :param name: File name
+    :return: number of samples in resampled audio
+    """
+    import math
+
+    import soundfile
+    from pydub import AudioSegment
+
+    from ..constants import SAMPLE_RATE
+    from ..utils.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(expanded_name)
+        samples = info.frames
+        sample_rate = info.samplerate
+
+    return math.ceil(SAMPLE_RATE * samples / sample_rate)

sonusai/mixture/class_balancing.py

@@ -0,0 +1,103 @@
+from ..datatypes import EffectList
+from ..datatypes import EffectedFile
+from ..datatypes import File
+
+
+def balance_sources(
+    effected_sources: list[EffectedFile],
+    files: list[File],
+    effects: list[EffectList],
+    class_balancing_effect: EffectList,
+    num_classes: int,
+    num_ir: int,
+    mixups: list[int] | None = None,
+) -> tuple[list[EffectedFile], list[EffectList]]:
+    import math
+
+    from .augmentation import get_mixups
+    from .sources import get_augmented_target_ids_by_class
+
+    first_cba_id = len(effects)
+
+    if mixups is None:
+        mixups = get_mixups(effects)
+
+    for mixup in mixups:
+        if mixup == 1:
+            continue
+
+        effected_sources_indices_by_class = get_augmented_target_ids_by_class(
+            augmented_targets=effected_sources,
+            targets=files,
+            target_augmentations=effects,
+            mixup=mixup,
+            num_classes=num_classes,
+        )
+
+        largest = max([len(item) for item in effected_sources_indices_by_class])
+        largest = math.ceil(largest / mixup) * mixup
+        for es_indices in effected_sources_indices_by_class:
+            additional_effects_needed = largest - len(es_indices)
+            file_ids = sorted({effected_sources[at_index].file_id for at_index in es_indices})
+
+            tfi_idx = 0
+            for _ in range(additional_effects_needed):
+                file_id = file_ids[tfi_idx]
+                tfi_idx = (tfi_idx + 1) % len(file_ids)
+                effect_id, effects = _get_unused_balancing_effect(
+                    effected_sources=effected_sources,
+                    files=files,
+                    effects=effects,
+                    class_balancing_effect=class_balancing_effect,
+                    file_id=file_id,
+                    mixup=mixup,
+                    num_ir=num_ir,
+                    first_cbe_id=first_cba_id,
+                )
+                effected_sources.append(EffectedFile(file_id=file_id, effect_id=effect_id))
+
+    return effected_sources, effects
+
+
+def _get_unused_balancing_effect(
+    effected_sources: list[EffectedFile],
+    files: list[File],
+    effects: list[EffectList],
+    class_balancing_effect: EffectList,
+    file_id: int,
+    mixup: int,
+    num_ir: int,
+    first_cbe_id: int,
+) -> tuple[int, list[EffectList]]:
+    """Get an unused balancing augmentation for a given target file index"""
+    from dataclasses import asdict
+
+    from .augmentation import get_augmentation_rules
+
+    balancing_augmentations = [item for item in range(len(effects)) if item >= first_cbe_id]
+    used_balancing_augmentations = [
+        effected_source.effect_id
+        for effected_source in effected_sources
+        if effected_source.file_id == file_id and effected_source.effect_id in balancing_augmentations
+    ]
+
+    augmentation_indices = [
+        item
+        for item in balancing_augmentations
+        if item not in used_balancing_augmentations and effects[item].mixup == mixup
+    ]
+    if len(augmentation_indices) > 0:
+        return augmentation_indices[0], effects
+
+    class_balancing_effect = get_class_balancing_effect(file=files[file_id], default_cbe=class_balancing_effect)
+    new_effect = get_augmentation_rules(rules=asdict(class_balancing_effect), num_ir=num_ir)[0]
+    new_effect.mixup = mixup
+    effects.append(new_effect)
+    return len(effects) - 1, effects
+
+
+def get_class_balancing_effect(file: File, default_cbe: EffectList) -> EffectList:
+    """Get the class balancing effect rule for the given target"""
+    if file.class_balancing_effect is not None:
+        return file.class_balancing_effect
+    return default_cbe

sonusai/mixture/constants.py

@@ -0,0 +1,3 @@
+MIXDB_VERSION = 3
+MIXDB_NAME = "mixdb.db"
+TEST_MIXDB_NAME = "mixdb_test.db"