Blame for speech.py

"""Extract prosody features from recorded audio.

Speak instead of drawing. The same two signals come out -- interpolated f0
and intensity at 200 Hz -- so the rest of the chain is unchanged: conformer
to mel, mel to words.

This is the other half of the demo's point. A drawn contour shows that
prosody alone carries recoverable words; real speech put through the
identical pipeline shows what the drawing is being compared against, and how
far a contour can be from a human one and still transcribe.

Ported from the reference implementation's `speech.py`, which in turn ports
the training preprocessor. The settings are not interchangeable with
sensible-looking alternatives: frame_length, nccf_thresh1, the zero-padding,
the six-frame trim and the two clamps all move the contour, and the
conformer only ever saw the combination below.
"""

from __future__ import annotations

import numpy as np

from .features import FEATURE_RATE_HZ, ProsodyFeatures

# Whisper's rate, and what the preprocessor ran at.
TARGET_SAMPLE_RATE = 16000

# pYAAPT settings from the training preprocessor's own config. Every one of
# these differed from the obvious default, and the conformer only ever saw
# contours produced with these.
FRAME_LENGTH_MS = 20.0
FRAME_SPACE_MS = 5.0          # 200 Hz
TDA_FRAME_LENGTH_MS = 25.0
NCCF_THRESH1 = 0.25
# f0_min/f0_max are deliberately not passed: the preprocessor leaves yaapt's
# own defaults (60/400 Hz) in place. Passing 50/500 widened the search and
# changed which frames came back voiced.

# The preprocessor drops the first six frames, which the zero-padding below
# makes spurious.
LEADING_FRAMES_DROPPED = 6
# Clamped after interpolation, as the preprocessor does.
F0_CEILING_HZ = 500.0
INTENSITY_CLAMP_DB = 20.0

# Guards on a posted recording. Audio is far bigger than a drawing, and the
# conformer export takes ten seconds regardless.
MAX_AUDIO_S = 10.0

# How far below the loudest frame still counts as speech.
#
# yaapt tracks a pitch in every frame it can and f0_interp fills the rest, so
# "voiced" is no help at all -- room noise comes back 100% voiced. Intensity
# separates them cleanly: measured on synthetic cases, noise sits 25-45 dB
# under the speech peak while the speech itself spans about 10 dB.
#
# Relative to the peak rather than absolute, because microphone gain varies
# by tens of dB between machines and an absolute floor would trim everything
# on a quiet one and nothing on a hot one.
#
# 18 dB, tightened from 25 after measuring what a stricter cut actually
# costs. The worry with a strict threshold is clipping a quiet trailing
# syllable, so it was tested against speech ending 12 dB down: at 18 dB that
# case still comes back within 0.04 s of its true length, while a loud room
# no longer survives. Below about 12 dB the threshold starts competing with
# the dynamic range of the speech itself.
SPEECH_FLOOR_BELOW_PEAK_DB = 18.0

# Frames of margin kept either side, so an onset is not clipped mid-syllable.
# Four, down from eight: at 20 ms a side the margin was contributing more
# excess than the threshold was, which is a strange way to be generous.
TRIM_MARGIN_FRAMES = 4


