sonusai 0.18.6__py3-none-any.whl → 0.18.7__py3-none-any.whl

sonusai/__init__.py

@@ -2,6 +2,12 @@ import logging
 from importlib import metadata
 from os.path import dirname
 
+from pyaaware import TorchForwardTransform
+from pyaaware import TorchInverseTransform
+
+ForwardTransform = TorchForwardTransform
+InverseTransform = TorchInverseTransform
+
 __version__ = metadata.version(__package__)
 BASEDIR = dirname(__file__)
 

sonusai/genmetrics.py

@@ -115,11 +115,11 @@ def main() -> None:
     mixdb = MixtureDatabase(location)
     supported = mixdb.supported_metrics
     if show_supported:
-        logger.info(f'\nSupported metrics: {", ".join(sorted(supported))}')
+        logger.info(f'\nSupported metrics:\n\n{supported.pretty}')
         sys.exit(0)
 
     if includes is None or 'all' in includes:
-        metrics = supported
+        metrics = supported.names
     else:
         metrics = set(includes)
         if 'mxwer' in metrics:
@@ -127,7 +127,7 @@ def main() -> None:
             for name in mixdb.asr_configs:
                 metrics.add(f'mxwer.{name}')
 
-    diff = metrics.difference(supported)
+    diff = metrics.difference(supported.names)
     if diff:
         logger.error(f'Unrecognized metric: {", ".join(diff)}')
         sys.exit(1)
@@ -141,7 +141,7 @@ def main() -> None:
             for name in mixdb.asr_configs:
                 _excludes.add(f'mxwer.{name}')
 
-    diff = _excludes.difference(supported)
+    diff = _excludes.difference(supported.names)
     if diff:
         logger.error(f'Unrecognized metric: {", ".join(diff)}')
         sys.exit(1)

sonusai/metrics/__init__.py

@@ -8,7 +8,8 @@ from .calc_pesq import calc_pesq
 from .calc_phase_distance import calc_phase_distance
 from .calc_sa_sdr import calc_sa_sdr
 from .calc_sample_weights import calc_sample_weights
-from .calc_snr_f import calc_snr_f
+from .calc_segsnr_f import calc_segsnr_f
+from .calc_segsnr_f import calc_segsnr_f_bin
 from .calc_speech import calc_speech
 from .calc_wer import calc_wer
 from .calc_wsdr import calc_wsdr

sonusai/metrics/calc_audio_stats.py

@@ -2,6 +2,14 @@ from sonusai.mixture.datatypes import AudioStatsMetrics
 from sonusai.mixture.datatypes import AudioT
 
 
+def _convert_str_with_factors_to_int(x: str) -> int:
+    if 'k' in x:
+        return int(1000 * float(x.replace('k', '')))
+    if 'M' in x:
+        return int(1000000 * float(x.replace('M', '')))
+    return int(x)
+
+
 def calc_audio_stats(audio: AudioT, win_len: float = None) -> AudioStatsMetrics:
     from sonusai.mixture import SAMPLE_RATE
     from sonusai.mixture import Transformer
@@ -38,5 +46,5 @@ def calc_audio_stats(audio: AudioT, win_len: float = None) -> AudioStatsMetrics:
         tr=float(stats['RMS Tr dB']),
         cr=float(stats['Crest factor']),
         fl=float(stats['Flat factor']),
-        pkc=int(stats['Pk count']),
+        pkc=_convert_str_with_factors_to_int(stats['Pk count']),
     )

sonusai/metrics/calc_segsnr_f.py

