feat: add loudness transform, alltransform, refactor into folders

Author: flavioschneider

audio_data_pytorch/__init__.py

@@ -1,4 +1,2 @@
-from .ljspeech_dataset import LJSpeechDataset
-from .transforms import Crop, OverlapChannels, RandomCrop, Resample, Scale
-from .wav_dataset import WAVDataset
-from .youtube_dataset import YoutubeDataset
+from .datasets import *  # noqa
+from .transforms import *  # noqa

audio_data_pytorch/datasets/__init__.py

@@ -0,0 +1,3 @@
+from .ljspeech_dataset import LJSpeechDataset
+from .wav_dataset import WAVDataset
+from .youtube_dataset import YoutubeDataset

audio_data_pytorch/ljspeech_dataset.py renamed to audio_data_pytorch/datasets/ljspeech_dataset.py

@@ -4,7 +4,7 @@
 import requests  # type: ignore
 from tqdm import tqdm
 
-from .utils import camel_to_snake
+from ..utils import camel_to_snake
 from .wav_dataset import WAVDataset
 
 

audio_data_pytorch/wav_dataset.py renamed to audio_data_pytorch/datasets/wav_dataset.py

@@ -8,7 +8,7 @@
 from torch.nn import functional as F
 from tqdm import tqdm
 
-from .utils import camel_to_snake, exists
+from ..utils import camel_to_snake, exists
 from .wav_dataset import WAVDataset
 
 

audio_data_pytorch/youtube_dataset.py renamed to audio_data_pytorch/datasets/youtube_dataset.py

@@ -0,0 +1,7 @@
+from .all import AllTransform
+from .crop import Crop
+from .loudness import Loudness
+from .overlap_channels import OverlapChannels
+from .randomcrop import RandomCrop
+from .resample import Resample
+from .scale import Scale

audio_data_pytorch/transforms.py

@@ -0,0 +1,47 @@
+from typing import Optional
+
+from torch import Tensor, nn
+
+from ..utils import exists
+from .crop import Crop
+from .loudness import Loudness
+from .overlap_channels import OverlapChannels
+from .randomcrop import RandomCrop
+from .resample import Resample
+from .scale import Scale
+
+
+class AllTransform(nn.Module):
+    def __init__(
+        self,
+        source_rate: Optional[int] = None,
+        target_rate: Optional[int] = None,
+        crop_size: Optional[int] = None,
+        random_crop_size: Optional[int] = None,
+        loudness: Optional[int] = None,
+        scale: Optional[float] = None,
+        overlap_channels: bool = False,
+    ):
+        super().__init__()
+
+        message = "Both source_rate and target_rate must be provided"
+        assert not exists(source_rate) ^ exists(target_rate), message
+
+        message = "Loudness requires target_rate"
+        assert not exists(loudness) or exists(target_rate), message
+
+        self.transform = nn.Sequential(
+            Resample(source=source_rate, target=target_rate)  # type: ignore
+            if exists(source_rate) and source_rate != target_rate
+            else nn.Identity(),
+            RandomCrop(random_crop_size) if exists(random_crop_size) else nn.Identity(),
+            Crop(crop_size) if exists(crop_size) else nn.Identity(),
+            OverlapChannels() if overlap_channels else nn.Identity(),
+            Loudness(sampling_rate=target_rate, target=loudness)  # type: ignore
+            if exists(loudness)
+            else nn.Identity(),
+            Scale(scale) if exists(scale) else nn.Identity(),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.transform(x)

audio_data_pytorch/transforms/__init__.py

@@ -0,0 +1,22 @@
+import torch
+from torch import Tensor, nn
+
+
+class Crop(nn.Module):
+    """Crops waveform to fixed size"""
+
+    def __init__(self, size: int, start: int = 0) -> None:
+        super().__init__()
+        self.size = size
+        self.start = start
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x[:, self.start :]
+        channels, length = x.shape
+
+        if length < self.size:
+            padding_length = self.size - length
+            padding = torch.zeros(channels, padding_length).to(x)
+            return torch.cat([x, padding], dim=1)
+        else:
+            return x[:, 0 : self.size]

audio_data_pytorch/transforms/all.py

@@ -0,0 +1,28 @@
+import pyloudnorm as pyln
+import torch
+from torch import Tensor, nn
+
+
+class Loudness(nn.Module):
+    """Normalizes to target loudness using BS.1770-4, requires pyloudnorm"""
+
+    def __init__(self, sampling_rate: int, target: float):
+        super().__init__()
+        self.sampling_rate = sampling_rate
+        self.target = target
+        self.meter = pyln.Meter(sampling_rate)
+
+    def forward(self, x: Tensor) -> Tensor:
+        channels, length = x.shape
+        # Measure sample loudness
+        x_numpy = x.numpy().T
+        loudness = self.meter.integrated_loudness(data=x_numpy)
+        # Don't normalize zeros sample (i.e. silence)
+        if loudness == -float("inf"):
+            return x
+        # Normalize sample loudness
+        x_normalized = pyln.normalize.loudness(
+            data=x_numpy, input_loudness=loudness, target_loudness=self.target
+        )
+        # Return normalized as torch Tensor
+        return torch.from_numpy(x_normalized.T)