sonusai 0.12.5__py3-none-any.whl → 0.12.7__py3-none-any.whl

sonusai/mixture/generate_mixtures.py

@@ -11,8 +11,7 @@ def generate_mixtures(mixdb: MixtureDatabase,
                       augmented_targets: AugmentedTargets,
                       noise_files: NoiseFiles,
                       noise_augmentations: Augmentations,
-                      mixups: List[int],
-                      show_progress: bool = False) -> Tuple[int, int]:
+                      mixups: List[int]) -> Tuple[int, int]:
     """Generate mixtures and append to mixture database
 
     :param mixdb: Mixture database
@@ -20,7 +19,6 @@ def generate_mixtures(mixdb: MixtureDatabase,
     :param noise_files: List of noise files
     :param noise_augmentations: List of noise augmentations
     :param mixups: List of mixup values
-    :param show_progress: Show progress bar
     :return: (Number of noise files used, number of noise samples used)
     """
     from sonusai import SonusAIError
@@ -30,20 +28,17 @@ def generate_mixtures(mixdb: MixtureDatabase,
                                      augmented_targets=augmented_targets,
                                      noise_files=noise_files,
                                      noise_augmentations=noise_augmentations,
-                                     mixups=mixups,
-                                     show_progress=show_progress)
+                                     mixups=mixups)
 
     if mixdb.noise_mix_mode == 'non-exhaustive':
         return _non_exhaustive_noise_mix(mixdb=mixdb,
                                          augmented_targets=augmented_targets,
-                                         mixups=mixups,
-                                         show_progress=show_progress)
+                                         mixups=mixups)
 
     if mixdb.noise_mix_mode == 'non-combinatorial':
         return _non_combinatorial_noise_mix(mixdb=mixdb,
                                             augmented_targets=augmented_targets,
-                                            mixups=mixups,
-                                            show_progress=show_progress)
+                                            mixups=mixups)
 
     raise SonusAIError(f'invalid noise_mix_mode: {mixdb.noise_mix_mode}')
 
@@ -52,8 +47,7 @@ def _exhaustive_noise_mix(mixdb: MixtureDatabase,
                           augmented_targets: AugmentedTargets,
                           noise_files: NoiseFiles,
                           noise_augmentations: Augmentations,
-                          mixups: List[int],
-                          show_progress: bool = False) -> Tuple[int, int]:
+                          mixups: List[int]) -> Tuple[int, int]:
     """Exhaustive noise mix mode
     Use every noise/augmentation with every target/augmentation.
 
@@ -62,13 +56,11 @@ def _exhaustive_noise_mix(mixdb: MixtureDatabase,
     :param noise_files: List of noise files
     :param noise_augmentations: List of noise augmentations
     :param mixups: List of mixup values
-    :param show_progress: Show progress bar
     :return: (Number of noise files used, number of noise samples used)
     """
     from random import randint
 
     import numpy as np
-    from tqdm import tqdm
 
     from sonusai.mixture import MRecord
     from sonusai.mixture import get_augmented_target_indices_for_mixup
