Repository: Lightning-Universe/lightning-flash

Issue:
Speech recognition data loading and augmentation
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

https://github.com/PyTorchLightning/lightning-flash/blob/4ebc45dd74df50e0a8b8d8e92efb97c6c0b9f6cb/flash/audio/speech_recognition/data.py#L72

At the moment the audio is getting loaded by soundfile, if it would be replaced with librosa more audio types are supported and resampling is done by default.

### Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

More audio formats like MP3 would be supported.
Furthermore an augmentation of data would be cool for training, especially with low ressource domains or languages.

### Pitch

<!-- A clear and concise description of what you want to happen. -->

Replacing soundfile with librosa and integrating audiomentations library.

### Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->
I am not familiar with lightning flash API, especially I dont know how to add augmentation for training only, when its the same method used for prediction and training.
I would like to add these two features, if you give some advice.
In the same step I could check adding Hubert, or replacing the hardcoded classes with automodel class, if it exists.

One more question:
What happens when I load an only pretrained model without feature extractor ? Does LF built its automatically ?


Discussion:
Hi @flozi00, thanks for your suggestions! Regarding augmentation support, transforms in flash are generally represented as a dict which maps hook names to the callable transforms to run on that hook. Any transforms you provide to the `SpeechRecognitionData` classmethod bound to the `pre_tensor_transform` hook will receive as input the output from the `load_sample` here: https://github.com/PyTorchLightning/lightning-flash/blob/4ebc45dd74df50e0a8b8d8e92efb97c6c0b9f6cb/flash/audio/speech_recognition/data.py#L75

So your data loading with some train augmentations would look like this:
```
from flash.core.data.transforms import ApplyToKeys

datamodule = SpeechRecognitionData.from_json(
    ...,
    train_transform={"pre_tensor_transform": ApplyToKeys("input", my_augmentation),
)
```

Currently, you're augmentation would have to apply to the output of soundfile, but if librosa enables more augmentations then that would be a welcome change :smiley: The main other change that would be needed is to convert back to the required numpy format before applying the processor from HF here: https://github.com/PyTorchLightning/lightning-flash/blob/4ebc45dd74df50e0a8b8d8e92efb97c6c0b9f6cb/flash/audio/speech_recognition/collate.py#L70

> What happens when I load an only pretrained model without feature extractor ? Does LF built its automatically ?

Not sure on that one, perhaps @SeanNaren can answer?

Looking forward to your PR :smiley:
Do you have an example of such an callable ?
I would use audiomentations: https://github.com/iver56/audiomentations

> but if librosa enables more augmentations then that would be a welcome change

librosa is the same as soundfile but powered by ffmpeg in background for files like .mp3
soundfile only supports files like .wav

audiomentations, soundfile and librosa have all the correct format for the processor
during the huggingface wav2vec sprint I tried a lot of these stuff and everything worked fine
> audiomentations, soundfile and librosa have all the correct format for the processor
> during the huggingface wav2vec sprint I tried a lot of these stuff and everything worked fine

That's awesome, switch to librosa should definitely work then, happy to have a PR for it :smiley:

> Do you have an example of such an callable ?
> I would use audiomentations: https://github.com/iver56/audiomentations

So in that case in the above example you could have:
```
my_augmentation = Compose([
    AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.015, p=0.5),
    TimeStretch(min_rate=0.8, max_rate=1.25, p=0.5),
    PitchShift(min_semitones=-4, max_semitones=4, p=0.5),
    Shift(min_fraction=-0.5, max_fraction=0.5, p=0.5),
])
```

Hope that helps :smiley:
That callable takes two inputs, the audio data and the sample rate
Is it possible to pass both arguments to the function or do I need to built my own one around it ?

https://github.com/iver56/audiomentations/blob/8bbd26537545c946d306b099aed46edc7dad727a/audiomentations/core/composition.py#L41-L54