def resample(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
    """Linear resample, which is what the reference implementation uses."""
    audio = np.asarray(audio, dtype=np.float64)
    if orig_sr == target_sr or audio.size == 0:
        return audio
    duration = audio.size / orig_sr
    target_n = max(1, int(round(duration * target_sr)))
    source_t = np.arange(audio.size) / orig_sr
    target_t = np.arange(target_n) / target_sr
    return np.interp(target_t, source_t, audio)


def extract_features(
    audio: np.ndarray, sr: int = TARGET_SAMPLE_RATE, trim: bool = True
) -> ProsodyFeatures:
    """Audio -> the f0 and intensity the conformer consumes.

    Returns raw f0_hz and intensity_db; call `.normalised()` for the model
    input, exactly as a drawing does.

    `trim` cuts the quiet lead-in and tail, which a live microphone always
    has and the training clips never did -- see `trim_to_speech`. Pass False
    to reproduce the reference implementation's output exactly, which is what
    the equivalence test does.
    """
    import parselmouth
    from amfm_decompy import basic_tools, pYAAPT

    from ._yaapt_patch import patched_interpolate

    audio = np.asarray(audio, dtype=np.float64)
    if audio.ndim > 1:
        audio = audio.squeeze()
    if audio.size == 0:
        return ProsodyFeatures(np.zeros(0), np.zeros(0))

    if sr != TARGET_SAMPLE_RATE:
        audio = resample(audio, sr, TARGET_SAMPLE_RATE)
        sr = TARGET_SAMPLE_RATE

    # The preprocessor patches pYAAPT's interpolation before every call and
    # reads samp_interp off the result rather than interpolating the voiced
    # values afterwards. Patch the class, not a stub, so samp_interp is
    # produced the same way.
    pYAAPT.PitchObj.interpolate = patched_interpolate

    # Zero-pad by half a frame each side before tracking. This shifts every
    # frame boundary, so without it the whole contour sits offset against
    # what training saw.
    pad = int(FRAME_LENGTH_MS / 1000 * sr) // 2
    padded = np.pad(audio, (pad, pad), "constant", constant_values=0)

    signal = basic_tools.SignalObj(padded, sr)
    pitch = pYAAPT.yaapt(
        signal,
        frame_length=FRAME_LENGTH_MS,
        frame_space=FRAME_SPACE_MS,
        nccf_thresh1=NCCF_THRESH1,
        tda_frame_length=TDA_FRAME_LENGTH_MS,
    )
    f0 = np.asarray(pitch.samp_interp, dtype=np.float64)
    f0 = f0[LEADING_FRAMES_DROPPED:]
    f0[f0 > F0_CEILING_HZ] = 0.0
    f0[f0 < 0.0] = 0.0

    # No minimum_pitch: the preprocessor passes only time_step, and
    # parselmouth's default pitch floor (100 Hz) sets the analysis window.
    # Passing 50 Hz instead doubles that window and smooths the envelope.
    sound = parselmouth.Sound(audio, sampling_frequency=sr, start_time=0)
    intensity = sound.to_intensity(time_step=1.0 / FEATURE_RATE_HZ)
    intensity_db = np.asarray(intensity.values.T.flatten(), dtype=np.float64)

    frames = min(len(f0), len(intensity_db))
    f0 = f0[:frames]
    intensity_db = intensity_db[:frames]
    intensity_db[intensity_db < INTENSITY_CLAMP_DB] = INTENSITY_CLAMP_DB

    if trim:
        f0, intensity_db = trim_to_speech(f0, intensity_db)

    return ProsodyFeatures(f0_hz=f0, intensity_db=intensity_db)


def trim_to_speech(
    f0: np.ndarray,
    intensity_db: np.ndarray,
    below_peak_db: float = SPEECH_FLOOR_BELOW_PEAK_DB,
    margin: int = TRIM_MARGIN_FRAMES,
) -> tuple[np.ndarray, np.ndarray]:
    """Cut the quiet lead-in and tail off a recording.

    A live microphone hands over whatever was in the room before and after
    the words, and the model has no way to tell that from speech: yaapt finds
    a pitch in noise, and f0_interp fills whatever it misses, so the contour
    arrives looking fully voiced throughout. Feeding that in means the mel --
    and the transcription -- describe several hundred milliseconds of nothing
    alongside the utterance.

    The reference implementation never needed this because its dataset clips
    came trimmed to the utterance already. A recording does not.

    Returns the arrays unchanged when there is no clear peak to measure
    against, which is the honest answer for a recording that is quiet
    throughout: better to pass it on and let the result show that than to
    invent a span.
    """
    if intensity_db.size == 0:
        return f0, intensity_db

    peak = float(intensity_db.max())
    # Nothing to separate: a recording that never rose above the clamp is
    # silence, and any threshold would pick an arbitrary window out of it.
    if peak <= INTENSITY_CLAMP_DB:
        return f0, intensity_db

    loud = np.flatnonzero(intensity_db > peak - below_peak_db)
    if loud.size == 0:
        return f0, intensity_db

    first = max(0, int(loud[0]) - margin)
    last = min(len(intensity_db) - 1, int(loud[-1]) + margin)
    return f0[first:last + 1], intensity_db[first:last + 1]


def decode_wav(data: bytes) -> tuple[np.ndarray, int]:
    """A WAV blob -> (mono float64 in -1..1, sample rate).

    The browser records to WebM/Opus by default, which needs a decoder this
    package should not carry; the page converts to WAV in an AudioContext
    first, so only the one format arrives here. `wave` is in the standard
    library, which keeps the dependency list where it is.
    """
    import io
    import wave

    try:
        with wave.open(io.BytesIO(data), "rb") as handle:
            channels = handle.getnchannels()
            width = handle.getsampwidth()
            rate = handle.getframerate()
            frames = handle.readframes(handle.getnframes())
    except Exception as error:
        # A malformed body is the caller's fault, so raise the kind of error
        # the route turns into a 400. Letting wave's own exception through
        # produced a 500, which says the server broke.
        raise ValueError(f"not a readable WAV: {error}") from error

    if width == 2:
        samples = np.frombuffer(frames, dtype="<i2").astype(np.float64) / 32768.0
    elif width == 4:
        samples = np.frombuffer(frames, dtype="<f4").astype(np.float64)
    elif width == 1:
        samples = (np.frombuffer(frames, dtype="u1").astype(np.float64) - 128) / 128.0
    else:
        raise ValueError(f"unsupported WAV sample width: {width} bytes")

    if channels > 1:
        samples = samples.reshape(-1, channels).mean(axis=1)

    seconds = samples.size / max(1, rate)
    if seconds > MAX_AUDIO_S:
        raise ValueError(
            f"recording is {seconds:.1f}s; the conformer export takes at "
            f"most {MAX_AUDIO_S:.0f}s"
        )
    return samples, rate