authtransforms 0.1.0__py3-none-any.whl

authtransforms/transforms.py

@@ -0,0 +1,648 @@
+"""
+Audio augmentation transforms for PyTorch, inspired by:
+https://jonathanbgn.com/2021/08/30/audio-augmentation.html
+
+Transforms follow the torchvision convention: callable objects with __call__
+that accept a (channels, samples) tensor and return a transformed tensor.
+"""
+
+import math
+import os
+import pathlib
+import random
+
+import torch
+import torchaudio
+import torchaudio.transforms as T
+import tempfile
+import random
+import subprocess
+import scipy.signal
+import numpy as np
+from pathlib import Path
+from scipy.signal import fftconvolve
+import numpy as np
+import librosa
+
+
+
+# ---------------------------------------------------------------------------
+# Core transforms
+# ---------------------------------------------------------------------------
+
+
+class RandomClip:
+    """Extract a random fixed-length clip from the audio.
+
+    Args:
+        sample_rate: Sample rate of the audio.
+        clip_length: Desired clip length in samples.
+        vad: If True, apply Voice Activity Detection to trim leading/trailing
+             silence after clipping.
+        vad_trigger_level: VAD sensitivity (higher = more aggressive trimming).
+    """
+
+    def __init__(self, sample_rate: int, clip_length: int, vad: bool = True, vad_trigger_level: float = 7.0):
+        self.clip_length = clip_length
+        self.vad = T.Vad(sample_rate=sample_rate, trigger_level=vad_trigger_level) if vad else None
+        #VAD is Voice Activity Detection
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        audio_length = audio.shape[-1]
+        if audio_length > self.clip_length:
+            offset = random.randint(0, audio_length - self.clip_length) #to select random int position where offset is the int
+            audio = audio[..., offset : offset + self.clip_length]  #Torch slicing
+        if self.vad is not None:
+            audio = self.vad(audio)
+        return audio
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(clip_length={self.clip_length})"
+
+
+class RandomSpeedChange:
+    """Randomly perturb audio playback speed.
+
+    Resamples back to the original sample rate so the output length stays
+    approximately the same.  Uses torchaudio's SoX effects backend.
+
+    Args:
+        sample_rate: Sample rate of the audio.
+        speed_factors: Sequence of speed multipliers to choose from.
+                       Default is (0.9, 1.0, 1.1) as used in the literature.
+    """
+
+    def __init__(self, sample_rate: int, speed_factors=(0.9, 1.0, 1.1)):
+        self.sample_rate = sample_rate
+        self.speed_factors = list(speed_factors)
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        speed_factor = random.choice(self.speed_factors)
+        if speed_factor == 1.0:
+            return audio
+
+        sox_effects = [                                  #This is the sox effect that changes the speed and sample_rate of the audio
+            ["speed", str(speed_factor)],
+            ["rate", str(self.sample_rate)],
+        ]
+        transformed, _ = torchaudio.sox_effects.apply_effects_tensor(
+            audio, self.sample_rate, sox_effects
+        )
+        return transformed
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(speed_factors={self.speed_factors})"
+
+
+class RandomBackgroundNoise:
+    """Mix the audio with a random noise file at a random SNR level.
+
+    Args:
+        sample_rate: Target sample rate.
+        noise_dir: Directory (searched recursively) containing .wav noise files.
+        min_snr_db: Minimum signal-to-noise ratio in dB.
+        max_snr_db: Maximum signal-to-noise ratio in dB.
+    """
+
+    def __init__(self, sample_rate: int, noise_dir: str, min_snr_db: float = 0, max_snr_db: float = 15):
+        self.sample_rate = sample_rate
+        self.min_snr_db = min_snr_db
+        self.max_snr_db = max_snr_db
+
+        noise_path = pathlib.Path(noise_dir)
+        if not noise_path.exists():                                #To check if the noise path exists
+            raise IOError(f"Noise directory `{noise_dir}` does not exist") 
+        self.noise_files = list(noise_path.glob("**/*.wav"))                # Find all the files inside the noise path
+        if not self.noise_files:
+            raise IOError(f"No .wav files found in `{noise_dir}`")
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        noise_file = random.choice(self.noise_files)
+        effects = [
+            ["remix", "1"],              # convert to mono (noise)
+            ["rate", str(self.sample_rate)],  # resample (noise)
+        ]
+        noise, _ = torchaudio.sox_effects.apply_effects_file(noise_file, effects, normalize=True)
+
+        audio_length = audio.shape[-1]
+        noise_length = noise.shape[-1]
+
+        if noise_length > audio_length:
+            offset = random.randint(0, noise_length - audio_length)  #To make the noise length equal to audio length
+            noise = noise[..., offset : offset + audio_length]
+        elif noise_length < audio_length:
+            # tile noise to cover the audio
+            repeats = math.ceil(audio_length / noise_length)
+            noise = noise.repeat(1, repeats)[..., :audio_length]
+
+        snr_db = random.uniform(self.min_snr_db, self.max_snr_db)   #The snr defines whether the noise or audio is gonna dominate , high snr signal                                                                     dominates, whereas low snr noise dominates
+        snr = math.exp(snr_db / 10)
+        audio_power = audio.norm(p=2)          #Computing both the audio and noise power
+        noise_power = noise.norm(p=2)
+        scale = snr * noise_power / (audio_power + 1e-9)
+
+        return (scale * audio + noise) / 2    #Mix the audio and noise
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"snr=[{self.min_snr_db}, {self.max_snr_db}] dB, "
+            f"n_files={len(self.noise_files)})"
+        )
+
+
+class RandomPitchShift:
+
+    # This changes the depth of the voice , variation in male/female/child voices
+    """Shift audio pitch by a random number of semitones.
+
+    Args:
+        sample_rate: Sample rate of the audio.
+        semitones: Range tuple (min, max) of semitone shifts.
+    """
+
+    def __init__(self, sample_rate: int, semitones: tuple = (-2, 2)):  #Semitone is the range of pitchshift
+        self.sample_rate = sample_rate
+        self.semitones = semitones
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        n = random.uniform(*self.semitones)
+        if n == 0:
+            return audio
+        sox_effects = [["pitch", str(int(n * 100))], ["rate", str(self.sample_rate)]] #1 semitone = 100 cents and the sample rate is resampled to original
+        transformed, _ = torchaudio.sox_effects.apply_effects_tensor(
+            audio, self.sample_rate, sox_effects
+        )
+        return transformed
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(semitones={self.semitones})"
+
+
+class RandomGain:
+    """Scale audio amplitude by a random gain factor (in dB).
+
+    Args:
+        min_gain_db: Minimum gain in dB.
+        max_gain_db: Maximum gain in dB.
+    """
+
+    def __init__(self, min_gain_db: float = -6, max_gain_db: float = 6):
+        self.min_gain_db = min_gain_db
+        self.max_gain_db = max_gain_db
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        gain_db = random.uniform(self.min_gain_db, self.max_gain_db)   # This is to make the audio louder or quieter
+        gain = 10 ** (gain_db / 20 )       #This is the formula that changes the amplitude to linear scales and multiply with audio. 
+        return audio * gain
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(gain_db=[{self.min_gain_db}, {self.max_gain_db}])"
+
+
+class AddGaussianNoise:
+    """Add Gaussian white noise at a random amplitude.
+
+    Args:
+        min_amplitude: Minimum noise standard deviation.
+        max_amplitude: Maximum noise standard deviation.
+    """
+
+    def __init__(self, min_amplitude: float = 0.001, max_amplitude: float = 0.015):
+        self.min_amplitude = min_amplitude         #min strength of noise
+        self.max_amplitude = max_amplitude         #max strength of noise
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        amplitude = random.uniform(self.min_amplitude, self.max_amplitude)
+        noise = torch.randn_like(audio) * amplitude   #randn_like generates random values 
+        return audio + noise
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"amplitude=[{self.min_amplitude}, {self.max_amplitude}])"
+        )
+
+
+class TimeShift:
+    """Shift the audio waveform in time, wrapping or zero-padding.
+
+    Args:
+        max_shift: Maximum shift as a fraction of the total length (0–1).
+        roll: If True, wrap the audio (circular shift). If False, zero-pad.
+    """
+
+    def __init__(self, max_shift: float = 0.2, roll: bool = False):
+        self.max_shift = max_shift
+        self.roll = roll
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        length = audio.shape[-1]
+        shift = random.randint(-int(length * self.max_shift), int(length * self.max_shift))
+        if shift == 0:
+            return audio
+        if self.roll:
+            return torch.roll(audio, shift, dims=-1)
+        # zero-pad version
+        result = torch.zeros_like(audio)
+        if shift > 0:
+            result[..., shift:] = audio[..., : length - shift]
+        else:
+            result[..., : length + shift] = audio[..., -shift:]
+        return result
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(max_shift={self.max_shift}, roll={self.roll})"
+
+
+class SpecAugment:
+    """Apply SpecAugment on the spectrogram (time & frequency masking).
+
+    This transform expects a raw audio tensor and internally computes a
+    spectrogram, applies masking, then returns the masked spectrogram.
+    Use it as the final step in a pipeline when your model consumes spectrograms.
+
+    Args:
+        sample_rate: Sample rate.
+        n_fft: FFT size.
+        n_mels: Number of mel filterbanks.
+        freq_mask_param: Maximum frequency mask width.
+        time_mask_param: Maximum time mask width.
+        n_freq_masks: Number of frequency masks.
+        n_time_masks: Number of time masks.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        n_fft: int = 400,
+        n_mels: int = 80,
+        freq_mask_param: int = 27,
+        time_mask_param: int = 100,
+        n_freq_masks: int = 2,
+        n_time_masks: int = 2,
+    ):
+        self.mel_spec = T.MelSpectrogram(sample_rate=sample_rate, n_fft=n_fft, n_mels=n_mels)
+        self.freq_masking = T.FrequencyMasking(freq_mask_param=freq_mask_param)
+        self.time_masking = T.TimeMasking(time_mask_param=time_mask_param)
+        self.n_freq_masks = n_freq_masks
+        self.n_time_masks = n_time_masks
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        spec = self.mel_spec(audio)
+        for _ in range(self.n_freq_masks):
+            spec = self.freq_masking(spec)
+        for _ in range(self.n_time_masks):
+            spec = self.time_masking(spec)
+        return spec
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(n_freq={self.n_freq_masks}, n_time={self.n_time_masks})"
+
+
+class RandomApply:
+    """Apply a transform with probability p (mirrors torchvision.transforms.RandomApply).
+
+    Args:
+        transform: A callable transform.
+        p: Probability of applying the transform (0–1).
+    """
+
+    def __init__(self, transform, p: float = 0.5):
+        self.transform = transform
+        self.p = p
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if random.random() < self.p:
+            return self.transform(audio)
+        return audio
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(p={self.p}, transform={self.transform})"
+
+
+class Normalize:
+    """Normalize audio to peak amplitude of 1.0 (or a target level).
+
+    Args:
+        target_db: Target peak level in dBFS. Default -3 dBFS. Use None for 0 dBFS.
+    """
+
+    def __init__(self, target_db: float = -3.0):
+        self.target_level = 10 ** (target_db / 20)
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        peak = audio.abs().max()
+        if peak < 1e-9:
+            return audio
+        return audio / peak * self.target_level
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(target_db={20 * math.log10(self.target_level):.1f})"
+
+
+class ToMono:
+    """Convert multi-channel audio to mono by averaging channels."""
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if audio.shape[0] > 1:
+            return audio.mean(dim=0, keepdim=True)
+        return audio
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+class CodecAugmentation:
+    """Apply random lossy codec compression (AAC, MP3, Opus, Vorbis, G.722).
+
+    Args:
+        sample_rate: Sample rate of the audio.
+        codecs: Sequence of codecs to sample from.
+        p: Probability of applying augmentation.
+        bitrate_range: Sequence of bitrates (bps) to choose from.
+        mp3_vbr_prob: Probability of using VBR for MP3.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        codecs: tuple = ("aac", "mp3", "opus", "vorbis","g722"),
+        p: float = 0.5,
+        bitrate_range: tuple = (12000, 16000, 24000, 32000),
+        mp3_vbr_prob: float = 0.3,
+    ):
+        self.sample_rate = sample_rate
+        self.codecs = list(codecs)
+        self.p = p
+        self.bitrate_range = list(bitrate_range)
+        self.mp3_vbr_prob = mp3_vbr_prob
+
+        self.codec_map = {
+            "aac": ("aac", ".m4a"),
+            "mp3": ("libmp3lame", ".mp3"),
+            "opus": ("libopus", ".ogg"),
+            "vorbis": ("libvorbis", ".ogg"),
+            "g722": ("g722", ".g722"),
+        }
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if random.random() > self.p:
+            return audio
+
+        codec = random.choice(self.codecs)
+        encoder, ext = self.codec_map[codec]
+        bitrate = random.choice(self.bitrate_range)
+        target_sr = 8000 if codec == "g722" else self.sample_rate
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as src, \
+             tempfile.NamedTemporaryFile(suffix=ext, delete=False) as enc, \
+             tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as dec:
+
+            try:
+                torchaudio.save(src.name, audio, self.sample_rate)
+
+                # Encode with ffmpeg
+                cmd = [
+                    "ffmpeg", "-y",
+                    "-i", src.name,
+                    "-ar", str(target_sr),
+                    "-c:a", encoder,
+                ]
+
+                if codec == "mp3" and random.random() < self.mp3_vbr_prob:
+                    cmd += ["-q:a", str(random.randint(2, 6))]
+                else:
+                    cmd += ["-b:a", str(bitrate)]
+
+                cmd.append(enc.name)
+                subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
+
+                # Decode back to WAV
+                subprocess.run([
+                    "ffmpeg", "-y",
+                    "-i", enc.name,
+                    "-ar", str(self.sample_rate),
+                    dec.name
+                ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
+
+                augmented, _ = torchaudio.load(dec.name)
+                return augmented
+
+            finally:
+                for f in [src.name, enc.name, dec.name]:
+                    if os.path.exists(f):
+                        os.remove(f)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(sample_rate={self.sample_rate}, codecs={self.codecs}, p={self.p})"
+    
+class FilterAugmentation:
+    """Apply a band-pass filter to simulate speech frequency range.
+    Args:
+        sample_rate: Sample rate of the audio.
+        low_hz: Lower cutoff frequency in Hz.
+        high_hz: Upper cutoff frequency in Hz.
+        p: Probability of applying augmentation.
+    """
+    def __init__(
+        self,
+        sample_rate: int,
+        low_hz: float = 50.0,
+        high_hz: float = 7000.0,
+        p: float = 0.5,
+    ):
+        self.sample_rate = sample_rate
+        self.low_hz = low_hz
+        self.high_hz = high_hz
+        self.p = p
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if random.random() > self.p:
+            return audio
+
+        # Normalise cutoffs to [0, 1] where 1 = Nyquist frequency
+        nyquist = self.sample_rate / 2
+        low = self.low_hz / nyquist
+        high = self.high_hz / nyquist
+
+        # Design a Butterworth band-pass filter
+        sos = scipy.signal.butter(
+            N=4,                        # filter order
+            Wn=[low, high],
+            btype="bandpass",
+            output="sos"                # numerically stable form
+        )
+
+        # Apply filter to each channel
+        audio_np = audio.numpy()
+        filtered = scipy.signal.sosfiltfilt(sos, audio_np)  # zero-phase filtering
+        filtered = filtered.copy()
+        return torch.from_numpy(filtered).float()
+
+    def __repr__(self) -> str:
+        return (f"{self.__class__.__name__}("
+                f"low_hz={self.low_hz}, high_hz={self.high_hz}, p={self.p})")
+    
+    
+class RoomImpulseResponse:
+    """
+    Convolve audio with a randomly sampled Room Impulse Response (RIR).
+
+    Args:
+        rir_dir:        Directory containing .wav RIR files.
+        sample_rate:    Expected sample rate of audio and RIR files.
+        normalize:      Normalise output to prevent clipping after convolution.
+        p:              Probability of applying augmentation.
+    """
+    def __init__(
+        self,
+        rir_dir: str,
+        sample_rate: int,
+        normalize: bool = True,
+        p: float = 0.5,
+    ):
+        self.sample_rate = sample_rate
+        self.normalize = normalize
+        self.p = p
+
+        # Pre-load all RIR paths at init — avoids repeated filesystem scans
+        self.rir_paths = list(Path(rir_dir).rglob("*.wav"))
+        if not self.rir_paths:
+            raise ValueError(f"No .wav files found in {rir_dir}")
+
+    def _load_rir(self, path: Path) -> np.ndarray:
+        import soundfile as sf
+        rir, sr = sf.read(path, always_2d=True)  # (samples, channels)
+        rir = rir.T                               # → (channels, samples)
+        if sr != self.sample_rate:
+            rir = librosa.resample(rir, orig_sr=sr, target_sr=self.sample_rate)
+        return rir.astype(np.float32)
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if random.random() > self.p:
+            return audio
+
+        rir_path = random.choice(self.rir_paths)
+        rir = self._load_rir(rir_path)
+
+        audio_np = audio.numpy()                  # (channels, samples)
+
+        # Match channel count — mono RIR applied to stereo audio
+        if rir.shape[0] == 1 and audio_np.shape[0] > 1:
+            rir = np.repeat(rir, audio_np.shape[0], axis=0)
+        elif rir.shape[0] > audio_np.shape[0]:
+            rir = rir[:audio_np.shape[0]]
+
+        # Convolve each channel independently
+        convolved = np.stack([
+            fftconvolve(audio_np[c], rir[c])[:audio_np.shape[1]]
+            for c in range(audio_np.shape[0])
+        ])
+
+        # Normalise to original peak — convolution can cause large amplitude spikes
+        if self.normalize:
+            peak = np.abs(convolved).max()
+            if peak > 0:
+                original_peak = np.abs(audio_np).max()
+                convolved = convolved * (original_peak / peak)
+
+        return torch.from_numpy(convolved).float()
+
+    def __repr__(self) -> str:
+        return (f"{self.__class__.__name__}("
+                f"rir_dir='{self.rir_paths[0].parent}', "
+                f"n_rirs={len(self.rir_paths)}, p={self.p})")
+    
+class TimeStretch:
+    """
+    Stretch or compress audio in time without changing pitch.
+    Uses a phase vocoder via librosa.
+
+    Args:
+        min_rate: Minimum stretch rate (< 1.0 = slower, longer audio).
+        max_rate: Maximum stretch rate (> 1.0 = faster, shorter audio).
+        p:        Probability of applying augmentation.
+    """
+    def __init__(
+        self,
+        min_rate: float = 0.8,
+        max_rate: float = 1.25,
+        p: float = 0.5,
+    ):
+        if min_rate <= 0 or max_rate <= 0:
+            raise ValueError("Rates must be positive")
+        if min_rate > max_rate:
+            raise ValueError("min_rate must be <= max_rate")
+        
+        self.min_rate = min_rate
+        self.max_rate = max_rate
+        self.p = p
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if random.random() > self.p:
+            return audio
+
+        rate = random.uniform(self.min_rate, self.max_rate)
+
+        audio_np = audio.numpy()  # (channels, samples)
+
+        # Stretch each channel independently
+        stretched = np.stack([
+            librosa.effects.time_stretch(audio_np[c], rate=rate)
+            for c in range(audio_np.shape[0])
+        ])
+
+        return torch.from_numpy(stretched).float()
+
+    def __repr__(self) -> str:
+        return (f"{self.__class__.__name__}("
+                f"min_rate={self.min_rate}, max_rate={self.max_rate}, p={self.p})")
+    
+class ResampleAugmentation:
+    """
+    Simulate bandwidth degradation by resampling to an intermediate
+    sample rate and back.
+
+    Args:
+        original_freq:      Sample rate of the input audio.
+        target_freq:        Intermediate sample rate to resample to.
+                            Lower = ResampleDown, Higher = ResampleUp.
+        p:                  Probability of applying augmentation.
+    """
+    def __init__(
+        self,
+        original_freq: int,
+        target_freq: int,
+        p: float = 0.5,
+    ):
+        self.original_freq = original_freq
+        self.target_freq = target_freq
+        self.p = p
+
+    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
+        if random.random() > self.p:
+            return audio
+
+        # Step 1: resample to intermediate rate
+        degraded = torchaudio.functional.resample(
+            audio,
+            orig_freq=self.original_freq,
+            new_freq=self.target_freq,
+        )
+
+        # Step 2: resample back to original rate
+        restored = torchaudio.functional.resample(
+            degraded,
+            orig_freq=self.target_freq,
+            new_freq=self.original_freq,
+        )
+
+        return restored
+
+    def __repr__(self) -> str:
+        direction = "Down" if self.target_freq < self.original_freq else "Up"
+        return (
+            f"{self.__class__.__name__}(Resample{direction}, "
+            f"{self.original_freq}Hz → {self.target_freq}Hz → {self.original_freq}Hz, "
+            f"p={self.p})"
+        )