@@ -79,52 +71,43 @@ def _exhaustive_noise_mix(mixdb: MixtureDatabase,
     augmented_target_indices_for_mixups = [get_augmented_target_indices_for_mixup(mixdb=mixdb,
                                                                                   augmented_targets=augmented_targets,
                                                                                   mixup=mixup) for mixup in mixups]
-    total = (len(noise_files) *
-             len(noise_augmentations) *
-             len([target for mixup in augmented_target_indices_for_mixups for target in mixup]) *
-             len(mixdb.spectral_masks) *
-             len(mixdb.all_snrs))
-
-    with tqdm(total=total, disable=not show_progress) as progress:
-        for noise_file_index in range(len(noise_files)):
-            for noise_augmentation_index in range(len(noise_augmentations)):
-                noise_offset = 0
-                noise_length = mixdb.augmented_noise_length(noise_file_index, noise_augmentation_index)
-
-                for augmented_target_indices_for_mixup in augmented_target_indices_for_mixups:
-                    for augmented_target_indices in augmented_target_indices_for_mixup:
-                        (target_file_index,
-                         target_augmentation_index,
-                         target_length) = _get_target_info(mixdb=mixdb,
-                                                           augmented_target_indices=augmented_target_indices,
-                                                           augmented_targets=augmented_targets)
-
-                        for spectral_mask_index in range(len(mixdb.spectral_masks)):
-                            for snr in mixdb.all_snrs:
-                                mixdb.mixtures.append(MRecord(
-                                    name=None,
-                                    target_file_index=target_file_index,
-                                    target_augmentation_index=target_augmentation_index,
-                                    noise_file_index=noise_file_index,
-                                    noise_offset=noise_offset,
-                                    noise_augmentation_index=noise_augmentation_index,
-                                    samples=target_length,
-                                    snr=snr.value,
-                                    spectral_mask_index=spectral_mask_index,
-                                    spectral_mask_seed=randint(0, np.iinfo('i').max),
-                                    random_snr=snr.is_random))
-
-                                noise_offset = int((noise_offset + target_length) % noise_length)
-                                used_noise_samples += target_length
-                                progress.update()
+    for noise_file_index in range(len(noise_files)):
+        for noise_augmentation_index in range(len(noise_augmentations)):
+            noise_offset = 0
+            noise_length = mixdb.augmented_noise_length(noise_file_index, noise_augmentation_index)
+
+            for augmented_target_indices_for_mixup in augmented_target_indices_for_mixups:
+                for augmented_target_indices in augmented_target_indices_for_mixup:
+                    (target_file_index,
+                     target_augmentation_index,
+                     target_length) = _get_target_info(mixdb=mixdb,
+                                                       augmented_target_indices=augmented_target_indices,
+                                                       augmented_targets=augmented_targets)
+
+                    for spectral_mask_index in range(len(mixdb.spectral_masks)):
+                        for snr in mixdb.all_snrs:
+                            mixdb.mixtures.append(MRecord(
+                                name=None,
+                                target_file_index=target_file_index,
+                                target_augmentation_index=target_augmentation_index,
+                                noise_file_index=noise_file_index,
+                                noise_offset=noise_offset,
+                                noise_augmentation_index=noise_augmentation_index,
+                                samples=target_length,
+                                snr=snr.value,
+                                spectral_mask_index=spectral_mask_index,
+                                spectral_mask_seed=randint(0, np.iinfo('i').max),
+                                random_snr=snr.is_random))
+
+                            noise_offset = int((noise_offset + target_length) % noise_length)
+                            used_noise_samples += target_length
 
     return used_noise_files, used_noise_samples
 
 
 def _non_exhaustive_noise_mix(mixdb: MixtureDatabase,
                               augmented_targets: AugmentedTargets,
-                              mixups: List[int],
-                              show_progress: bool = False) -> Tuple[int, int]:
+                              mixups: List[int]) -> Tuple[int, int]:
     """Non-exhaustive noise mix mode
     Cycle through every target/augmentation without necessarily using all
     noise/augmentation combinations (reduced data set).
@@ -132,13 +115,11 @@ def _non_exhaustive_noise_mix(mixdb: MixtureDatabase,
     :param mixdb: Mixture database
     :param augmented_targets: List of augmented targets
     :param mixups: List of mixup values
-    :param show_progress: Show progress bar
     :return: (Number of noise files used, number of noise samples used)
     """
     from random import randint
 
     import numpy as np
-    from tqdm import tqdm
 
     from sonusai.mixture import MRecord
     from sonusai.mixture import get_augmented_target_indices_for_mixup
