Add resample preprocessor

Author: istupakov

preprocessors/__init__.py

@@ -1,12 +1,14 @@
 from .gigaam import GigaamPreprocessor
 from .kaldi import KaldiPreprocessor
 from .nemo import NemoPreprocessor
+from .resample import ResamplePreprocessor
 from .whisper import WhisperPreprocessor80, WhisperPreprocessor128
 
 __all__ = [
     "GigaamPreprocessor",
     "KaldiPreprocessor",
     "NemoPreprocessor",
+    "ResamplePreprocessor",
     "WhisperPreprocessor80",
     "WhisperPreprocessor128",
 ]

preprocessors/build.py

@@ -24,3 +24,4 @@ def build():
     save_model(preprocessors.NemoPreprocessor, preprocessors_dir.joinpath("nemo.onnx"))
     save_model(preprocessors.WhisperPreprocessor80, preprocessors_dir.joinpath("whisper80.onnx"))
     save_model(preprocessors.WhisperPreprocessor128, preprocessors_dir.joinpath("whisper128.onnx"))
+    save_model(preprocessors.ResamplePreprocessor, preprocessors_dir.joinpath("resample.onnx"))

preprocessors/resample.py

@@ -0,0 +1,54 @@
+import math
+
+import torch
+import torchaudio
+from onnx import numpy_helper
+from onnxscript import FLOAT, INT64, script
+from onnxscript import opset17 as op
+
+
+def make_kernel(orig_freq: int):
+    new_freq = 16_000
+    gcd = math.gcd(orig_freq, new_freq)
+    kernel, width = torchaudio.functional.functional._get_sinc_resample_kernel(orig_freq, new_freq, gcd, dtype=torch.float32)
+    return kernel.numpy()[:, None], width, orig_freq // gcd, new_freq // gcd
+
+
+kernel08, width08, orig_freq08, new_freq08 = make_kernel(8_000)
+kernel22, width22, orig_freq22, new_freq22 = make_kernel(22_050)
+kernel44, width44, orig_freq44, new_freq44 = make_kernel(44_100)
+kernel48, width48, orig_freq48, new_freq48 = make_kernel(48_000)
+
+
+@script(doc_string="Resampling waveform to 16 kHz")
+def ResamplePreprocessor(
+    waveforms: FLOAT["batch_size", "N"],
+    waveforms_lens: INT64["batch_size"],
+    sample_rate: INT64["1"],
+) -> tuple[FLOAT["batch_size", "M"], INT64["batch_size"]]:
+    waveforms = op.Unsqueeze(waveforms, axes=[1, 2])
+
+    if sample_rate[0] == 8_000:
+        kernel = op.Constant(value=numpy_helper.from_array(kernel08, "kernel"))
+        conv = op.Conv(waveforms, kernel, pads=(0, width08, 0, width08 + orig_freq08), strides=(1, orig_freq08))
+        waveforms_lens = (new_freq08 * waveforms_lens + orig_freq08 - 1) / orig_freq08
+    elif sample_rate[0] == 22_050:
+        kernel = op.Constant(value=numpy_helper.from_array(kernel22, "kernel"))
+        conv = op.Conv(waveforms, kernel, pads=(0, width22, 0, width22 + orig_freq22), strides=(1, orig_freq22))
+        waveforms_lens = (new_freq22 * waveforms_lens + orig_freq22 - 1) / orig_freq22
+    elif sample_rate[0] == 44_100:
+        kernel = op.Constant(value=numpy_helper.from_array(kernel44, "kernel"))
+        conv = op.Conv(waveforms, kernel, pads=(0, width44, 0, width44 + orig_freq44), strides=(1, orig_freq44))
+        waveforms_lens = (new_freq44 * waveforms_lens + orig_freq44 - 1) / orig_freq44
+    elif sample_rate[0] == 48_000:
+        kernel = op.Constant(value=numpy_helper.from_array(kernel48, "kernel"))
+        conv = op.Conv(waveforms, kernel, pads=(0, width48, 0, width48 + orig_freq48), strides=(1, orig_freq48))
+        waveforms_lens = (new_freq48 * waveforms_lens + orig_freq48 - 1) / orig_freq48
+    else:
+        conv = waveforms
+
+    resampled_lens = op.Identity(waveforms_lens)
+    max_len = op.ReduceMax(resampled_lens, keepdims=0)
+    mask = op.Unsqueeze(op.Range(0, max_len, 1), [0]) < op.Unsqueeze(resampled_lens, [1])
+    resampled = op.Where(mask, op.Flatten(op.Transpose(conv, perm=(0, 3, 2, 1)))[:, :max_len], 0)
+    return resampled, resampled_lens

