preprocess.py

import librosa
import numpy as np
import torch

try:
    import parselmouth
    from parselmouth.praat import call
except ModuleNotFoundError:  # pragma: no cover - optional at runtime
    parselmouth = None
    call = None


TARGET_SAMPLE_RATE = 16000
CLIP_DURATION_SECONDS = 4.0
CLIP_NUM_SAMPLES = int(TARGET_SAMPLE_RATE * CLIP_DURATION_SECONDS)
FEATURE_DIM = 418
NUM_MFCC = 40
NUM_MELS = 128


def load_and_standardize_audio(
    audio_path: str,
    sample_rate: int = TARGET_SAMPLE_RATE,
    required_num_samples: int = CLIP_NUM_SAMPLES,
) -> np.ndarray:
    waveform, _ = librosa.load(audio_path, sr=sample_rate, mono=True)
    if len(waveform) < required_num_samples:
        waveform = np.pad(waveform, (0, required_num_samples - len(waveform)), mode="constant")
    else:
        waveform = waveform[:required_num_samples]

    peak = float(np.max(np.abs(waveform))) if len(waveform) else 0.0
    if peak > 0:
        waveform = waveform / peak
    return waveform.astype(np.float32)


def load_and_standardise(
    audio_path: str,
    sr: int = TARGET_SAMPLE_RATE,
    target_len: int = CLIP_NUM_SAMPLES,
) -> np.ndarray:
    return load_and_standardize_audio(
        audio_path=audio_path,
        sample_rate=sr,
        required_num_samples=target_len,
    )


def augment_waveform(
    waveform: np.ndarray,
    sample_rate: int = TARGET_SAMPLE_RATE,
    augment: bool = True,
) -> np.ndarray:
    del sample_rate  # kept for notebook-signature parity
    if not augment:
        return waveform

    gain = 10 ** (np.random.uniform(-0.3, 0.3) / 20)
    augmented = waveform * gain

    snr_db = np.random.uniform(25, 40)
    signal_power = np.mean(augmented ** 2) + 1e-10
    noise_power = signal_power / (10 ** (snr_db / 10))
    augmented = augmented + np.random.randn(len(augmented)).astype(np.float32) * np.sqrt(noise_power)

    peak = np.max(np.abs(augmented))
    if peak > 0:
        augmented = augmented / peak
    return augmented.astype(np.float32)


