|
|
"""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 |
|
|
|
