Repository: iver56/audiomentations (python)

Issue:
Improve TimeMask defaults
The current default masks too much on average

Pull request description:
Improve TimeMask
* Add `fade_duration` argument and remove `fade`
* Add `mask_location` argument
* Improve performance by ~50% compared to v0.40.0
* Adjust default value of `max_band_part` and `max_band_part` - Close #249 

Patch:
diff --git a/audiomentations/augmentations/time_mask.py b/audiomentations/augmentations/time_mask.py
index 71e5481..7a784f4 100644
--- a/audiomentations/augmentations/time_mask.py
+++ b/audiomentations/augmentations/time_mask.py
@@ -1,25 +1,33 @@
 import random
-
 import numpy as np
+from typing import Literal
 from numpy.typing import NDArray
 
 from audiomentations.core.transforms_interface import BaseWaveformTransform
 from audiomentations.core.utils import get_crossfade_mask_pair
 
+# 0.00025 s corresponds to 2 samples at 8 kHz
+DURATION_EPSILON = 0.00025
+
 
 class TimeMask(BaseWaveformTransform):
     """
-    Make a randomly chosen part of the audio silent.
-    Inspired by https://arxiv.org/pdf/1904.08779.pdf
+    Make a chosen part of the audio silent (time-masking augmentation).
+
+    The silent part can optionally be faded in/out to avoid hard transients.
+
+    Inspired by *SpecAugment: A Simple Data Augmentation Method for Automatic
+    Speech Recognition* (https://arxiv.org/pdf/1904.08779.pdf)
     """
 
     supports_multichannel = True
 
     def __init__(
         self,
-        min_band_part: float = 0.0,
-        max_band_part: float = 0.5,
-        fade: bool = True,
+        min_band_part: float = 0.01,
+        max_band_part: float = 0.2,
+        fade_duration: float = 0.005,
+        mask_location: Literal["start", "end", "random"] = "random",
         p: float = 0.5,
     ):
         """
@@ -27,44 +35,89 @@ class TimeMask(BaseWaveformTransform):
             total sound length. Must be between 0.0 and 1.0
         :param max_band_part: Maximum length of the silent part as a fraction of the
             total sound length. Must be between 0.0 and 1.0
-        :param fade: When set to True, a smooth fade-in and fade-out is added to the silent part.
-            This can smooth out unwanted abrupt changes between consecutive samples, which might
-            otherwise sound like transients/clicks/pops.
+        :param fade_duration: Duration (in seconds) of the fade-in and fade-out applied
+            at the edges of the silent region to smooth transitions and avoid abrupt
+            changes, which can otherwise produce impulses or clicks in the audio.
+            If you need hard edges or clicks, set this to `0.0` to disable fading.
+            Positive values must be at least 0.00025 seconds.
+        :param mask_location: Where to place the silent region.
+            * "start": silence begins at index 0
+            * "end": silence ends at the last sample
+            * "random" (default): silence starts at a random position
         :param p: The probability of applying this transform
         """
         super().__init__(p)
-        if min_band_part < 0.0 or min_band_part > 1.0:
+
+        if not (0.0 <= min_band_part <= 1.0):
             raise ValueError("min_band_part must be between 0.0 and 1.0")
-        if max_band_part < 0.0 or max_band_part > 1.0:
+        if not (0.0 <= max_band_part <= 1.0):
             raise ValueError("max_band_part must be between 0.0 and 1.0")
         if min_band_part > max_band_part:
             raise ValueError("min_band_part must not be greater than max_band_part")
+
+        if fade_duration < 0:
+            raise ValueError("fade_duration must be non-negative")
+        if 0 < fade_duration < DURATION_EPSILON:
+            raise ValueError(
+                f"When fade_duration is positive it must be >= {DURATION_EPSILON}"
+            )
+
+        if mask_location not in {"start", "end", "random"}:
+            raise ValueError('mask_location must be "start", "end" or "random"')
+
         self.min_band_part = min_band_part
         self.max_band_part = max_band_part
-        self.fade = fade
+        self.fade_duration = fade_duration
+        self.mask_location = mask_location
 
     def randomize_parameters(self, samples: NDArray[np.float32], sample_rate: int):
         super().randomize_parameters(samples, sample_rate)
-        if self.parameters["should_apply"]:
-            num_samples = samples.shape[-1]
-            self.parameters["t"] = random.randint(
-                int(num_samples * self.min_band_part),
-                int(num_samples * self.max_band_part),
-            )
-            self.parameters["t0"] = random.randint(
-                0, num_samples - self.parameters["t"]
-            )
+        if not self.parameters["should_apply"]:
+            return
+
+        num_samples = samples.shape[-1]
+
+        # Mask length
+        t = random.randint(
+            int(num_samples * self.min_band_part),
+            int(num_samples * self.max_band_part),
+        )
+
+        # Start index based on mask_location
+        loc = self.mask_location
+        if loc == "start":
+            t0 = 0
+        elif loc == "end":
+            t0 = num_samples - t
+        else:  # "random"
+            t0 = random.randint(0, num_samples - t)
+
+        self.parameters.update({"t": t, "t0": t0})
 
