| | """Preprocessing and normalization to prepare audio for Kintsugi Depression and Anxiety model.""" |
| | from typing import Union, BinaryIO |
| | import numpy as np |
| | import os |
| | import torch |
| | import torchaudio |
| | from transformers import AutoFeatureExtractor |
| |
|
| | from config import EXPECTED_SAMPLE_RATE, logmel_energies |
| |
|
| |
|
| | def load_audio(source: Union[BinaryIO, str, os.PathLike]) -> torch.Tensor: |
| | """Load audio file, verify mono channel count, and resample if necessary. |
| | |
| | Parameters |
| | ---------- |
| | source: open file or path to file |
| | |
| | Returns |
| | ------- |
| | Time domain audio samples as a 1 x num_samples float tensor sampled at 16 kHz. |
| | |
| | """ |
| | audio, fs = torchaudio.load(source) |
| | if audio.shape[0] != 1: |
| | raise ValueError(f"Provided audio has {audio.shape[0]} != 1 channels.") |
| | if fs != EXPECTED_SAMPLE_RATE: |
| | audio = torchaudio.functional.resample(audio, fs, EXPECTED_SAMPLE_RATE) |
| | return audio |
| |
|
| |
|
| | class Preprocessor: |
| | def __init__(self, |
| | normalize_features: bool = True, |
| | chunk_seconds: int = 30, |
| | max_overlap_frac: float = 0.0, |
| | pad_last_chunk_to_full: bool = True, |
| | ): |
| | """Create preprocessor object. |
| | |
| | Parameters |
| | ---------- |
| | normalize_features: Whether the Whisper preprocessor should normalize features |
| | chunk_seconds: Size of model's receptive field in seconds |
| | max_overlap_frac: Fraction of each chunk allowed to overlap previous chunk for inputs longer than chunk_seconds |
| | pad_last_chunk_to_full: Whether to pad audio to an integer multiple of chunk_seconds |
| | |
| | """ |
| | self.preprocessor = AutoFeatureExtractor.from_pretrained("openai/whisper-small.en") |
| | self.normalize_features = normalize_features |
| | self.chunk_seconds = chunk_seconds |
| | self.max_overlap_frac = max_overlap_frac |
| | self.pad_last_chunk_to_full = pad_last_chunk_to_full |
| |
|
| | def preprocess_with_audio_normalization( |
| | self, |
| | audio: torch.Tensor, |
| | ) -> torch.Tensor: |
| | """Run Whisper preprocessor and normalization expected by the model. |
| | |
| | Note: some normalization steps can be avoided, but are included to match |
| | feature extraction used during training. |
| | |
| | Parameters |
| | ---------- |
| | audio: Raw audio samples as a 1 x num_samples float tensor sampled at 16 kHz |
| | |
| | Returns |
| | ------- |
| | Normalized mel filter bank features as a float tensor of shape |
| | num_chunks x 80 mel filter bands x 3000 time frames |
| | |
| | """ |
| | |
| | audio = torch.squeeze(audio, 0) |
| | audio = audio - torch.mean(audio) |
| | audio = audio / torch.max(torch.abs(audio)) |
| |
|
| | chunk_samples = EXPECTED_SAMPLE_RATE * self.chunk_seconds |
| |
|
| | if self.pad_last_chunk_to_full: |
| | |
| | if self.max_overlap_frac > 0: |
| | raise ValueError( |
| | f"pad_last_chunk_to_full is only supported for non-overlapping windows" |
| | ) |
| | num_chunks = np.ceil(len(audio) / chunk_samples) |
| | pad_size = int(num_chunks * chunk_samples - len(audio)) |
| | audio = torch.nn.functional.pad(audio, (0, pad_size)) |
| |
|
| | overflow_len = len(audio) - chunk_samples |
| |
|
| | min_hop_samples = int( |
| | (1 - self.max_overlap_frac) * chunk_samples |
| | ) |
| |
|
| | n_windows = 1 + overflow_len // min_hop_samples |
| | window_starts = np.linspace(0, overflow_len, max(n_windows, 1)).astype(int) |
| |
|
| | features = self.preprocessor( |
| | [ |
| | audio[start : start + chunk_samples].numpy(force=True) |
| | for start in window_starts |
| | ], |
| | return_tensors="pt", |
| | sampling_rate=EXPECTED_SAMPLE_RATE, |
| | do_normalize=self.normalize_features, |
| | ) |
| | for key in ("input_features", "input_values"): |
| | if hasattr(features, key): |
| | features = getattr(features, key) |
| | break |
| |
|
| | mean_features = torch.mean(features, dim=-1) |
| | |
| | rescale_factor = logmel_energies.unsqueeze(0) - mean_features |
| | rescale_factor = rescale_factor.unsqueeze(2) |
| | features += rescale_factor |
| | return features |
| |
|