@@ -0,0 +1,84 @@
+import numpy as np
+
+from sonusai.mixture.datatypes import AudioF
+from sonusai.mixture.datatypes import Segsnr
+from sonusai.mixture.datatypes import SnrFBinMetrics
+from sonusai.mixture.datatypes import SnrFMetrics
+
+
+def calc_segsnr_f(segsnr_f: Segsnr) -> SnrFMetrics:
+    """Calculate metrics of snr_f truth data.
+
+    Includes mean and standard deviation of the linear values (usually energy)
+    and mean and standard deviation of the dB values (10 * log10).
+    """
+    if np.count_nonzero(segsnr_f) == 0:
+        # If all entries are zeros
+        return SnrFMetrics(0, 0, -np.inf, 0)
+
+    tmp = np.ma.array(segsnr_f, mask=np.logical_not(np.isfinite(segsnr_f)))
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are infinite
+        return SnrFMetrics(np.inf, 0, np.inf, 0)
+
+    snr_mean = np.mean(tmp, axis=0)
+    snr_std = np.std(tmp, axis=0)
+
+    tmp = 10 * np.ma.log10(tmp)
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are masked, special case where all inputs are either 0 or infinite
+        snr_db_mean = -np.inf
+        snr_db_std = np.inf
+    else:
+        snr_db_mean = np.mean(tmp, axis=0)
+        snr_db_std = np.std(tmp, axis=0)
+
+    return SnrFMetrics(snr_mean,
+                       snr_std,
+                       snr_db_mean,
+                       snr_db_std)
+
+
+def calc_segsnr_f_bin(target_f: AudioF, noise_f: AudioF) -> SnrFBinMetrics:
+    """Calculate per-bin segmental SNR metrics.
+
+    Includes per-bin mean and standard deviation of the linear values
+    and mean and standard deviation of the dB values.
+    """
+    if target_f.ndim != 2 and noise_f.ndim != 2:
+        raise ValueError('target_f and noise_f must have 2 dimensions')
+
+    segsnr_f = (np.abs(target_f) ** 2) / (np.abs(noise_f) ** 2)
+
+    frames, bins = segsnr_f.shape
+    if np.count_nonzero(segsnr_f) == 0:
+        # If all entries are zeros
+        return SnrFBinMetrics(np.zeros(bins),
+                              np.zeros(bins),
+                              -np.inf * np.ones(bins),
+                              np.zeros(bins))
+
+    tmp = np.ma.array(segsnr_f, mask=np.logical_not(np.isfinite(segsnr_f)))
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are infinite
+        return SnrFBinMetrics(np.inf * np.ones(bins),
+                              np.zeros(bins),
+                              np.inf * np.ones(bins),
+                              np.zeros(bins))
+
+    snr_mean = np.mean(tmp, axis=0)
+    snr_std = np.std(tmp, axis=0)
+
+    tmp = 10 * np.ma.log10(tmp)
+    if np.ma.count_masked(tmp) == np.ma.size(tmp, axis=0):
+        # If all entries are masked, special case where all inputs are either 0 or infinite
+        snr_db_mean = -np.inf * np.ones(bins)
+        snr_db_std = np.inf * np.ones(bins)
+    else:
+        snr_db_mean = np.mean(tmp, axis=0)
+        snr_db_std = np.std(tmp, axis=0)
+
+    return SnrFBinMetrics(np.ma.getdata(snr_mean),
+                          np.ma.getdata(snr_std),
+                          np.ma.getdata(snr_db_mean),
+                          np.ma.getdata(snr_db_std))

sonusai/metrics/calc_speech.py

@@ -6,7 +6,7 @@ from .calc_pesq import calc_pesq
 
 
 def calc_speech(hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int = SAMPLE_RATE) -> SpeechMetrics:
-    """Calculate speech metrics pesq, c_sig, c_bak, c_ovl, seg_snr.
+    """Calculate speech metrics pesq, c_sig, c_bak, and c_ovl.
     
     These are all related and thus included in one function. Reference: matlab script "compute_metrics.m".
 