def extract_features_418_from_waveform(
    waveform: np.ndarray,
    sample_rate: int = TARGET_SAMPLE_RATE,
    augment: bool = False,
) -> np.ndarray:
    """
    Notebook-faithful 418-dim acoustic feature vector.

    The F0 group uses delta-F0 and z-score normalized F0 instead of raw Hz values
    to reduce pitch leakage into the tonal-language deployment domain.
    """
    if augment:
        waveform = augment_waveform(waveform, sample_rate=sample_rate, augment=True)

    features: list[float] = []
    n_fft = min(2048, len(waveform))
    hop_length = int(sample_rate * 0.010)

    mfcc = librosa.feature.mfcc(
        y=waveform,
        sr=sample_rate,
        n_mfcc=NUM_MFCC,
        n_fft=int(sample_rate * 0.025),
        hop_length=hop_length,
    )
    mfcc_delta = librosa.feature.delta(mfcc)
    mfcc_delta2 = librosa.feature.delta(mfcc, order=2)
    for matrix in (mfcc, mfcc_delta, mfcc_delta2):
        features.extend(np.mean(matrix, axis=1).tolist())
        features.extend(np.std(matrix, axis=1).tolist())

    try:
        sound = parselmouth.Sound(waveform, sampling_frequency=float(sample_rate))
        pitch = sound.to_pitch()
        f0 = pitch.selected_array["frequency"]
        voiced_f0 = f0[f0 > 0]
        if len(voiced_f0) > 5:
            f0_zscore = (voiced_f0 - np.mean(voiced_f0)) / (np.std(voiced_f0) + 1e-8)
            delta_f0 = np.diff(voiced_f0)
            pitch_features = [
                float(np.mean(f0_zscore)),
                float(np.std(f0_zscore)),
                float(np.mean(np.abs(delta_f0))),
                float(np.std(delta_f0)),
                float(np.sum(delta_f0 > 0)) / max(len(delta_f0), 1),
                float(np.polyfit(np.linspace(0, 1, len(voiced_f0)), f0_zscore, 1)[0]),
            ]
        else:
            pitch_features = [0.0] * 6
    except Exception:
        pitch_features = [0.0] * 6
    features.extend(pitch_features)

    rms = librosa.feature.rms(y=waveform)[0]
    features.extend([float(np.mean(rms)), float(np.std(rms))])

    zero_crossing_rate = librosa.feature.zero_crossing_rate(y=waveform)[0]
    features.extend([float(np.mean(zero_crossing_rate)), float(np.std(zero_crossing_rate))])

    spectral_centroid = librosa.feature.spectral_centroid(
        y=waveform,
        sr=sample_rate,
        n_fft=n_fft,
        hop_length=hop_length,
    )[0]
    spectral_bandwidth = librosa.feature.spectral_bandwidth(
        y=waveform,
        sr=sample_rate,
        n_fft=n_fft,
        hop_length=hop_length,
    )[0]
    spectral_rolloff = librosa.feature.spectral_rolloff(
        y=waveform,
        sr=sample_rate,
        n_fft=n_fft,
        hop_length=hop_length,
        roll_percent=0.85,
    )[0]
    onset_strength = librosa.onset.onset_strength(
        y=waveform,
        sr=sample_rate,
        hop_length=hop_length,
    )
    spectral_flatness = librosa.feature.spectral_flatness(
        y=waveform,
        n_fft=n_fft,
        hop_length=hop_length,
    )[0]
    for descriptor in (
        spectral_centroid,
        spectral_bandwidth,
        spectral_rolloff,
        onset_strength,
        spectral_flatness,
    ):
        features.extend([float(np.mean(descriptor)), float(np.std(descriptor))])

    mel_spectrogram = librosa.feature.melspectrogram(
        y=waveform,
        sr=sample_rate,
        n_mels=NUM_MELS,
        n_fft=n_fft,
        hop_length=hop_length,
    )
    features.extend(np.mean(mel_spectrogram, axis=1).tolist())

    try:
        sound = parselmouth.Sound(waveform, sampling_frequency=float(sample_rate))
        harmonicity = call(sound, "To Harmonicity (cc)", 0.01, 75, 0.1, 1.0)
        hnr = call(harmonicity, "Get mean", 0, 0)
        point_process = call(sound, "To PointProcess (periodic, cc)", 75, 500)
        voice_quality_features = [
            call(point_process, "Get jitter (local)", 0, 0, 0.0001, 0.02, 1.3),
            call(point_process, "Get jitter (local, absolute)", 0, 0, 0.0001, 0.02, 1.3),
            call(point_process, "Get jitter (rap)", 0, 0, 0.0001, 0.02, 1.3),
            call([sound, point_process], "Get shimmer (local)", 0, 0, 0.0001, 0.02, 1.3, 1.6),
            call([sound, point_process], "Get shimmer (apq3)", 0, 0, 0.0001, 0.02, 1.3, 1.6),
        ]
        voice_quality_features = [hnr, *voice_quality_features]
        voice_quality_features = [
            0.0 if (value is None or not np.isfinite(value)) else float(value)
            for value in voice_quality_features
        ]
    except Exception:
        voice_quality_features = [0.0] * 6
    features.extend(voice_quality_features)

    chroma_cens = librosa.feature.chroma_cens(y=waveform, sr=sample_rate, n_chroma=12)
    features.extend(np.mean(chroma_cens, axis=1).tolist())
    features.extend(np.std(chroma_cens, axis=1).tolist())

    vector = np.asarray(features, dtype=np.float32)
    if vector.shape[0] != FEATURE_DIM:
        raise ValueError(f"Feature dim mismatch: expected {FEATURE_DIM}, got {vector.shape[0]}.")
    return np.nan_to_num(vector, nan=0.0, posinf=0.0, neginf=0.0)


def extract_features_418(audio_path: str, augment: bool = False) -> np.ndarray:
    waveform = load_and_standardize_audio(audio_path)
    return extract_features_418_from_waveform(
        waveform=waveform,
        sample_rate=TARGET_SAMPLE_RATE,
        augment=augment,
    )


def preprocess_audio_file_for_mada(
    audio_path: str,
    scaler=None,
) -> torch.Tensor:
    features = extract_features_418(audio_path).reshape(1, -1)
    if scaler is not None:
        features = scaler.transform(features)
    return torch.tensor(features, dtype=torch.float32)