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

sonusai/mixture/generation.py

@@ -93,7 +93,8 @@ def initialize_db(location: str, test: bool = False) -> None:
     CREATE TABLE impulse_response_file (
     id INTEGER PRIMARY KEY NOT NULL,
     file TEXT NOT NULL,
-    tags TEXT NOT NULL)
+    tags TEXT NOT NULL,
+    delay INTEGER NOT NULL)
     """)
 
     con.execute("""
@@ -119,8 +120,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("""
@@ -361,11 +361,12 @@ def populate_impulse_response_file_table(
 
     con = db_connection(location=location, readonly=False, test=test)
     con.executemany(
-        "INSERT INTO impulse_response_file (file, tags) VALUES (?, ?)",
+        "INSERT INTO impulse_response_file (file, tags, delay) VALUES (?, ?, ?)",
         [
             (
                 impulse_response_file.file,
                 json.dumps(impulse_response_file.tags),
+                impulse_response_file.delay,
             )
             for impulse_response_file in impulse_response_files
         ],
@@ -389,8 +390,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 +403,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 +414,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 +428,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 +489,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 +510,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 +521,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,
@@ -543,17 +531,13 @@ def update_mixture_table(location: str, m_id: int, with_data: bool = False, test
 
 def _augmented_noise_audio(mixdb: MixtureDatabase, mixture: Mixture) -> AudioT:
     from .audio import read_audio
-    from .audio import read_ir
     from .augmentation import apply_augmentation
-    from .augmentation import apply_impulse_response
 
     noise = mixdb.noise_file(mixture.noise.file_id)
     noise_augmentation = mixture.noise.augmentation
 
     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_augmentation(mixdb, audio, noise_augmentation.pre)
 
     return audio
 
@@ -567,8 +551,9 @@ def _initialize_targets_audio(mixdb: MixtureDatabase, mixture: Mixture) -> tuple
         target_audio = mixdb.read_target_audio(target.file_id)
         targets_audio.append(
             apply_augmentation(
+                mixdb=mixdb,
                 audio=target_audio,
-                augmentation=target.augmentation,
+                augmentation=target.augmentation.pre,
                 frame_length=mixdb.feature_step_samples,
             )
         )
@@ -582,7 +567,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
 
@@ -691,7 +679,7 @@ def _exhaustive_noise_mix(
             noise_offset = 0
             noise_length = estimate_augmented_length_from_length(
                 length=noise_files[noise_file_id].samples,
-                tempo=noise_augmentation.tempo,
+                tempo=noise_augmentation.pre.tempo,
             )
 
             for augmented_target_ids_for_mixup in augmented_target_ids_for_mixups:
@@ -933,7 +921,7 @@ def _get_next_noise_indices(
 
     noise_augmentation = augmentation_from_rule(noise_augmentations[noise_augmentation_id], num_ir)
     noise_length = estimate_augmented_length_from_length(
-        length=noise_files[noise_file_id].samples, tempo=noise_augmentation.tempo
+        length=noise_files[noise_file_id].samples, tempo=noise_augmentation.pre.tempo
     )
     return noise_file_id, noise_augmentation_id, noise_augmentation, noise_length
 
@@ -957,7 +945,7 @@ def _get_next_noise_offset(
 
     noise_augmentation = augmentation_from_rule(noise_augmentations[noise_file_id], num_ir)
     noise_length = estimate_augmented_length_from_length(
-        length=noise_files[noise_file_id].samples, tempo=noise_augmentation.tempo
+        length=noise_files[noise_file_id].samples, tempo=noise_augmentation.pre.tempo
     )
     if noise_offset + target_length >= noise_length:
         if noise_offset == 0:
@@ -998,7 +986,7 @@ def _get_target_info(
         target_length = max(
             estimate_augmented_length_from_length(
                 length=target_files[tfi].samples,
-                tempo=target_augmentation.tempo,
+                tempo=target_augmentation.pre.tempo,
                 frame_length=feature_step_samples,
             ),
             target_length,

sonusai/mixture/helpers.py

@@ -117,11 +117,11 @@ def mixture_all_speech_metadata(mixdb: MixtureDatabase, mixture: Mixture) -> lis
                 # Check for tempo augmentation and adjust Interval start and end data as needed
                 entries = []
                 for entry in item:
-                    if target.augmentation.tempo is not None:
+                    if target.augmentation.pre.tempo is not None:
                         entries.append(
                             Interval(
-                                entry.start / target.augmentation.tempo,
-                                entry.end / target.augmentation.tempo,
+                                entry.start / target.augmentation.pre.tempo,
+                                entry.end / target.augmentation.pre.tempo,
                                 entry.label,
                             )
                         )
@@ -135,22 +135,26 @@ 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):
         target_file = mixdb.target_file(target.file_id)
-        target_augmentation = target.augmentation
         metadata += f"target {mi} name: {target_file.name}\n"
         metadata += f"target {mi} augmentation: {target.augmentation.to_dict()}\n"
-        metadata += f"target {mi} ir: {mixdb.impulse_response_file(target_augmentation.ir)}\n"
         metadata += f"target {mi} target_gain: {target.gain if not mixture.is_noise_only else 0}\n"
         metadata += f"target {mi} class indices: {target_file.class_indices}\n"
         for key in target_file.truth_configs:
@@ -162,7 +166,6 @@ def mixture_metadata(mixdb: MixtureDatabase, m_id: int) -> str:
     noise_augmentation = mixture.noise.augmentation
     metadata += f"noise name: {noise.name}\n"
     metadata += f"noise augmentation: {noise_augmentation.to_dict()}\n"
-    metadata += f"noise ir: {mixdb.impulse_response_file(noise_augmentation.ir)}\n"
     metadata += f"noise offset: {mixture.noise_offset}\n"
     metadata += f"snr: {mixture.snr}\n"
     metadata += f"random_snr: {mixture.snr.is_random}\n"
@@ -173,17 +176,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(
@@ -254,24 +265,25 @@ def get_target(mixdb: MixtureDatabase, mixture: Mixture, targets_audio: list[Aud
     :param targets_audio: List of augmented target audio data (one per target in the mixup)
     :return: Sum of augmented target audio data
     """
-    # Apply impulse responses to targets
+    # Apply post-truth augmentation effects to targets and sum
     import numpy as np
 
-    from .audio import read_ir
-    from .augmentation import apply_impulse_response
-
-    targets_ir = []
-    for idx, target in enumerate(targets_audio):
-        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))))
+    from .augmentation import apply_augmentation
+
+    targets_post = []
+    for idx, target_audio in enumerate(targets_audio):
+        target = mixture.targets[idx]
+        targets_post.append(
+            apply_augmentation(
+                mixdb=mixdb,
+                audio=target_audio,
+                augmentation=target.augmentation.post,
+                frame_length=mixdb.feature_step_samples,
             )
-        else:
-            targets_ir.append(target)
+        )
 
     # Return sum of targets