@@ -38,11 +38,11 @@ def calc_speech(hypothesis: np.ndarray, reference: np.ndarray, sample_rate: int
     _pesq = calc_pesq(hypothesis=hypothesis, reference=reference, sample_rate=sample_rate)
 
     # Now compute the composite measures
-    c_sig = np.clip(3.093 - 1.029 * llr_mean + 0.603 * _pesq - 0.009 * wss_dist, 1, 5)
-    c_bak = np.clip(1.634 + 0.478 * _pesq - 0.007 * wss_dist + 0.063 * seg_snr, 1, 5)
-    c_ovl = np.clip(1.594 + 0.805 * _pesq - 0.512 * llr_mean - 0.007 * wss_dist, 1, 5)
+    csig = np.clip(3.093 - 1.029 * llr_mean + 0.603 * _pesq - 0.009 * wss_dist, 1, 5)
+    cbak = np.clip(1.634 + 0.478 * _pesq - 0.007 * wss_dist + 0.063 * seg_snr, 1, 5)
+    covl = np.clip(1.594 + 0.805 * _pesq - 0.512 * llr_mean - 0.007 * wss_dist, 1, 5)
 
-    return SpeechMetrics(_pesq, c_sig, c_bak, c_ovl)
+    return SpeechMetrics(_pesq, csig, cbak, covl)
 
 
 def _calc_weighted_spectral_slope_measure(hypothesis: np.ndarray,

sonusai/mixture/__init__.py

@@ -66,6 +66,8 @@ from .datatypes import GeneralizedIDs
 from .datatypes import ImpulseResponseData
 from .datatypes import ImpulseResponseFiles
 from .datatypes import ListAudiosT
+from .datatypes import MetricDoc
+from .datatypes import MetricDocs
 from .datatypes import Mixture
 from .datatypes import MixtureDatabaseConfig
 from .datatypes import Mixtures
@@ -105,6 +107,7 @@ from .helpers import augmented_noise_samples
 from .helpers import augmented_target_samples
 from .helpers import check_audio_files_exist
 from .helpers import forward_transform
+from .helpers import frames_from_samples
 from .helpers import get_audio_from_transform
 from .helpers import get_ft
 from .helpers import get_segsnr

sonusai/mixture/datatypes.py

@@ -1,7 +1,9 @@
 from dataclasses import dataclass
 from typing import Any
+from typing import Iterable
 from typing import NamedTuple
 from typing import Optional
+from typing import SupportsIndex
 from typing import TypeAlias
 
 import numpy as np
@@ -336,17 +338,24 @@ SpeechMetadata: TypeAlias = str | list[Interval] | None
 
 
 class SnrFMetrics(NamedTuple):
-    mean: Optional[float] = None
-    var: Optional[float] = None
-    db_mean: Optional[float] = None
+    avg: Optional[float] = None
+    std: Optional[float] = None
+    db_avg: Optional[float] = None
     db_std: Optional[float] = None
 
 
+class SnrFBinMetrics(NamedTuple):
+    avg: Optional[np.ndarray] = None
+    std: Optional[np.ndarray] = None
+    db_avg: Optional[np.ndarray] = None
+    db_std: Optional[np.ndarray] = None
+
+
 class SpeechMetrics(NamedTuple):
     pesq: Optional[float] = None
-    c_sig: Optional[float] = None
-    c_bak: Optional[float] = None
-    c_ovl: Optional[float] = None
+    csig: Optional[float] = None
+    cbak: Optional[float] = None
+    covl: Optional[float] = None
 
 
 class AudioStatsMetrics(NamedTuple):
@@ -360,3 +369,53 @@ class AudioStatsMetrics(NamedTuple):
     cr: Optional[float] = None
     fl: Optional[float] = None
     pkc: Optional[float] = None
+
+
+@dataclass
+class MetricDoc:
+    category: str
+    name: str
+    description: str
+
+
+class MetricDocs(list[MetricDoc]):
+    def __init__(self, __iterable: Iterable[MetricDoc]) -> None:
+        super().__init__(item for item in __iterable)
+
+    def __setitem__(self, __key: SupportsIndex, __value: MetricDoc) -> None:  # type: ignore
+        super().__setitem__(__key, __value)
+
+    def insert(self, __index: SupportsIndex, __object: MetricDoc) -> None:
+        super().insert(__index, __object)
+
+    def append(self, __object: MetricDoc) -> None:
+        super().append(__object)
+
+    def extend(self, __iterable: Iterable[MetricDoc]) -> None:
+        if isinstance(__iterable, type(self)):
+            super().extend(__iterable)
+        else:
+            super().extend(item for item in __iterable)
+
+    @property
+    def pretty(self) -> str:
+        max_category_len = ((max([len(item.category) for item in self]) + 9) // 10) * 10
+        max_name_len = 2 + ((max([len(item.name) for item in self]) + 1) // 2) * 2
+        categories: list[str] = []
+        for item in self:
+            if item.category not in categories:
+                categories.append(item.category)
+
+        result = ''
+        for category in categories:
+            result += f'{category}\n'
+            result += '-' * max_category_len + '\n'
+            for item in [sub for sub in self if sub.category == category]:
+                result += f'  {item.name:<{max_name_len}}{item.description}\n'
+            result += '\n'
+
+        return result
+
+    @property
+    def names(self) -> set[str]:
+        return set(item.name for item in self)

sonusai/mixture/feature.py

@@ -1,6 +1,5 @@
 from typing import Optional
 
-from sonusai.mixture.datatypes import AudioF
 from sonusai.mixture.datatypes import AudioT
 from sonusai.mixture.datatypes import Feature
 
@@ -58,15 +57,13 @@ def get_feature_from_audio(audio: AudioT,
 def get_audio_from_feature(feature: Feature,
                            feature_mode: str,
                            num_classes: Optional[int] = 1,
-                           truth_mutex: Optional[bool] = False,
-                           trim: Optional[bool] = True) -> AudioT:
+                           truth_mutex: Optional[bool] = False) -> AudioT:
     """Apply inverse transform to feature data to generate audio data
 
     :param feature: Feature data [frames, strides, feature_parameters]
     :param feature_mode: Feature mode
     :param num_classes: Number of classes
     :param truth_mutex: Whether to calculate 'other' label
-    :param trim: Whether to trim the audio data
     :return: Audio data [samples]
     """
     import numpy as np
@@ -76,6 +73,7 @@ def get_audio_from_feature(feature: Feature,
     from .datatypes import TransformConfig
     from .helpers import inverse_transform
     from sonusai.utils.stacked_complex import unstack_complex
+    from sonusai.utils.compress import power_uncompress
 
     fg = FeatureGenerator(feature_mode=feature_mode,
                           num_classes=num_classes,
@@ -83,23 +81,10 @@ def get_audio_from_feature(feature: Feature,
 
     feature_complex = unstack_complex(feature)
     if feature_mode[0:1] == 'h':
-        feature_complex = _power_uncompress(feature_complex)
+        feature_complex = power_uncompress(feature_complex)
     return np.squeeze(inverse_transform(transform=feature_complex,
                                         config=TransformConfig(N=fg.itransform_N,
                                                                R=fg.itransform_R,
                                                                bin_start=fg.bin_start,
                                                                bin_end=fg.bin_end,
-                                                               ttype=fg.itransform_ttype),
-                                        trim=trim))
-
-
-def _power_uncompress(feature: AudioF) -> AudioF:
-    import numpy as np
-
-    mag = np.abs(feature)
-    phase = np.angle(feature)
-    mag = mag ** (1. / 0.3)
-    real_uncompress = mag * np.cos(phase)
-    imag_uncompress = mag * np.sin(phase)
-
-    return real_uncompress + 1j * imag_uncompress
+                                                               ttype=fg.itransform_ttype)))

sonusai/mixture/helpers.py

@@ -2,9 +2,9 @@ from typing import Any
 from typing import Optional
 
 from praatio.utilities.constants import Interval
-from pyaaware import ForwardTransform
-from pyaaware import InverseTransform
 
+from sonusai import ForwardTransform
+from sonusai import InverseTransform
 from sonusai.mixture import EnergyT
 from sonusai.mixture.datatypes import AudioF
 from sonusai.mixture.datatypes import AudioT
@@ -285,7 +285,10 @@ def read_mixture_data(name: str, items: list[str] | str) -> Any:
 
     def _get_dataset(file: h5py.File, d_name: str) -> Any:
         if d_name in file:
-            return np.array(file[d_name])
+            data = np.array(file[d_name])
+            if data.size == 1:
+                return data.item()
+            return data
         return None
 
     if not isinstance(items, list):
@@ -371,8 +374,8 @@ def get_ft(mixdb: MixtureDatabase, mixture: Mixture, mixture_audio: AudioT, trut
 
     mixture_f = get_mixture_f(mixdb=mixdb, mixture=mixture, mixture_audio=mixture_audio)
 
-    transform_frames = mixdb.mixture_transform_frames(mixture.samples)
-    feature_frames = mixdb.mixture_feature_frames(mixture.samples)
+    transform_frames = frames_from_samples(mixture.samples, mixdb.ft_config.R)
+    feature_frames = frames_from_samples(mixture.samples, mixdb.feature_step_samples)
 
     feature = np.empty((feature_frames, mixdb.fg_stride, mixdb.feature_parameters), dtype=np.float32)
     truth_f = np.empty((feature_frames, mixdb.num_classes), dtype=np.complex64)
@@ -418,20 +421,21 @@ def get_segsnr_t(mixdb: MixtureDatabase, mixture: Mixture, target_audio: AudioT,
     :return: segsnr_t data
     """
     import numpy as np
-    from pyaaware import AawareForwardTransform
+    import torch
+    from sonusai import ForwardTransform
 
     from sonusai import SonusAIError
 
-    fft = AawareForwardTransform(N=mixdb.ft_config.N,
-                                 R=mixdb.ft_config.R,
-                                 bin_start=mixdb.ft_config.bin_start,
-                                 bin_end=mixdb.ft_config.bin_end,
-                                 ttype=mixdb.ft_config.ttype)
+    fft = ForwardTransform(N=mixdb.ft_config.N,
+                           R=mixdb.ft_config.R,
+                           bin_start=mixdb.ft_config.bin_start,
+                           bin_end=mixdb.ft_config.bin_end,
+                           ttype=mixdb.ft_config.ttype)
 
     segsnr_t = np.empty(mixture.samples, dtype=np.float32)
 
-    _, target_energy = fft.execute_all(target_audio)
-    _, noise_energy = fft.execute_all(noise_audio)
+    target_energy = fft.execute_all(torch.from_numpy(target_audio))[1].numpy()
+    noise_energy = fft.execute_all(torch.from_numpy(noise_audio))[1].numpy()
 
     offsets = range(0, mixture.samples, mixdb.ft_config.R)
     if len(target_energy) != len(offsets):
@@ -505,8 +509,11 @@ def get_transform_from_audio(audio: AudioT, transform: ForwardTransform) -> tupl
     :param transform: ForwardTransform object
     :return: Frequency domain data [frames, bins], Energy [frames]
     """
-    f, e = transform.execute_all(audio)
-    return f.transpose(), e
+    import torch
+
+    f, e = transform.execute_all(torch.from_numpy(audio))
+
+    return f.numpy(), e.numpy()
 
 
 def forward_transform(audio: AudioT, config: TransformConfig) -> AudioF:
@@ -518,54 +525,50 @@ def forward_transform(audio: AudioT, config: TransformConfig) -> AudioF:
     :param config: Transform configuration
     :return: Frequency domain data [frames, bins]
     """
-    from pyaaware import AawareForwardTransform
+    from sonusai import ForwardTransform
 
     audio_f, _ = get_transform_from_audio(audio=audio,
-                                          transform=AawareForwardTransform(N=config.N,
-                                                                           R=config.R,
-                                                                           bin_start=config.bin_start,
-                                                                           bin_end=config.bin_end,
-                                                                           ttype=config.ttype))
+                                          transform=ForwardTransform(N=config.N,
+                                                                     R=config.R,
+                                                                     bin_start=config.bin_start,
+                                                                     bin_end=config.bin_end,
+                                                                     ttype=config.ttype))
     return audio_f
 
 
-def get_audio_from_transform(data: AudioF, transform: InverseTransform, trim: bool = True) -> tuple[AudioT, EnergyT]:
+def get_audio_from_transform(data: AudioF, transform: InverseTransform) -> tuple[AudioT, EnergyT]:
     """Apply inverse transform to input transform data to generate audio data
 
     :param data: Frequency domain data [frames, bins]
     :param transform: InverseTransform object
-    :param trim: Removes starting samples so output waveform will be time-aligned with input waveform to the transform
     :return: Time domain data [samples], Energy [frames]
     """
-    t, e = transform.execute_all(data.transpose())
-    if trim:
-        t = t[transform.N - transform.R:]
+    import torch
 
-    return t, e
+    t, e = transform.execute_all(torch.from_numpy(data))
 
+    return t.numpy(), e.numpy()
 
-def inverse_transform(transform: AudioF, config: TransformConfig, trim: bool = True) -> AudioT:
+
+def inverse_transform(transform: AudioF, config: TransformConfig) -> AudioT:
     """Transform frequency domain data into time domain using the inverse transform config from the feature
 
     A new transform is used for each call; i.e., state is not maintained between calls to inverse_transform().
 
     :param transform: Frequency domain data [frames, bins]
     :param config: Transform configuration
-    :param trim: Removes starting samples so output waveform will be time-aligned with input waveform to the
-                 transform
     :return: Time domain data [samples]
     """
     import numpy as np
-    from pyaaware import AawareInverseTransform
+    from sonusai import InverseTransform
 
     audio, _ = get_audio_from_transform(data=transform,
-                                        transform=AawareInverseTransform(N=config.N,
-                                                                         R=config.R,
-                                                                         bin_start=config.bin_start,
-                                                                         bin_end=config.bin_end,
-                                                                         ttype=config.ttype,
-                                                                         gain=np.float32(1)),
-                                        trim=trim)
+                                        transform=InverseTransform(N=config.N,
+                                                                   R=config.R,
+                                                                   bin_start=config.bin_start,
+                                                                   bin_end=config.bin_end,
+                                                                   ttype=config.ttype,
+                                                                   gain=np.float32(1)))
     return audio
 
 
@@ -641,3 +644,9 @@ def get_textgrid_tier_from_target_file(target_file: str, tier: str) -> Optional[
         return list(entries)
     else:
         return entries[0].label
+
+
+def frames_from_samples(samples: int, step_samples: int) -> int:
+    import numpy as np
+
+    return int(np.ceil(samples / step_samples))