src/onnx_asr/asr.py

@@ -10,7 +10,7 @@
 import numpy.typing as npt
 
 from .preprocessors import Preprocessor
-from .utils import pad_list, read_wav_files
+from .utils import SampleRates, pad_list, read_wav_files
 
 
 class Asr(ABC):
@@ -21,26 +21,35 @@ class Asr(ABC):
     def _get_model_files(quantization: str | None = None) -> dict[str, str]: ...
 
     @abstractmethod
-    def _recognize_batch(self, waveforms: list[npt.NDArray[np.float32]], language: str | None = None) -> list[str]: ...
+    def _recognize_batch(
+        self, waveforms: list[npt.NDArray[np.float32]], sample_rate: SampleRates, language: str | None
+    ) -> list[str]: ...
 
     def recognize(
-        self, waveform: str | npt.NDArray[np.float32] | list[str | npt.NDArray[np.float32]], language: str | None = None
+        self,
+        waveform: str | npt.NDArray[np.float32] | list[str | npt.NDArray[np.float32]],
+        *,
+        sample_rate: SampleRates = 16_000,
+        language: str | None = None,
     ) -> str | list[str]:
         """Recognize speech (single or batch).
 
         Args:
             waveform: Path to wav file (only PCM_U8, PCM_16, PCM_24 and PCM_32 formats with 16 kHz sample rate are supported)
                       or Numpy array with PCM waveform.
                       A list of file paths or numpy arrays for batch recognition are also supported.
+            sample_rate: Sample rate for Numpy arrays in waveform.
             language: Speech language (only for Whisper models).
 
         Returns:
             Speech recognition results (single string or list for batch recognition).
 
         """
         if isinstance(waveform, list):
-            return self._recognize_batch(read_wav_files(waveform), language)
-        return self._recognize_batch(read_wav_files([waveform]), language)[0]
+            if not waveform:
+                return []
+            return self._recognize_batch(*read_wav_files(waveform, sample_rate), language)
+        return self._recognize_batch(*read_wav_files([waveform], sample_rate), language)[0]
 
 
 class _AsrWithDecoding(Asr):
@@ -65,8 +74,12 @@ def _decode_tokens(self, tokens: list[int]) -> str:
         text = "".join([self._vocab[i] for i in tokens])
         return re.sub(self.DECODE_SPACE_PATTERN, lambda x: " " if x.group(1) else "", text)
 
-    def _recognize_batch(self, waveforms: list[npt.NDArray[np.float32]], language: str | None = None) -> list[str]:
-        return list(map(self._decode_tokens, self._decoding(*self._encode(*self._preprocessor(*pad_list(waveforms))))))
+    def _recognize_batch(
+        self, waveforms: list[npt.NDArray[np.float32]], sample_rate: SampleRates, language: str | None = None
+    ) -> list[str]:
+        return list(
+            map(self._decode_tokens, self._decoding(*self._encode(*self._preprocessor(*pad_list(waveforms), sample_rate))))
+        )
 
 
 class _AsrWithCtcDecoding(_AsrWithDecoding):

src/onnx_asr/loader.py

@@ -146,6 +146,7 @@ def _download_model(model: str, files: list[str]) -> str:
 def load_model(
     model: str | ModelNames | ModelTypes,
     path: str | Path | None = None,
+    *,
     quantization: str | None = None,
     providers: Sequence[str | tuple[str, dict]] | None = None,
 ) -> Asr:

src/onnx_asr/models/whisper.py

