Upload speech.py with huggingface_hub

ebcf60e verified 6 days ago

9.78 kB

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

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