-    return np.sum(targets_ir, axis=0)
+    return np.sum(targets_post, axis=0)
 
 
 def get_transform_from_audio(audio: AudioT, transform: ForwardTransform) -> tuple[AudioF, EnergyT]:
@@ -385,7 +397,7 @@ def augmented_target_samples(
         [
             estimate_augmented_length_from_length(
                 length=target_files[fi].samples,
-                tempo=target_augmentations[ai].tempo,
+                tempo=target_augmentations[ai].pre.tempo,
                 frame_length=feature_step_samples,
             )
             for fi, ai in it
@@ -405,7 +417,7 @@ def augmented_noise_samples(noise_files: list[NoiseFile], noise_augmentations: l
 def augmented_noise_length(noise_file: NoiseFile, noise_augmentation: Augmentation) -> int:
     from .augmentation import estimate_augmented_length_from_length
 
-    return estimate_augmented_length_from_length(length=noise_file.samples, tempo=noise_augmentation.tempo)
+    return estimate_augmented_length_from_length(length=noise_file.samples, tempo=noise_augmentation.pre.tempo)
 
 
 def get_textgrid_tier_from_target_file(target_file: str, tier: str) -> SpeechMetadata | None:

sonusai/mixture/ir_delay.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+
+def get_impulse_response_delay(file: str) -> int:
+    from sonusai.utils import temp_seed
+
+    from .audio import raw_read_audio
+
+    ir, sample_rate = raw_read_audio(file)
+
+    with temp_seed(42):
+        wgn_ref = np.random.normal(loc=0, scale=0.2, size=int(np.ceil(0.05 * sample_rate))).astype(np.float32)
+
+    wgn_conv = np.convolve(ir, wgn_ref)
+
+    return int(np.round(tdoa(wgn_conv, wgn_ref, interp=16, phat=True)))
+
+
+def tdoa(signal: np.ndarray, reference: np.ndarray, interp: int = 1, phat: bool = False, fs: int | float = 1) -> float:
+    """Estimates the shift of array signal with respect to reference using generalized cross-correlation.
+
+    :param signal: The array whose tdoa is measured
+    :param reference: The reference array
+    :param interp: Interpolation factor for the output array
+    :param phat: Apply the PHAT weighting
+    :param fs: The sampling frequency of the input arrays
+    :return: The estimated delay between the two arrays
+    """
+    n_reference = reference.shape[0]
+
+    r_12 = correlate(signal, reference, interp=interp, phat=phat)
+
+    delay = (np.argmax(np.abs(r_12)) / interp - (n_reference - 1)) / fs
+
+    return float(delay)
+
+
+def correlate(x1: np.ndarray, x2: np.ndarray, interp: int = 1, phat: bool = False) -> np.ndarray:
+    """Compute the cross-correlation between x1 and x2
+
+    :param x1: Input array 1
+    :param x2: Input array 2
+    :param interp: Interpolation factor for the output array
+    :param phat: Apply the PHAT weighting
+    :return: The cross-correlation between the two arrays
+    """
+    n_x1 = x1.shape[0]
+    n_x2 = x2.shape[0]
+
+    n = n_x1 + n_x2 - 1
+
+    fft1 = np.fft.rfft(x1, n=n)
+    fft2 = np.fft.rfft(x2, n=n)
+
+    if phat:
+        eps1 = np.mean(np.abs(fft1)) * 1e-10
+        fft1 /= np.abs(fft1) + eps1
+        eps2 = np.mean(np.abs(fft2)) * 1e-10
+        fft2 /= np.abs(fft2) + eps2
+
+    out = np.fft.irfft(fft1 * np.conj(fft2), n=int(n * interp))
+
+    return np.concatenate([out[-interp * (n_x2 - 1) :], out[: (interp * n_x1)]])