@@ -11,7 +11,7 @@
 
 from onnx_asr.asr import Asr
 from onnx_asr.preprocessors.preprocessor import Preprocessor
-from onnx_asr.utils import pad_list
+from onnx_asr.utils import SampleRates, pad_list
 
 
 @typing.no_type_check
@@ -66,8 +66,8 @@ def _get_model_files(quantization: str | None = None) -> dict[str, str]:
             "added_tokens": "added_tokens.json",
         }
 
-    def _preprocess(self, waveforms: list[npt.NDArray[np.float32]]) -> npt.NDArray[np.float32]:
-        input_features, _ = self._preprocessor(*pad_list(waveforms))
+    def _preprocess(self, waveforms: list[npt.NDArray[np.float32]], sample_rate: SampleRates) -> npt.NDArray[np.float32]:
+        input_features, _ = self._preprocessor(*pad_list(waveforms), sample_rate)
         return input_features
 
     @abstractmethod
@@ -77,8 +77,10 @@ def _decode_tokens(self, tokens: npt.NDArray) -> str:
         text = "".join(token for id in tokens if (token := self._vocab[id]) and not token.startswith("<|"))
         return bytearray([self._byte_decoder[c] for c in text]).decode("utf-8", errors="replace").removeprefix(" ")
 
-    def _recognize_batch(self, waveforms: list[npt.NDArray[np.float32]], language: str | None = None) -> list[str]:
-        input_features = self._preprocess(waveforms)
+    def _recognize_batch(
+        self, waveforms: list[npt.NDArray[np.float32]], sample_rate: SampleRates, language: str | None = None
+    ) -> list[str]:
+        input_features = self._preprocess(waveforms, sample_rate)
         input_tokens = np.repeat(self._decoder_input, len(waveforms), axis=0)
 
         if language:
@@ -149,8 +151,8 @@ def _get_model_files(quantization: str | None = None) -> dict[str, str]:
             "decoder": f"**/decoder_model{suffix}.onnx",
         } | _Whisper._get_model_files(suffix)
 