Resolving patch:
diff --git a/flash/audio/speech_recognition/data.py b/flash/audio/speech_recognition/data.py
--- a/flash/audio/speech_recognition/data.py
+++ b/flash/audio/speech_recognition/data.py
@@ -35,7 +35,7 @@
 from flash.core.utilities.imports import _AUDIO_AVAILABLE, requires_extras
 
 if _AUDIO_AVAILABLE:
-    import soundfile as sf
+    import librosa
     from datasets import Dataset as HFDataset
     from datasets import load_dataset
     from transformers import Wav2Vec2CTCTokenizer
@@ -44,11 +44,16 @@
 
 
 class SpeechRecognitionDeserializer(Deserializer):
+    def __init__(self, sampling_rate: int):
+        super().__init__()
+
+        self.sampling_rate = sampling_rate
+
     def deserialize(self, sample: Any) -> Dict:
         encoded_with_padding = (sample + "===").encode("ascii")
         audio = base64.b64decode(encoded_with_padding)
         buffer = io.BytesIO(audio)
-        data, sampling_rate = sf.read(buffer)
+        data, sampling_rate = librosa.load(buffer, sr=self.sampling_rate)
         return {
             DefaultDataKeys.INPUT: data,
             DefaultDataKeys.METADATA: {"sampling_rate": sampling_rate},
@@ -61,7 +66,8 @@ def example_input(self) -> str:
 
 
 class BaseSpeechRecognition:
-    def _load_sample(self, sample: Dict[str, Any]) -> Any:
+    @staticmethod
+    def _load_sample(sample: Dict[str, Any], sampling_rate: int) -> Any:
         path = sample[DefaultDataKeys.INPUT]
         if (
             not os.path.isabs(path)
@@ -69,16 +75,17 @@ def _load_sample(self, sample: Dict[str, Any]) -> Any:
             and "root" in sample[DefaultDataKeys.METADATA]
         ):
             path = os.path.join(sample[DefaultDataKeys.METADATA]["root"], path)
-        speech_array, sampling_rate = sf.read(path)
+        speech_array, sampling_rate = librosa.load(path, sr=sampling_rate)
         sample[DefaultDataKeys.INPUT] = speech_array
         sample[DefaultDataKeys.METADATA] = {"sampling_rate": sampling_rate}
         return sample
 
 
 class SpeechRecognitionFileDataSource(DataSource, BaseSpeechRecognition):
-    def __init__(self, filetype: Optional[str] = None):
+    def __init__(self, sampling_rate: int, filetype: Optional[str] = None):
         super().__init__()
         self.filetype = filetype
+        self.sampling_rate = sampling_rate
 
     def load_data(
         self,
@@ -107,32 +114,44 @@ def load_data(
         ]
 
     def load_sample(self, sample: Dict[str, Any], dataset: Any = None) -> Any:
-        return self._load_sample(sample)
+        return self._load_sample(sample, self.sampling_rate)
 
 
 class SpeechRecognitionCSVDataSource(SpeechRecognitionFileDataSource):
-    def __init__(self):
-        super().__init__(filetype="csv")
+    def __init__(self, sampling_rate: int):
+        super().__init__(sampling_rate, filetype="csv")
 
 
 class SpeechRecognitionJSONDataSource(SpeechRecognitionFileDataSource):
-    def __init__(self):
-        super().__init__(filetype="json")
+    def __init__(self, sampling_rate: int):
+        super().__init__(sampling_rate, filetype="json")
 
 
 class SpeechRecognitionDatasetDataSource(DatasetDataSource, BaseSpeechRecognition):
+    def __init__(self, sampling_rate: int):
+        super().__init__()
+
+        self.sampling_rate = sampling_rate
+
     def load_data(self, data: Dataset, dataset: Optional[Any] = None) -> Union[Sequence[Mapping[str, Any]]]:
         if isinstance(data, HFDataset):
             data = list(zip(data["file"], data["text"]))
         return super().load_data(data, dataset)
 
+    def load_sample(self, sample: Dict[str, Any], dataset: Any = None) -> Any:
+        if isinstance(sample[DefaultDataKeys.INPUT], (str, Path)):
+            sample = self._load_sample(sample, self.sampling_rate)
+        return sample
+
 
 class SpeechRecognitionPathsDataSource(PathsDataSource, BaseSpeechRecognition):
-    def __init__(self):
-        super().__init__(("wav", "ogg", "flac", "mat"))
+    def __init__(self, sampling_rate: int):
+        super().__init__(("wav", "ogg", "flac", "mat", "mp3"))
+
+        self.sampling_rate = sampling_rate
 
     def load_sample(self, sample: Dict[str, Any], dataset: Any = None) -> Any:
-        return self._load_sample(sample)
+        return self._load_sample(sample, self.sampling_rate)
 
 
 class SpeechRecognitionPreprocess(Preprocess):
@@ -143,6 +162,7 @@ def __init__(
         val_transform: Optional[Dict[str, Callable]] = None,
         test_transform: Optional[Dict[str, Callable]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
+        sampling_rate: int = 16000,
     ):
         super().__init__(
             train_transform=train_transform,
@@ -150,13 +170,13 @@ def __init__(
             test_transform=test_transform,
             predict_transform=predict_transform,
             data_sources={
-                DefaultDataSources.CSV: SpeechRecognitionCSVDataSource(),
-                DefaultDataSources.JSON: SpeechRecognitionJSONDataSource(),
-                DefaultDataSources.FILES: SpeechRecognitionPathsDataSource(),
-                DefaultDataSources.DATASETS: SpeechRecognitionDatasetDataSource(),
+                DefaultDataSources.CSV: SpeechRecognitionCSVDataSource(sampling_rate),
+                DefaultDataSources.JSON: SpeechRecognitionJSONDataSource(sampling_rate),
+                DefaultDataSources.FILES: SpeechRecognitionPathsDataSource(sampling_rate),
+                DefaultDataSources.DATASETS: SpeechRecognitionDatasetDataSource(sampling_rate),
             },
             default_data_source=DefaultDataSources.FILES,
-            deserializer=SpeechRecognitionDeserializer(),
+            deserializer=SpeechRecognitionDeserializer(sampling_rate),
         )
 
     def get_state_dict(self) -> Dict[str, Any]:
diff --git a/flash/core/utilities/imports.py b/flash/core/utilities/imports.py
--- a/flash/core/utilities/imports.py
+++ b/flash/core/utilities/imports.py
@@ -88,7 +88,7 @@ def _compare_version(package: str, op, version) -> bool:
 _PIL_AVAILABLE = _module_available("PIL")
 _OPEN3D_AVAILABLE = _module_available("open3d")
 _SEGMENTATION_MODELS_AVAILABLE = _module_available("segmentation_models_pytorch")
-_SOUNDFILE_AVAILABLE = _module_available("soundfile")
+_LIBROSA_AVAILABLE = _module_available("librosa")
 _TORCH_SCATTER_AVAILABLE = _module_available("torch_scatter")
 _TORCH_SPARSE_AVAILABLE = _module_available("torch_sparse")
 _TORCH_GEOMETRIC_AVAILABLE = _module_available("torch_geometric")
@@ -148,7 +148,7 @@ class Image(metaclass=MetaImage):
 )
 _SERVE_AVAILABLE = _FASTAPI_AVAILABLE and _PYDANTIC_AVAILABLE and _CYTOOLZ_AVAILABLE and _UVICORN_AVAILABLE
 _POINTCLOUD_AVAILABLE = _OPEN3D_AVAILABLE and _TORCHVISION_AVAILABLE
-_AUDIO_AVAILABLE = all([_TORCHAUDIO_AVAILABLE, _SOUNDFILE_AVAILABLE, _TRANSFORMERS_AVAILABLE])
+_AUDIO_AVAILABLE = all([_TORCHAUDIO_AVAILABLE, _LIBROSA_AVAILABLE, _TRANSFORMERS_AVAILABLE])
 _GRAPH_AVAILABLE = _TORCH_SCATTER_AVAILABLE and _TORCH_SPARSE_AVAILABLE and _TORCH_GEOMETRIC_AVAILABLE
 
 _EXTRAS_AVAILABLE = {


Test patch: