audio io

enhancer/utils/io.py

@@ -0,0 +1,68 @@
+import os
+import librosa
+from typing import Optional
+from matplotlib.pyplot import axis
+import numpy as np
+import torch
+
+class Audio:
+
+    def __init__(
+        self,
+        sampling_rate:int=16000,
+        mono:bool=True,
+        return_tensor=True
+    ) -> None:
+        
+        self.sampling_rate = sampling_rate
+        self.mono = mono
+        self.return_tensor = return_tensor
+
+    def __call__(
+        self,
+        audio,
+        sampling_rate:Optional[int]=None,
+        offset:Optional[float] = None,
+        duration:Optional[float] = None
+    ):
+        if isinstance(audio,str):
+            if os.path.exists(audio):
+                audio,sampling_rate = librosa.load(audio,sr=sampling_rate,mono=False,
+                offset=offset,duration=duration)
+            else:
+                raise FileNotFoundError(f"File {audio} deos not exist")
+        elif isinstance(audio,np.ndarray):
+            if len(audio.shape) == 1:
+                audio = audio.reshape(1,-1)
+        else:
+            raise ValueError("audio should be either filepath or numpy ndarray")
+
+        if self.mono:
+            audio = self.convert_mono(audio)
+
+        resampled_audio =  self.resample_audio(audio,sampling_rate)
+        if self.return_tensor:
+            return torch.tensor(resampled_audio)
+        else:
+            return resampled_audio
+
+    def convert_mono(
+        self,
+        audio
+
+    ):
+        num_channels,num_samples = audio.shape
+        if num_channels>1 and self.mono:
+            return audio.mean(axis=0).reshape(1,num_samples)
+        return audio
+
+
+    def resample_audio(
+        self,
+        audio,
+        sampling_rate
+    ):
+        if self.sampling_rate!=sampling_rate:
+            audio = librosa.resample(audio,orig_sr=sampling_rate,target_sr=self.sampling_rate)
+
+        return audio