-    def apply(self, samples: NDArray[np.float32], sample_rate: int) -> NDArray[np.float32]:
+    def apply(
+        self, samples: NDArray[np.float32], sample_rate: int
+    ) -> NDArray[np.float32]:
         new_samples = samples.copy()
-        t = self.parameters["t"]
-        t0 = self.parameters["t0"]
-        mask = np.zeros(t, dtype=np.float32)
-        if self.fade:
-            fade_length = min(int(sample_rate * 0.01), int(t * 0.1))
-            if fade_length >= 2:
-                fade_in, fade_out = get_crossfade_mask_pair(fade_length, equal_energy=False)
-                mask[:fade_length] = fade_out
-                mask[-fade_length:] = fade_in
-        new_samples[..., t0 : t0 + t] *= mask
+        t: int = self.parameters["t"]
+        t0: int = self.parameters["t0"]
+
+        fade_len = 0
+        if self.fade_duration > 0.0:
+            fade_len = int(round(sample_rate * self.fade_duration))
+            fade_len = min(fade_len, t // 2)
+
+        if fade_len >= 2:
+            fade_in, fade_out = get_crossfade_mask_pair(fade_len, equal_energy=False)
+
+            left = slice(t0, t0 + fade_len)
+            mid = slice(t0 + fade_len, t0 + t - fade_len)
+            right = slice(t0 + t - fade_len, t0 + t)
+
+            new_samples[..., left] *= fade_out
+            if mid.start < mid.stop:
+                new_samples[..., mid] = 0.0
+            new_samples[..., right] *= fade_in
+        else:
+            new_samples[..., t0 : t0 + t] = 0.0
+
         return new_samples
diff --git a/demo/demo.py b/demo/demo.py
index 132c6c9..c21c8db 100644
--- a/demo/demo.py
+++ b/demo/demo.py
@@ -487,8 +487,10 @@ if __name__ == "__main__":
             "name": "ShiftSeconds",
         },
         {"instance": TanhDistortion(p=1.0), "num_runs": 5},