@@ -152,54 +133,47 @@ def _non_exhaustive_noise_mix(mixdb: MixtureDatabase,
     augmented_target_indices_for_mixups = [get_augmented_target_indices_for_mixup(mixdb=mixdb,
                                                                                   augmented_targets=augmented_targets,
                                                                                   mixup=mixup) for mixup in mixups]
-    total = (len([target for mixup in augmented_target_indices_for_mixups for target in mixup]) *
-             len(mixdb.spectral_masks) *
-             len(mixdb.all_snrs))
-
-    with tqdm(total=total, disable=not show_progress) as progress:
-        for mixup in augmented_target_indices_for_mixups:
-            for augmented_target_indices in mixup:
-                (target_file_index,
-                 target_augmentation_index,
-                 target_length) = _get_target_info(mixdb=mixdb,
-                                                   augmented_target_indices=augmented_target_indices,
-                                                   augmented_targets=augmented_targets)
-
-                for spectral_mask_index in range(len(mixdb.spectral_masks)):
-                    for snr in mixdb.all_snrs:
-                        used_noise_files.add(f'{noise_file_index}_{noise_augmentation_index}')
-                        (noise_file_index,
-                         noise_augmentation_index,
-                         noise_offset) = _get_next_noise_offset(mixdb=mixdb,
-                                                                target_length=target_length,
-                                                                noise_file_index=noise_file_index,
-                                                                noise_augmentation_index=noise_augmentation_index,
-                                                                noise_offset=noise_offset)
-
-                        mixdb.mixtures.append(MRecord(
-                            name=None,
-                            target_file_index=target_file_index,
-                            target_augmentation_index=target_augmentation_index,
-                            noise_file_index=noise_file_index,
-                            noise_augmentation_index=noise_augmentation_index,
-                            noise_offset=noise_offset,
-                            samples=target_length,
-                            snr=snr.value,
-                            spectral_mask_index=spectral_mask_index,
-                            spectral_mask_seed=randint(0, np.iinfo('i').max),
-                            random_snr=snr.is_random))
-
-                        noise_offset += target_length
-                        used_noise_samples += target_length
-                        progress.update()
+    for mixup in augmented_target_indices_for_mixups:
+        for augmented_target_indices in mixup:
+            (target_file_index,
+             target_augmentation_index,
+             target_length) = _get_target_info(mixdb=mixdb,
+                                               augmented_target_indices=augmented_target_indices,
+                                               augmented_targets=augmented_targets)
+
+            for spectral_mask_index in range(len(mixdb.spectral_masks)):
+                for snr in mixdb.all_snrs:
+                    used_noise_files.add(f'{noise_file_index}_{noise_augmentation_index}')
+                    (noise_file_index,
+                     noise_augmentation_index,
+                     noise_offset) = _get_next_noise_offset(mixdb=mixdb,
+                                                            target_length=target_length,
+                                                            noise_file_index=noise_file_index,
+                                                            noise_augmentation_index=noise_augmentation_index,
+                                                            noise_offset=noise_offset)
+
+                    mixdb.mixtures.append(MRecord(
+                        name=None,
+                        target_file_index=target_file_index,
+                        target_augmentation_index=target_augmentation_index,
+                        noise_file_index=noise_file_index,
+                        noise_augmentation_index=noise_augmentation_index,
+                        noise_offset=noise_offset,
+                        samples=target_length,
+                        snr=snr.value,
+                        spectral_mask_index=spectral_mask_index,
+                        spectral_mask_seed=randint(0, np.iinfo('i').max),
+                        random_snr=snr.is_random))
+
+                    noise_offset += target_length
+                    used_noise_samples += target_length
 
     return len(used_noise_files), used_noise_samples
 
 
 def _non_combinatorial_noise_mix(mixdb: MixtureDatabase,
                                  augmented_targets: AugmentedTargets,
-                                 mixups: List[int],
-                                 show_progress: bool = False) -> Tuple[int, int]:
+                                 mixups: List[int]) -> Tuple[int, int]:
     """Non-combinatorial noise mix mode
     Combine a target/augmentation with a single cut of a noise/augmentation
     non-exhaustively (each target/augmentation does not use each noise/augmentation).
@@ -208,14 +182,12 @@ def _non_combinatorial_noise_mix(mixdb: MixtureDatabase,
     :param mixdb: Mixture database
     :param augmented_targets: List of augmented targets
     :param mixups: List of mixup values
-    :param show_progress: Show progress bar
     :return: (Number of noise files used, number of noise samples used)
     """
     from random import choice
     from random import randint
 
     import numpy as np
-    from tqdm import tqdm
 
     from sonusai.mixture import MRecord
     from sonusai.mixture import get_augmented_target_indices_for_mixup