-    def _preprocess(self, waveforms: list[npt.NDArray[np.float32]]) -> npt.NDArray[np.float32]:
-        input_features = super()._preprocess(waveforms)
+    def _preprocess(self, waveforms: list[npt.NDArray[np.float32]], sample_rate: SampleRates) -> npt.NDArray[np.float32]:
+        input_features = super()._preprocess(waveforms, sample_rate)
         (last_hidden_state,) = self._encoder.run(
             ["last_hidden_state"],
             {"input_features": input_features},

src/onnx_asr/preprocessors/__init__.py

@@ -1,5 +1,6 @@
 """ASR preprocessor implementations."""
 
 from .preprocessor import Preprocessor
+from .resampler import Resampler
 
-__all__ = ["Preprocessor"]
+__all__ = ["Preprocessor", "Resampler"]

src/onnx_asr/preprocessors/preprocessor.py

@@ -8,6 +8,10 @@
 import numpy.typing as npt
 import onnxruntime as rt
 
+from onnx_asr.utils import SampleRates
+
+from .resampler import Resampler
+
 
 class Preprocessor:
     """ASR preprocessor implementation."""
@@ -22,11 +26,15 @@ def __init__(self, name: str, **kwargs: Any):
         """
         filename = str(Path(name).with_suffix(".onnx"))
         self._preprocessor = rt.InferenceSession(files(__package__).joinpath(filename).read_bytes(), **kwargs)
+        self._resampler = Resampler(**kwargs)
 
     def __call__(
-        self, waveforms: npt.NDArray[np.float32], waveforms_lens: npt.NDArray[np.int64]
+        self, waveforms: npt.NDArray[np.float32], waveforms_lens: npt.NDArray[np.int64], sample_rate: SampleRates = 16_000
     ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.int64]]:
         """Convert waveforms to model features."""
+        if sample_rate != 16_000:
+            waveforms, waveforms_lens = self._resampler(waveforms, waveforms_lens, sample_rate)
+
         features, features_lens = self._preprocessor.run(
             ["features", "features_lens"], {"waveforms": waveforms, "waveforms_lens": waveforms_lens}
         )

src/onnx_asr/preprocessors/resampler.py

@@ -0,0 +1,33 @@
+"""Waveform resampler implementations."""
+
+from importlib.resources import files
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import onnxruntime as rt
+
+from onnx_asr.utils import SampleRates
+
+
+class Resampler:
+    """Waveform resampler to 16 kHz implementation."""
+
+    def __init__(self, **kwargs: Any):
+        """Create waveform resampler.
+
+        Args:
+            kwargs: Additional parameters for onnxruntime.InferenceSession.
+
+        """
+        self._preprocessor = rt.InferenceSession(files(__package__).joinpath("resample.onnx").read_bytes(), **kwargs)
+
+    def __call__(
+        self, waveforms: npt.NDArray[np.float32], waveforms_lens: npt.NDArray[np.int64], sample_rate: SampleRates
+    ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.int64]]:
+        """Resample waveform to 16 kHz."""
+        resampled, resampled_lens = self._preprocessor.run(
+            ["resampled", "resampled_lens"],
+            {"waveforms": waveforms, "waveforms_lens": waveforms_lens, "sample_rate": [sample_rate]},
+        )
+        return resampled, resampled_lens

src/onnx_asr/utils.py

@@ -1,10 +1,18 @@
 """Utils for ASR."""
 
 import wave
+from typing import Literal, TypeGuard, get_args
 
 import numpy as np
 import numpy.typing as npt
 
+SampleRates = Literal[8_000, 16_000, 22_050, 44_100, 48_000]
+
+
+def is_supported_sample_rate(sample_rate: int) -> TypeGuard[SampleRates]:
+    """Sample rate is supported."""
+    return sample_rate in get_args(SampleRates)
+
 
 class SupportedOnlyMonoAudioError(ValueError):
     """Supported only mono audio error."""
@@ -19,7 +27,15 @@ class WrongSampleRateError(ValueError):
 
     def __init__(self) -> None:
         """Create error."""
-        super().__init__("Supported only 16 kHz sample rate.")
+        super().__init__("Supported only 8, 16, 22.05, 44.1 and 48 kHz sample rate.")
+
+
+class DifferentSampleRatesError(ValueError):
+    """Different sample rates error."""
+
+    def __init__(self) -> None:
+        """Create error."""
+        super().__init__("All sample rates in a batch must be the same.")
 
 
 def read_wav(filename: str) -> tuple[npt.NDArray[np.float32], int]:
@@ -41,23 +57,31 @@ def read_wav(filename: str) -> tuple[npt.NDArray[np.float32], int]:
         return buffer.reshape(f.getnframes(), f.getnchannels()).astype(np.float32) / max_value - zero_value, f.getframerate()
 
 
-def read_wav_files(waveforms: list[npt.NDArray[np.float32] | str]) -> list[npt.NDArray[np.float32]]:
+def read_wav_files(
+    waveforms: list[npt.NDArray[np.float32] | str], numpy_sample_rate: SampleRates
+) -> tuple[list[npt.NDArray[np.float32]], SampleRates]:
     """Convert list of waveform or filenames to list of waveforms."""
     results = []
+    sample_rates = []
     for x in waveforms:
         if isinstance(x, str):
             waveform, sample_rate = read_wav(x)
-            if sample_rate != 16_000:
-                raise WrongSampleRateError()
             if waveform.shape[1] != 1:
                 raise SupportedOnlyMonoAudioError()
             results.append(waveform[:, 0])
+            sample_rates.append(sample_rate)
         else:
             if x.ndim != 1:
                 raise SupportedOnlyMonoAudioError()
             results.append(x)
+            sample_rates.append(numpy_sample_rate)
+
+    if len(set(sample_rates)) > 1:
+        raise DifferentSampleRatesError()
 
-    return results
+    if is_supported_sample_rate(sample_rates[0]):
+        return results, sample_rates[0]
+    raise WrongSampleRateError()
 
 
 def pad_list(arrays: list[npt.NDArray[np.float32]], axis: int = 0) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.int64]]: