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

sonusai/metrics_summary.py

@@ -0,0 +1,322 @@
+"""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 signal
+
+import numpy as np
+import pandas as pd
+
+
+def signal_handler(_sig, _frame):
+    import sys
+
+    from sonusai import logger
+
+    logger.info("Canceled due to keyboard interrupt")
+    sys.exit(1)
+
+
+signal.signal(signal.SIGINT, signal_handler)
+
+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.metrics import calc_wer
+    from sonusai.mixture import SAMPLE_RATE
+    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).snr
+    ti = mixdb.mixture(m_id).targets[0].file_id
+    ni = mixdb.mixture(m_id).noise.file_id
+    t0file = basename(mixdb.target_file(ti).name)
+    nfile = basename(mixdb.noise_file(ni).name)
+
+    all_metrics = mixdb.mixture_metrics(m_id, all_metric_names)
+
+    # replace lists with first value (ignore mixup)
+    scalar_metrics = {
+        key: all_metrics[key][0] if isinstance(all_metrics[key], list) else all_metrics[key]
+        for key in scalar_metric_names
+    }
+    string_metrics = {
+        key: all_metrics[key][0] if isinstance(all_metrics[key], list) 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 sonusai_ver
+    from sonusai.utils import trim_docstring
+
+    args = docopt(trim_docstring(__doc__), version=sonusai_ver, 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 import create_timestamp
+    from sonusai.utils import par_track
+    from sonusai.utils import track
+
+    try:
+        mixdb = MixtureDatabase(location)
+        print(f"Found SonusAI mixture database with {mixdb.num_mixtures} mixtures.")
+    except:
+        print(f"Could not open SonusAI mixture database in {location}, exiting ...")
+        return
+
+    # 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 ...")
+        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"))
+        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, list):
+            if len(metval) > 1:
+                logger.warning(f"Mixid {mixids[0]} metric {metric} has a list with more than 1 element, using first.")
+            metval = metval[0]  # remove any list
+        if isinstance(metval, float):
+            logger.debug("Metric is scalar float, 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
+    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 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).to_csv(snrcsv_name, index=False, **table_args)
+    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 .csv
+    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).to_string(float_format=lambda x: f"{x:.2f}", index=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__":
+    main()

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
@@ -87,6 +90,7 @@ from .datatypes import TruthParameter
 from .datatypes import UniversalSNR
 from .feature import get_audio_from_feature
 from .feature import get_feature_from_audio
+from .generation import generate_mixtures
 from .generation import get_all_snrs_from_config
 from .generation import initialize_db
 from .generation import populate_class_label_table
@@ -99,7 +103,7 @@ from .generation import populate_target_file_table
 from .generation import populate_top_table
 from .generation import populate_truth_parameters_table
 from .generation import update_mixid_width
-from .generation import update_mixture_table
+from .generation import update_mixture
 from .helpers import augmented_noise_samples
 from .helpers import augmented_target_samples
 from .helpers import check_audio_files_exist
@@ -110,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

@@ -44,49 +44,173 @@ def validate_input_file(input_filepath: str | Path) -> None:
         raise OSError(f"This installation cannot process .{ext} files")
 
 
-@lru_cache
-def get_sample_rate(name: str | Path) -> int:
+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
     """
-    from .soundfile_audio import get_sample_rate
-
-    return get_sample_rate(name)
+    if use_cache:
+        return _get_sample_rate(name)
+    return _get_sample_rate.__wrapped__(name)
 
 
 @lru_cache
-def read_audio(name: str | Path) -> AudioT:
+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 .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:
     """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
     """
-    from .soundfile_audio import read_audio
-
-    return read_audio(name)
+    if use_cache:
+        return _read_audio(name)
+    return _read_audio.__wrapped__(name)
 
 
 @lru_cache
-def read_ir(name: str | Path) -> ImpulseResponseData:
+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
+    """
+    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 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
     """
-    from .soundfile_audio import read_ir
+    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:
+    """Get the number of samples resampled to the SonusAI sample rate in the given file
 
-    return read_ir(name)
+    :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:
+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
     """
-    from .soundfile_audio import get_num_samples
+    import math
 
-    return get_num_samples(name)
+    import soundfile
+    from pydub import AudioSegment
+
+    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)