@@ -228,43 +200,37 @@ def _non_combinatorial_noise_mix(mixdb: MixtureDatabase,
     augmented_target_indices_for_mixups = [get_augmented_target_indices_for_mixup(mixdb=mixdb,
                                                                                   augmented_targets=augmented_targets,
                                                                                   mixup=mixup) for mixup in mixups]
-    total = (len([target for mixup in augmented_target_indices_for_mixups for target in mixup]) *
-             len(mixdb.spectral_masks) *
-             len(mixdb.all_snrs))
-
-    with tqdm(total=total, disable=not show_progress) as progress:
-        for mixup in augmented_target_indices_for_mixups:
-            for augmented_target_indices in mixup:
-                (target_file_index,
-                 target_augmentation_index,
-                 target_length) = _get_target_info(mixdb=mixdb,
-                                                   augmented_target_indices=augmented_target_indices,
-                                                   augmented_targets=augmented_targets)
-
-                for spectral_mask_index in range(len(mixdb.spectral_masks)):
-                    for snr in mixdb.all_snrs:
-                        used_noise_files.add(f'{noise_file_index}_{noise_augmentation_index}')
-                        (noise_file_index,
-                         noise_augmentation_index,
-                         noise_length) = _get_next_noise_indices(mixdb=mixdb,
-                                                                 noise_file_index=noise_file_index,
-                                                                 noise_augmentation_index=noise_augmentation_index)
-
-                        mixdb.mixtures.append(MRecord(
-                            name=None,
-                            target_file_index=target_file_index,
-                            target_augmentation_index=target_augmentation_index,
-                            noise_file_index=noise_file_index,
-                            noise_augmentation_index=noise_augmentation_index,
-                            noise_offset=choice(range(noise_length)),
-                            samples=target_length,
-                            snr=snr.value,
-                            spectral_mask_index=spectral_mask_index,
-                            spectral_mask_seed=randint(0, np.iinfo('i').max),
-                            random_snr=snr.is_random))
-
-                        used_noise_samples += target_length
-                        progress.update()
+    for mixup in augmented_target_indices_for_mixups:
+        for augmented_target_indices in mixup:
+            (target_file_index,
+             target_augmentation_index,
+             target_length) = _get_target_info(mixdb=mixdb,
+                                               augmented_target_indices=augmented_target_indices,
+                                               augmented_targets=augmented_targets)
+
+            for spectral_mask_index in range(len(mixdb.spectral_masks)):
+                for snr in mixdb.all_snrs:
+                    used_noise_files.add(f'{noise_file_index}_{noise_augmentation_index}')
+                    (noise_file_index,
+                     noise_augmentation_index,
+                     noise_length) = _get_next_noise_indices(mixdb=mixdb,
+                                                             noise_file_index=noise_file_index,
+                                                             noise_augmentation_index=noise_augmentation_index)
+
+                    mixdb.mixtures.append(MRecord(
+                        name=None,
+                        target_file_index=target_file_index,
+                        target_augmentation_index=target_augmentation_index,
+                        noise_file_index=noise_file_index,
+                        noise_augmentation_index=noise_augmentation_index,
+                        noise_offset=choice(range(noise_length)),
+                        samples=target_length,
+                        snr=snr.value,
+                        spectral_mask_index=spectral_mask_index,
+                        spectral_mask_seed=randint(0, np.iinfo('i').max),
+                        random_snr=snr.is_random))
+
+                    used_noise_samples += target_length
 
     return len(used_noise_files), used_noise_samples
 
@@ -309,7 +275,7 @@ def _get_next_noise_offset(mixdb: MixtureDatabase,
 def _get_target_info(mixdb: MixtureDatabase,
                      augmented_target_indices: List[int],
                      augmented_targets: AugmentedTargets) -> Tuple[List[int], List[int], int]:
-    from sonusai.mixture import estimate_augmented_length_from_audio
+    from sonusai.mixture import estimate_augmented_length_from_length
 
     target_file_index = []
     target_augmentation_index = []
@@ -321,8 +287,8 @@ def _get_target_info(mixdb: MixtureDatabase,
         target_file_index.append(tfi)
         target_augmentation_index.append(tai)
 
-        target_length = max(estimate_augmented_length_from_audio(audio=mixdb.raw_target_audio(tfi),
-                                                                 augmentation=mixdb.target_augmentations[tai],
-                                                                 length_common_denominator=mixdb.feature_step_samples),
+        target_length = max(estimate_augmented_length_from_length(length=mixdb.targets[tfi].samples,
+                                                                  augmentation=mixdb.target_augmentations[tai],
+                                                                  length_common_denominator=mixdb.feature_step_samples),
                             target_length)
     return target_file_index, target_augmentation_index, target_length

sonusai/mixture/mixdb.py

@@ -431,17 +431,20 @@ class MixtureDatabase:
     def truth_settings(self, value: TruthSettings) -> None:
         self._config.truth_settings = value
 
-    def augmented_target_length(self, file_index: int, augmentation_index: int) -> int:
+    def augmented_target_length(self, file_index: int,
+                                augmentation_index: int) -> int:
         from sonusai.mixture import estimate_augmented_length_from_length
 
-        return estimate_augmented_length_from_length(self.targets[file_index].samples,
-                                                     self.target_augmentations[augmentation_index])
+        return estimate_augmented_length_from_length(length=self.targets[file_index].samples,
+                                                     augmentation=self.target_augmentations[augmentation_index],
+                                                     length_common_denominator=self.feature_step_samples)
 
-    def augmented_noise_length(self, file_index: int, augmentation_index: int) -> int:
+    def augmented_noise_length(self, file_index: int,
+                               augmentation_index: int) -> int:
         from sonusai.mixture import estimate_augmented_length_from_length
 
-        return estimate_augmented_length_from_length(self.noises[file_index].samples,
-                                                     self.noise_augmentations[augmentation_index])
+        return estimate_augmented_length_from_length(length=self.noises[file_index].samples,
+                                                     augmentation=self.noise_augmentations[augmentation_index])
 
     def raw_target_audio(self, file_index: int) -> AudioT:
         if self._use_cache:
@@ -472,7 +475,58 @@ class MixtureDatabase:
     def raw_target_audio_cache_info(self) -> NamedTuple:
         return self._cache_raw_target_audio.cache_info()
 
+    def augmented_target_audio(self, file_index: int, augmentation_index: int) -> AudioT:
+        """Get augmented target audio
+
+        :param file_index: Target audio file index
+        :param augmentation_index: Target augmentation index
+        :return: Augmented target audio
+        """
+        if self._use_cache:
+            return self._cache_augmented_target_audio(file_index, augmentation_index)
+        else:
+            return self._no_cache_augmented_target_audio(file_index, augmentation_index)
+
+    @lru_cache
+    def _cache_augmented_target_audio(self, file_index: int, augmentation_index: int) -> AudioT:
+        """Get augmented target audio (using cache)
+
+        :param file_index: Target audio file name
+        :param augmentation_index: Target augmentation index
+        :return: Augmented target audio
+        """
+        return self._no_cache_augmented_target_audio(file_index, augmentation_index)
+
+    def _no_cache_augmented_target_audio(self, file_index: int, augmentation_index: int) -> AudioT:
+        """Get augmented target audio (without using cache)
+
+        :param file_index: Target audio file index
+        :param augmentation_index: Target augmentation index
+        :return: Augmented target audio
+        """
+        from sonusai.mixture import apply_augmentation
+        from sonusai.mixture import apply_ir
+        from sonusai.mixture import read_audio
+
+        audio = read_audio(self.targets[file_index].name)
+        augmentation = self.target_augmentations[augmentation_index]
+        audio = apply_augmentation(audio, augmentation, self.feature_step_samples)
+        if augmentation.ir is not None:
+            audio = apply_ir(audio, self.ir_data[int(augmentation.ir)])
+
+        return audio
+
+    @property
+    def augmented_target_audio_cache_info(self) -> NamedTuple:
+        return self._cache_augmented_target_audio.cache_info()
+
     def augmented_noise_audio(self, file_index: int, augmentation_index: int) -> AudioT:
+        """Get augmented noise audio
+
+        :param file_index: Noise audio file index
+        :param augmentation_index: Noise augmentation index
+        :return: Augmented noise audio
+        """
         if self._use_cache:
             return self._cache_augmented_noise_audio(file_index, augmentation_index)
         else:
@@ -634,6 +688,20 @@ class MixtureDatabase:
 
                     line_num += 1
 
+    @property
+    def augmented_target_samples(self) -> int:
+        from itertools import product
+
+        it = list(product(*[range(len(self.targets)), range(len(self.target_augmentations))]))
+        return sum([self.augmented_target_length(fi, ai) for fi, ai in it])
+
+    @property
+    def augmented_noise_samples(self) -> int:
+        from itertools import product
+
+        it = list(product(*[range(len(self.noises)), range(len(self.noise_augmentations))]))
+        return sum([self.augmented_noise_length(fi, ai) for fi, ai in it])
+
     def total_samples(self, mixids: GeneralizedIDs = '*') -> int:
         return sum([self.mixture_samples(mixid) for mixid in self.mixids_to_list(mixids)])
 
@@ -817,7 +885,8 @@ class MixtureDatabase:
         mrecord = self.mixtures[mixid]
         targets = []
         for idx in range(len(mrecord.target_file_index)):
-            target = apply_augmentation(audio=self.raw_target_audio(mrecord.target_file_index[idx]),
+            target = self.raw_target_audio(mrecord.target_file_index[idx])
+            target = apply_augmentation(audio=target,
                                         augmentation=self.target_augmentations[mrecord.target_augmentation_index[idx]],
                                         length_common_denominator=self.feature_step_samples)
             target = apply_gain(audio=target, gain=mrecord.target_snr_gain)
@@ -996,7 +1065,7 @@ class MixtureDatabase:
         :param mixid: Mixture ID
         :param targets: List of augmented target audio data (one per target in the mixup) for the given mixid
         :param noise: Augmented noise audio data for the given mixid
-        :param force: Force computing data from original sources regardless of whether or not cached data exists
+        :param force: Force computing data from original sources regardless of whether cached data exists
         :return: truth_t data
         """
         import numpy as np