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

sonusai/mixture/audio.py

@@ -44,44 +44,77 @@ 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
     """
+    if use_cache:
+        return _get_sample_rate(name)
+    return _get_sample_rate.__wrapped__(name)
+
+
+@lru_cache
+def _get_sample_rate(name: str | Path) -> int:
     from .soundfile_audio import get_sample_rate
 
     return get_sample_rate(name)
 
 
-@lru_cache
-def read_audio(name: str | Path) -> AudioT:
+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:
     from .soundfile_audio import read_audio
 
     return read_audio(name)
 
 
-@lru_cache
-def read_ir(name: str | Path) -> ImpulseResponseData:
+def read_ir(name: str | Path, use_cache: bool = True) -> ImpulseResponseData:
     """Read impulse response data
 
     :param name: File name
+    :param use_cache: If true, use LRU caching
     :return: ImpulseResponseData object
     """
+    if use_cache:
+        return _read_ir(name)
+    return _read_ir.__wrapped__(name)
+
+
+@lru_cache
+def _read_ir(name: str | Path) -> ImpulseResponseData:
     from .soundfile_audio import read_ir
 
     return read_ir(name)
 
 
+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:
+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

sonusai/mixture/generation.py

@@ -119,8 +119,7 @@ def initialize_db(location: str, test: bool = False) -> None:
     id INTEGER PRIMARY KEY NOT NULL,
     file_id INTEGER NOT NULL,
     augmentation TEXT NOT NULL,
-    FOREIGN KEY(file_id) REFERENCES target_file (id),
-    UNIQUE(file_id, augmentation))
+    FOREIGN KEY(file_id) REFERENCES target_file (id))
     """)
 
     con.execute("""
@@ -389,8 +388,7 @@ def update_mixid_width(location: str, num_mixtures: int, test: bool = False) ->
     con.close()
 
 
-def populate_mixture_table(
-    location: str,
+def generate_mixtures(
     noise_mix_mode: str,
     augmented_targets: list[AugmentedTarget],
     target_files: list[TargetFile],
@@ -403,13 +401,8 @@ def populate_mixture_table(
     num_classes: int,
     feature_step_samples: int,
     num_ir: int,
-    test: bool = False,
-) -> tuple[int, int]:
-    """Generate mixtures and populate mixture table"""
-    from .helpers import from_mixture
-    from .helpers import from_target
-    from .mixdb import db_connection
-
+) -> tuple[int, int, list[Mixture]]:
+    """Generate mixtures"""
     if noise_mix_mode == "exhaustive":
         func = _exhaustive_noise_mix
     elif noise_mix_mode == "non-exhaustive":
@@ -419,7 +412,7 @@ def populate_mixture_table(
     else:
         raise ValueError(f"invalid noise_mix_mode: {noise_mix_mode}")
 
-    used_noise_files, used_noise_samples, mixtures = func(
+    return func(
         augmented_targets=augmented_targets,
         target_files=target_files,
         target_augmentations=target_augmentations,
@@ -433,20 +426,41 @@ def populate_mixture_table(
         num_ir=num_ir,
     )
 
+
+def populate_mixture_table(
+    location: str,
+    mixtures: list[Mixture],
+    test: bool = False,
+    logging: bool = False,
+    show_progress: bool = False,
+) -> None:
+    """Populate mixture table"""
+    from sonusai import logger
+    from sonusai.utils import track
+
+    from .helpers import from_mixture
+    from .helpers import from_target
+    from .mixdb import db_connection
+
     con = db_connection(location=location, readonly=False, test=test)
+
     # Populate target table
+    if logging:
+        logger.info("Populating target table")
+    targets: list[tuple[int, str]] = []
     for mixture in mixtures:
         for target in mixture.targets:
-            con.execute(
-                """
-                INSERT OR IGNORE INTO target (file_id, augmentation)
-                VALUES (?, ?)
-                """,
-                from_target(target),
-            )
+            entry = from_target(target)
+            if entry not in targets:
+                targets.append(entry)
+    for target in track(targets, disable=not show_progress):
+        con.execute("INSERT INTO target (file_id, augmentation) VALUES (?, ?)", target)
 
     # Populate mixture table
-    for m_id, mixture in enumerate(mixtures):
+    if logging:
+        logger.info("Populating mixture table")
+    for mixture in track(mixtures, disable=not show_progress):
+        m_id = int(mixture.name)
         con.execute(
             """
             INSERT INTO mixture (id, name, noise_file_id, noise_augmentation, noise_offset, noise_snr_gain, random_snr,
@@ -473,20 +487,13 @@ def populate_mixture_table(
     con.commit()
     con.close()
 
-    return used_noise_files, used_noise_samples
 
-
-def update_mixture_table(location: str, m_id: int, with_data: bool = False, test: bool = False) -> GenMixData:
+def update_mixture(mixdb: MixtureDatabase, mixture: Mixture, with_data: bool = False) -> tuple[Mixture, GenMixData]:
     """Update mixture record with name and gains"""
     from .audio import get_next_noise
     from .augmentation import apply_gain
-    from .datatypes import GenMixData
-    from .helpers import from_mixture
     from .helpers import get_target
-    from .mixdb import db_connection
 
-    mixdb = MixtureDatabase(location, test)
-    mixture = mixdb.mixture(m_id)
     mixture, targets_audio = _initialize_targets_audio(mixdb, mixture)
 
     noise_audio = _augmented_noise_audio(mixdb, mixture)
@@ -501,29 +508,8 @@ def update_mixture_table(location: str, m_id: int, with_data: bool = False, test
 
     mixture.name = f"{int(mixture.name):0{mixdb.mixid_width}}"
 
-    con = db_connection(location=location, readonly=False, test=test)
-    con.execute(
-        """
-        UPDATE mixture SET  name=?,
-                            noise_file_id=?,
-                            noise_augmentation=?,
-                            noise_offset=?,
-                            noise_snr_gain=?,
-                            random_snr=?,
-                            snr=?,
-                            samples=?,
-                            spectral_mask_id=?,
-                            spectral_mask_seed=?,
-                            target_snr_gain=?
-        WHERE ? = mixture.id
-        """,
-        (*from_mixture(mixture), m_id + 1),
-    )
-    con.commit()
-    con.close()
-
     if not with_data:
-        return GenMixData()
+        return mixture, GenMixData()
 
     # Apply SNR gains
     targets_audio = [apply_gain(audio=target_audio, gain=mixture.target_snr_gain) for target_audio in targets_audio]
@@ -533,7 +519,7 @@ def update_mixture_table(location: str, m_id: int, with_data: bool = False, test
     target_audio = get_target(mixdb, mixture, targets_audio)
     mixture_audio = target_audio + noise_audio
 
-    return GenMixData(
+    return mixture, GenMixData(
         mixture=mixture_audio,
         targets=targets_audio,
         target=target_audio,
@@ -553,7 +539,7 @@ def _augmented_noise_audio(mixdb: MixtureDatabase, mixture: Mixture) -> AudioT:
     audio = read_audio(noise.name)
     audio = apply_augmentation(audio, noise_augmentation)
     if noise_augmentation.ir is not None:
-        audio = apply_impulse_response(audio, read_ir(mixdb.impulse_response_file(noise_augmentation.ir)))
+        audio = apply_impulse_response(audio, read_ir(mixdb.impulse_response_file(noise_augmentation.ir)))  # pyright: ignore [reportArgumentType]
 
     return audio
 
@@ -582,7 +568,10 @@ def _initialize_targets_audio(mixdb: MixtureDatabase, mixture: Mixture) -> tuple
 
 
 def _initialize_mixture_gains(
-    mixdb: MixtureDatabase, mixture: Mixture, target_audio: AudioT, noise_audio: AudioT
+    mixdb: MixtureDatabase,
+    mixture: Mixture,
+    target_audio: AudioT,
+    noise_audio: AudioT,
 ) -> Mixture:
     import numpy as np
 

sonusai/mixture/helpers.py

@@ -135,14 +135,20 @@ def mixture_all_speech_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> lis
     return results
 
 
-def mixture_metadata(mixdb: MixtureDatabase, m_id: int) -> str:
+def mixture_metadata(mixdb: MixtureDatabase, m_id: int | None = None, mixture: Mixture | None = None) -> str:
     """Create a string of metadata for a Mixture
 
     :param mixdb: Mixture database
     :param m_id: Mixture ID
+    :param mixture: Mixture record
     :return: String of metadata
     """
-    mixture = mixdb.mixture(m_id)
+    if m_id is not None:
+        mixture = mixdb.mixture(m_id)
+
+    if mixture is None:
+        raise ValueError("No mixture specified.")
+
     metadata = ""
     speech_metadata = mixture_all_speech_metadata(mixdb, mixture)
     for mi, target in enumerate(mixture.targets):
@@ -173,17 +179,25 @@ def mixture_metadata(mixdb: MixtureDatabase, m_id: int) -> str:
     return metadata
 
 
-def write_mixture_metadata(mixdb: MixtureDatabase, m_id: int) -> None:
+def write_mixture_metadata(mixdb: MixtureDatabase, m_id: int | None = None, mixture: Mixture | None = None) -> None:
     """Write mixture metadata to a text file
 
     :param mixdb: Mixture database
     :param m_id: Mixture ID
+    :param mixture: Mixture record
     """
     from os.path import join
 
-    name = join(mixdb.location, "mixture", mixdb.mixture(m_id).name, "metadata.txt")
+    if m_id is not None:
+        name = mixdb.mixture(m_id).name
+    elif mixture is not None:
+        name = mixture.name
+    else:
+        raise ValueError("No mixture specified.")
+
+    name = join(mixdb.location, "mixture", name, "metadata.txt")
     with open(file=name, mode="w") as f:
-        f.write(mixture_metadata(mixdb, m_id))
+        f.write(mixture_metadata(mixdb, m_id, mixture))
 
 
 def from_mixture(
@@ -246,12 +260,13 @@ def to_target(entry: TargetRecord) -> Target:
     )
 
 
-def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: list[AudioT]) -> AudioT:
+def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: list[AudioT], use_cache: bool = True) -> AudioT:
     """Get the augmented target audio data for the given mixture record
 
     :param mixdb: Mixture database
     :param mixture: Mixture record
     :param targets_audio: List of augmented target audio data (one per target in the mixup)
+    :param use_cache: If true, use LRU caching
     :return: Sum of augmented target audio data
     """
     # Apply impulse responses to targets
@@ -265,7 +280,7 @@ def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: list[Aud
         ir_idx = mixture.targets[idx].augmentation.ir
         if ir_idx is not None:
             targets_ir.append(
-                apply_impulse_response(audio=target, ir=read_ir(mixdb.impulse_response_file(int(ir_idx))))
+                apply_impulse_response(audio=target, ir=read_ir(mixdb.impulse_response_file(int(ir_idx)), use_cache))  # pyright: ignore [reportArgumentType]
             )
         else:
             targets_ir.append(target)

sonusai/mixture/mixdb.py

@@ -61,7 +61,7 @@ def db_connection(
     if not create and readonly:
         name += "?mode=ro"
 
-    connection = sqlite3.connect("file:" + name, uri=True)
+    connection = sqlite3.connect("file:" + name, uri=True, timeout=20)
 
     if verbose:
         connection.set_trace_callback(print)
@@ -84,7 +84,7 @@ class SQLiteContextManager:
 
 
 class MixtureDatabase:
-    def __init__(self, location: str, test: bool = False) -> None:
+    def __init__(self, location: str, test: bool = False, use_cache: bool = True) -> None:
         import json
         from os.path import exists
 
@@ -92,6 +92,7 @@ class MixtureDatabase:
 
         self.location = location
         self.test = test
+        self.use_cache = use_cache
 
         if not exists(db_file(self.location, self.test)):
             raise OSError(f"Could not find mixture database in {self.location}")
@@ -121,7 +122,7 @@ class MixtureDatabase:
             class_weights_threshold=self.class_weights_thresholds,
             feature=self.feature,
             impulse_response_files=self.impulse_response_files,
-            mixtures=self.mixtures,
+            mixtures=self.mixtures(),
             noise_mix_mode=self.noise_mix_mode,
             noise_files=self.noise_files,
             num_classes=self.num_classes,
@@ -488,7 +489,7 @@ class MixtureDatabase:
             return truth_configs
 
     def target_truth_configs(self, t_id: int) -> TruthConfigs:
-        return _target_truth_configs(self.db, t_id)
+        return _target_truth_configs(self.db, t_id, self.use_cache)
 
     @cached_property
     def random_snrs(self) -> list[float]:
@@ -556,7 +557,7 @@ class MixtureDatabase:
         :param sm_id: Spectral mask ID
         :return: Spectral mask
         """
-        return _spectral_mask(self.db, sm_id)
+        return _spectral_mask(self.db, sm_id, self.use_cache)
 
     @cached_property
     def target_files(self) -> list[TargetFile]:
@@ -619,7 +620,7 @@ class MixtureDatabase:
         :param t_id: Target file ID
         :return: Target file
         """
-        return _target_file(self.db, t_id)
+        return _target_file(self.db, t_id, self.use_cache)
 
     @cached_property
     def num_target_files(self) -> int:
@@ -657,7 +658,7 @@ class MixtureDatabase:
         :param n_id: Noise file ID
         :return: Noise file
         """
-        return _noise_file(self.db, n_id)
+        return _noise_file(self.db, n_id, self.use_cache)
 
     @cached_property
     def num_noise_files(self) -> int:
@@ -706,7 +707,7 @@ class MixtureDatabase:
         """
         if ir_id is None:
             return None
-        return _impulse_response_file(self.db, ir_id)
+        return _impulse_response_file(self.db, ir_id, self.use_cache)
 
     @cached_property
     def num_impulse_response_files(self) -> int:
@@ -717,7 +718,6 @@ class MixtureDatabase:
         with self.db() as c:
             return int(c.execute("SELECT count(impulse_response_file.id) FROM impulse_response_file").fetchone()[0])
 
-    @cached_property
     def mixtures(self) -> list[Mixture]:
         """Get mixtures from db
 
@@ -760,7 +760,7 @@ class MixtureDatabase:
         :param m_id: Zero-based mixture ID
         :return: Mixture record
         """
-        return _mixture(self.db, m_id)
+        return _mixture(self.db, m_id, self.use_cache)
 
     @cached_property
     def mixid_width(self) -> int:
@@ -805,7 +805,7 @@ class MixtureDatabase:
         """
         from .audio import read_audio
 
-        return read_audio(self.target_file(t_id).name)
+        return read_audio(self.target_file(t_id).name, self.use_cache)
 
     def augmented_noise_audio(self, mixture: Mixture) -> AudioT:
         """Get augmented noise audio
@@ -819,12 +819,12 @@ class MixtureDatabase:
         from .augmentation import apply_impulse_response
 
         noise = self.noise_file(mixture.noise.file_id)
-        audio = read_audio(noise.name)
+        audio = read_audio(noise.name, self.use_cache)
         audio = apply_augmentation(audio, mixture.noise.augmentation)
         if mixture.noise.augmentation.ir is not None:
             audio = apply_impulse_response(
                 audio,
-                read_ir(self.impulse_response_file(mixture.noise.augmentation.ir)),
+                read_ir(self.impulse_response_file(mixture.noise.augmentation.ir), self.use_cache),  # pyright: ignore [reportArgumentType]
             )
 
         return audio
@@ -1332,7 +1332,7 @@ class MixtureDatabase:
         return sorted(set(self.speaker_metadata_tiers + self.textgrid_metadata_tiers))
 
     def speaker(self, s_id: int | None, tier: str) -> str | None:
-        return _speaker(self.db, s_id, tier)
+        return _speaker(self.db, s_id, tier, self.use_cache)
 
     def speech_metadata(self, tier: str) -> list[str]:
         from .helpers import get_textgrid_tier_from_target_file
@@ -1464,7 +1464,7 @@ class MixtureDatabase:
 
         return sorted(result)
 
-    def mixture_metrics(self, m_id: int, metrics: list[str], force: bool = False) -> list[Any]:
+    def mixture_metrics(self, m_id: int, metrics: list[str], force: bool = False) -> dict[str, Any]:
         """Get metrics data for the given mixture ID
 
         :param m_id: Zero-based mixture ID
@@ -1916,21 +1916,34 @@ class MixtureDatabase:
 
             raise AttributeError(f"Unrecognized metric: '{m}'")
 
-        result: list[float | int | str | Segsnr | None] = []
+        result: dict[str, Any] = {}
         for metric in metrics:
-            result.append(calc(metric))
+            result[metric] = calc(metric)
+
+            # Check for metrics dependencies and add them even if not explicitly requested.
+            if metric.startswith("mxwer"):
+                dependencies = ("mxasr." + metric[6:], "tasr." + metric[6:])
+                for dependency in dependencies:
+                    result[dependency] = calc(dependency)
 
         return result
 
 
-@lru_cache
-def _spectral_mask(db: partial, sm_id: int) -> SpectralMask:
+def _spectral_mask(db: partial, sm_id: int, use_cache: bool = True) -> SpectralMask:
     """Get spectral mask with ID from db
 
     :param db: Database context
     :param sm_id: Spectral mask ID
+    :param use_cache: If true, use LRU caching
     :return: Spectral mask
     """
+    if use_cache:
+        return __spectral_mask(db, sm_id)
+    return __spectral_mask.__wrapped__(db, sm_id)
+
+
+@lru_cache
+def __spectral_mask(db: partial, sm_id: int) -> SpectralMask:
     from .db_datatypes import SpectralMaskRecord
 
     with db() as c:
@@ -1953,12 +1966,26 @@ def _spectral_mask(db: partial, sm_id: int) -> SpectralMask:
         )
 
 
+def _target_file(db: partial, t_id: int, use_cache: bool = True) -> TargetFile:
+    """Get target file with ID from db
+
+    :param db: Database context
+    :param t_id: Target file ID
+    :param use_cache: If true, use LRU caching
+    :return: Target file
+    """
+    if use_cache:
+        return __target_file(db, t_id, use_cache)
+    return __target_file.__wrapped__(db, t_id, use_cache)
+
+
 @lru_cache
-def _target_file(db: partial, t_id: int) -> TargetFile:
+def __target_file(db: partial, t_id: int, use_cache: bool = True) -> TargetFile:
     """Get target file with ID from db
 
     :param db: Database context
     :param t_id: Target file ID
+    :param use_cache: If true, use LRU caching
     :return: Target file
     """
     import json
@@ -1982,19 +2009,26 @@ def _target_file(db: partial, t_id: int) -> TargetFile:
             samples=target_file.samples,
             class_indices=json.loads(target_file.class_indices),
             level_type=target_file.level_type,
-            truth_configs=_target_truth_configs(db, t_id),
+            truth_configs=_target_truth_configs(db, t_id, use_cache),
             speaker_id=target_file.speaker_id,
         )
 
 
-@lru_cache
-def _noise_file(db: partial, n_id: int) -> NoiseFile:
+def _noise_file(db: partial, n_id: int, use_cache: bool = True) -> NoiseFile:
     """Get noise file with ID from db
 
     :param db: Database context
     :param n_id: Noise file ID
+    :param use_cache: If true, use LRU caching
     :return: Noise file
     """
+    if use_cache:
+        return __noise_file(db, n_id)
+    return __noise_file.__wrapped__(db, n_id)
+
+
+@lru_cache
+def __noise_file(db: partial, n_id: int) -> NoiseFile:
     with db() as c:
         noise = c.execute(
             """
@@ -2007,14 +2041,21 @@ def _noise_file(db: partial, n_id: int) -> NoiseFile:
         return NoiseFile(name=noise[0], samples=noise[1])
 
 
-@lru_cache
-def _impulse_response_file(db: partial, ir_id: int) -> str:
+def _impulse_response_file(db: partial, ir_id: int, use_cache: bool = True) -> str:
     """Get impulse response file with ID from db
 
     :param db: Database context
     :param ir_id: Impulse response file ID
-    :return: Noise
+    :param use_cache: If true, use LRU caching
+    :return: Impulse response
     """
+    if use_cache:
+        return __impulse_response_file(db, ir_id)
+    return __impulse_response_file.__wrapped__(db, ir_id)
+
+
+@lru_cache
+def __impulse_response_file(db: partial, ir_id: int) -> str:
     with db() as c:
         return str(
             c.execute(
@@ -2028,14 +2069,21 @@ def _impulse_response_file(db: partial, ir_id: int) -> str:
         )
 
 
-@lru_cache
-def _mixture(db: partial, m_id: int) -> Mixture:
+def _mixture(db: partial, m_id: int, use_cache: bool = True) -> Mixture:
     """Get mixture record with ID from db
 
     :param db: Database context
     :param m_id: Zero-based mixture ID
+    :param use_cache: If true, use LRU caching
     :return: Mixture record
     """
+    if use_cache:
+        return __mixture(db, m_id)
+    return __mixture.__wrapped__(db, m_id)
+
+
+@lru_cache
+def __mixture(db: partial, m_id: int) -> Mixture:
     from .db_datatypes import MixtureRecord
     from .db_datatypes import TargetRecord
     from .helpers import to_mixture
@@ -2068,8 +2116,14 @@ def _mixture(db: partial, m_id: int) -> Mixture:
         return to_mixture(mixture, targets)
 
 
+def _speaker(db: partial, s_id: int | None, tier: str, use_cache: bool = True) -> str | None:
+    if use_cache:
+        return __speaker(db, s_id, tier)
+    return __speaker.__wrapped__(db, s_id, tier)
+
+
 @lru_cache
-def _speaker(db: partial, s_id: int | None, tier: str) -> str | None:
+def __speaker(db: partial, s_id: int | None, tier: str) -> str | None:
     if s_id is None:
         return None
 
@@ -2082,8 +2136,14 @@ def _speaker(db: partial, s_id: int | None, tier: str) -> str | None:
         return data[0]
 
 
+def _target_truth_configs(db: partial, t_id: int, use_cache: bool = True) -> TruthConfigs:
+    if use_cache:
+        return __target_truth_configs(db, t_id)
+    return __target_truth_configs.__wrapped__(db, t_id)
+
+
 @lru_cache
-def _target_truth_configs(db: partial, t_id: int) -> TruthConfigs:
+def __target_truth_configs(db: partial, t_id: int) -> TruthConfigs:
     import json
 
     from .datatypes import TruthConfig