-        {"instance": TimeMask(fade=True, p=1.0), "num_runs": 5, "name": "TimeMaskWithFade"},
-        {"instance": TimeMask(fade=False, p=1.0), "num_runs": 5, "name": "TimeMaskWithoutFade"},
+        {"instance": TimeMask(fade_duration=0.05, p=1.0), "num_runs": 5, "name": "TimeMaskWithFade"},
+        {"instance": TimeMask(fade_duration=0, p=1.0), "num_runs": 5, "name": "TimeMaskWithoutFade"},
+        {"instance": TimeMask(mask_location="start", p=1.0), "num_runs": 5, "name": "TimeMaskStart"},
+        {"instance": TimeMask(mask_location="end", p=1.0), "num_runs": 5, "name": "TimeMaskEnd"},
         {
             "instance": TimeStretch(
                 min_rate=0.8, max_rate=1.25, method="signalsmith_stretch", p=1.0
diff --git a/docs/waveform_transforms/time_mask.md b/docs/waveform_transforms/time_mask.md
index 073de03..de64e33 100644
--- a/docs/waveform_transforms/time_mask.md
+++ b/docs/waveform_transforms/time_mask.md
@@ -35,21 +35,39 @@ augmented_sound = transform(my_waveform_ndarray, sample_rate=16000)
 ## TimeMask API
 
 [`min_band_part`](#min_band_part){ #min_band_part }: `float` • range: [0.0, 1.0]
-:   :octicons-milestone-24: Default: `0.0`. Minimum length of the silent part as a
+:   :octicons-milestone-24: Default: `0.01`. Minimum length of the silent part as a
     fraction of the total sound length.
 
 [`max_band_part`](#max_band_part){ #max_band_part }: `float` • range: [0.0, 1.0]
-:   :octicons-milestone-24: Default: `0.5`. Maximum length of the silent part as a
+:   :octicons-milestone-24: Default: `0.2`. Maximum length of the silent part as a
     fraction of the total sound length.
 
+[`fade_duration`](#fade_duration){ #fade_duration }: `float` • unit: seconds • range: 0.0 or [0.00025, ∞)
+: :octicons-milestone-24: Default: `0.005`. Duration of the fade-in and fade-out applied
+    at the edges of the silent region to smooth transitions and avoid abrupt
+    changes, which can otherwise produce impulses or clicks in the audio.
+    If you need hard edges or clicks, set this to `0.0` to disable fading.
+    Positive values must be at least 0.00025.
+
+[`mask_location`](#mask_location){ #mask_location }: `str` • choices: `"start"`, `"end"`, `"random"`
+: :octicons-milestone-24: Default: `random`. Where to place the silent region.
+    
+    * `"start"`: silence begins at index 0
+    * `"end"`: silence ends at the last sample
+    * `"random"`: silence starts at a random position
+
+[`p`](#p){ #p }: `float` • range: [0.0, 1.0]
+:   :octicons-milestone-24: Default: `0.5`. The probability of applying this transform.
+
+## Old TimeMask API (<=v0.40.0)
+
+:warning: This only applies to version 0.40.0 and older
+
 [`fade`](#fade){ #fade }: `bool`
 :   :octicons-milestone-24: Default: `False`. When set to `True`, a linear fade-in and fade-out is added to the silent part.
     This can smooth out unwanted abrupt changes between consecutive samples, which might
     otherwise sound like transients/clicks/pops.
 
-[`p`](#p){ #p }: `float` • range: [0.0, 1.0]
-:   :octicons-milestone-24: Default: `0.5`. The probability of applying this transform.
-
 ## Source code :octicons-mark-github-16:
 
 [audiomentations/augmentations/time_mask.py :octicons-link-external-16:](https://github.com/iver56/audiomentations/blob/main/audiomentations/augmentations/time_mask.py){target=_blank}


Test patch:
diff --git a/tests/test_time_mask.py b/tests/test_time_mask.py
index 636883e..f2638a0 100644
--- a/tests/test_time_mask.py
+++ b/tests/test_time_mask.py
@@ -1,14 +1,14 @@
 import numpy as np
 import pytest
 
-from audiomentations import TimeMask, Compose
+from audiomentations import TimeMask
 
 
 def test_apply_time_mask():
     sample_len = 1024
     samples_in = np.random.normal(0, 1, size=sample_len).astype(np.float32)
     sample_rate = 16000
-    augmenter = TimeMask(min_band_part=0.2, max_band_part=0.5, p=1.0)
+    augmenter = TimeMask(min_band_part=0.2, max_band_part=0.5, fade_duration=0.0, p=1.0)
 
     samples_out = augmenter(samples=samples_in, sample_rate=sample_rate)
     assert samples_out.dtype == np.float32
@@ -19,6 +19,31 @@ def test_apply_time_mask():
     assert std_out < std_in
 
 
+def test_apply_time_mask_start():
+    sample_len = 1000
+    samples_in = np.random.normal(0, 1, size=sample_len).astype(np.float32)
+    sample_rate = 8000
+    augmenter = TimeMask(min_band_part=0.1, max_band_part=0.1, fade_duration=0.0, mask_location="start", p=1.0)
+
+    samples_out = augmenter(samples=samples_in, sample_rate=sample_rate)
+    assert samples_out.dtype == np.float32
+    assert len(samples_out) == sample_len
+    assert not np.any(samples_out[0:100])
+    assert samples_out[101] == samples_in[101]
+
+def test_apply_time_mask_end():
+    sample_len = 1000
+    samples_in = np.random.normal(0, 1, size=sample_len).astype(np.float32)
+    sample_rate = 8000
+    augmenter = TimeMask(min_band_part=0.1, max_band_part=0.1, fade_duration=0.0, mask_location="end", p=1.0)
+
+    samples_out = augmenter(samples=samples_in, sample_rate=sample_rate)
+    assert samples_out.dtype == np.float32
+    assert len(samples_out) == sample_len
+    assert not np.any(samples_out[-100:])
+    assert samples_out[-101] == samples_in[-101]
+
+
 def test_invalid_params():
     with pytest.raises(ValueError):
         TimeMask(min_band_part=0.5, max_band_part=1.5)
@@ -29,6 +54,15 @@ def test_invalid_params():
     with pytest.raises(ValueError):
         TimeMask(min_band_part=0.6, max_band_part=0.5)
 
+    with pytest.raises(ValueError):
+        TimeMask(fade_duration=-0.1)
+
+    with pytest.raises(ValueError):
+        TimeMask(fade_duration=0.00001)
+
+    with pytest.raises(ValueError):
+        TimeMask(mask_location="beyond_infinity")
+
 
 def test_apply_time_mask_multichannel():
     sample_len = 1024
@@ -49,7 +83,7 @@ def test_apply_time_mask_with_fade():
     sample_len = 1024
     samples_in = np.random.normal(0, 1, size=sample_len).astype(np.float32)
     sample_rate = 16000
-    augmenter = TimeMask(min_band_part=0.2, max_band_part=0.5, fade=True, p=1.0)
+    augmenter = TimeMask(min_band_part=0.2, max_band_part=0.5, p=1.0)
 
     samples_out = augmenter(samples=samples_in, sample_rate=sample_rate)
     assert samples_out.dtype == np.float32
@@ -64,7 +98,7 @@ def test_apply_time_mask_with_fade_short_signal():
     sample_len = 100
     samples_in = np.random.normal(0, 1, size=sample_len).astype(np.float32)
     sample_rate = 16000
-    augmenter = TimeMask(min_band_part=0.2, max_band_part=0.5, fade=True, p=1.0)
+    augmenter = TimeMask(min_band_part=0.2, max_band_part=0.5, p=1.0)
 
     samples_out = augmenter(samples=samples_in, sample_rate=sample_rate)
     assert samples_out.dtype == np.float32