python/diarizen_sdk/preprocess.py

"""Audio preprocessing for DiariZen segmentation model."""

import numpy as np


def preprocess_audio(
    audio: np.ndarray,
    sample_rate: int,
    target_sr: int = 16000,
    duration_s: float = 4.0,
    eps: float = 1e-5,
) -> np.ndarray:
    """Resample, trim, and LayerNorm-normalize audio for the CNN NPU frontend.

    Args:
        audio: 1-D float32 waveform.
        sample_rate: Original sample rate.
        target_sr: Target sample rate (default 16000).
        duration_s: Window duration in seconds (default 4.0).
        eps: Epsilon for LayerNorm.

    Returns:
        Normalized waveform of shape (1, target_sr * duration_s), float32.
    """
    target_samples = int(target_sr * duration_s)

    # Simple linear resample
    if sample_rate != target_sr:
        ratio = target_sr / sample_rate
        out_len = int(len(audio) * ratio)
        indices = np.linspace(0, len(audio) - 1, out_len)
        audio = np.interp(indices, np.arange(len(audio)), audio).astype(np.float32)

    # Trim or pad to target length
    if len(audio) < target_samples:
        audio = np.pad(audio, (0, target_samples - len(audio)))
    else:
        audio = audio[:target_samples]

    # LayerNorm normalization
    mean = audio.mean()
    var = ((audio - mean) ** 2).mean()
    audio_norm = (audio - mean) / np.sqrt(var + eps)

    return audio_norm.reshape(1, target_samples).astype(np.float32)