prosody2text-code / features.py
Daporte's picture
Upload features.py with huggingface_hub
73a450d verified
Raw History Blame Contribute Delete
13 kB
"""Turn drawn strokes into the prosody features a speech model expects.
The transcription models in `finetune_whisper` were trained on two signals
extracted from real speech at 200 Hz -- interpolated F0 and intensity -- which
is exactly what a stylus produces: vertical position is pitch, pressure is
loudness. This module resamples strokes onto that 200 Hz grid and applies the
same normalisation the training pipeline used, so a drawing can be fed to the
model as if it were speech.
The constants here are not free choices. They mirror the training pipeline's
own FEATURE_SCALES and the checkpoint's `model_config.json`; if those change,
these must change with them. This module is a verbatim port of the reference
implementation's `features.py`, minus code paths the demo never takes, so a
drawing is turned into model input exactly as in the paper.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from .f0_interp import interpolate_f0
# Prosody features are extracted at 200 Hz (time_step 0.005 s) and the mel
# target runs at 100 Hz; the model's length_ratio=0.5 encodes that 2:1.
FEATURE_RATE_HZ = 200.0
# Physical-range normalisation, copied from FeaturesToMelDataset.FEATURE_SCALES.
F0_SCALE = 500.0 # Hz, human pitch range ~50-500
INTENSITY_SCALE = 100.0 # dB SPL, typical range 0-100
# Pen pressure is a 0..1 ratio; speech intensity is dB SPL. These bounds are
# measured, not chosen: across the extracted templates (2600 frames of real
# speech from the training distribution) intensity runs p5=51.8, p50=68.6,
# p95=81.0 dB. Mapping 0..1 onto 50..85 dB covers that spread and puts the
# speech median at pressure 0.53, so a natural mid-pressure stroke produces
# typical speech loudness.
#
# The earlier 45-80 dB guess put the median at 0.67 and clipped the top: a
# comfortable touch read ~6 dB quieter than real speech, and the loudest
# frames (84.5 dB) could not be reached at all.
#
# Re-derive with: uv run python scripts/calibrate_intensity.py
INTENSITY_FLOOR_DB = 50.0
INTENSITY_CEIL_DB = 85.0
# Where speech actually lives. These are measured from the extracted
# templates (2600 frames of real speech from the training distribution)
# rather than taken from textbook adult F0 ranges: p1 is MIDI 38.7 and p99.5
# is MIDI 64.9, so real contours reach well outside the "average adult"
# 85-255 Hz band.
#
# An earlier F2-C#4 (41-61) surface clipped 4.8% of frames, and the clipped
# frames are the pitch extremes that carry the most information. B1-G4 clips
# none of the measured 2600.
#
# Re-derive with: uv run python scripts/trace_fidelity.py (which reports
# clipping) after rebuilding templates.json.
SPEECH_MIDI_LOW = 35 # B1, 61.7 Hz
SPEECH_MIDI_HIGH = 67 # G4, 392.0 Hz
# The widest surface the model can be drawn on at all, as opposed to the
# range real speech occupies. F0_SCALE is 500 Hz, which is MIDI 71 -- above
# it the normalised input exceeds 1.0, a value training never contained.
# Downward there is no such wall: a low pitch simply normalises small.
#
# Drawing out here is out of distribution and will transcribe worse, but it
# is a usable surface rather than a broken one, and a wide sweep is sometimes
# what the drawing wants.
# Extended downward rather than upward. The ceiling is a hard wall -- MIDI 71
# is 494 Hz and F0_SCALE is 500, so there is almost nothing above it -- while
# downward a low pitch simply normalises small. Logged drawings sit at a p99
# of MIDI 58, so headroom above 60 is canvas nobody uses; the room worth
# having is underneath.
DRAWABLE_MIDI_LOW = 12 # C0, 16.4 Hz
DRAWABLE_MIDI_HIGH = 71 # B4, 493.9 Hz
def midi_to_hz(midi: float | np.ndarray) -> np.ndarray:
return 440.0 * 2.0 ** ((np.asarray(midi, dtype=np.float64) - 69.0) / 12.0)
@dataclass
class ProsodyFeatures:
"""f0 and intensity on the model's 200 Hz grid.
`raw` holds physical units (Hz, dB) for inspection; `normalised` is what
the model consumes.
"""
f0_hz: np.ndarray
intensity_db: np.ndarray
@property
def duration_s(self) -> float:
return len(self.f0_hz) / FEATURE_RATE_HZ
@property
def num_frames(self) -> int:
return len(self.f0_hz)
def normalised(self) -> np.ndarray:
"""[T, 2] float32 array scaled exactly as in training."""
f0 = np.asarray(self.f0_hz, dtype=np.float32) / F0_SCALE
intensity = np.asarray(self.intensity_db, dtype=np.float32) / INTENSITY_SCALE
return np.stack([f0, intensity], axis=-1).astype(np.float32)
def pressure_to_db(
pressure: np.ndarray,
floor_db: float = INTENSITY_FLOOR_DB,
ceil_db: float = INTENSITY_CEIL_DB,
) -> np.ndarray:
"""Map 0..1 stylus pressure onto a speech-like dB SPL range."""
p = np.clip(np.asarray(pressure, dtype=np.float64), 0.0, 1.0)
return floor_db + (ceil_db - floor_db) * p
def resample_stroke(
times_s: np.ndarray,
values: np.ndarray,
rate_hz: float = FEATURE_RATE_HZ,
) -> tuple[np.ndarray, np.ndarray]:
"""Resample irregular pen samples onto a uniform grid.
Tablet events arrive at an irregular rate, so the raw samples cannot be
used directly: the model assumes a fixed 200 Hz step. Returns the grid
times and the linearly interpolated values.
"""
times_s = np.asarray(times_s, dtype=np.float64)
values = np.asarray(values, dtype=np.float64)
if times_s.size == 0:
return np.zeros(0), np.zeros(0)
if times_s.size == 1:
return times_s.copy(), values.copy()
# Pen events are timestamped by arrival and can repeat or (rarely) go
# backwards; np.interp needs a strictly increasing x.
order = np.argsort(times_s, kind="stable")
times_s, values = times_s[order], values[order]
keep = np.concatenate(([True], np.diff(times_s) > 0))
times_s, values = times_s[keep], values[keep]
if times_s.size == 1:
return times_s.copy(), values.copy()
duration = times_s[-1] - times_s[0]
n = max(2, int(round(duration * rate_hz)) + 1)
grid = times_s[0] + np.arange(n) / rate_hz
# Clamp so the last grid point never extrapolates past the final sample.
grid = np.minimum(grid, times_s[-1])
return grid, np.interp(grid, times_s, values)
def smooth(values: np.ndarray, window_frames: int) -> np.ndarray:
"""Moving-average smoothing with edge padding.
Drawn lines are geometrically smooth but can still have corners that no
speech F0 contour would contain. A short smoothing window keeps the
contour closer to what the model saw in training.
"""
if window_frames <= 1 or values.size == 0:
return np.asarray(values, dtype=np.float64)
window_frames = min(window_frames, values.size)
kernel = np.ones(window_frames) / window_frames
pad = window_frames // 2
padded = np.pad(np.asarray(values, dtype=np.float64), pad, mode="edge")
out = np.convolve(padded, kernel, mode="same")
return out[pad : pad + values.size]
def features_from_stroke(
times_s: np.ndarray,
midi: np.ndarray,
pressure: np.ndarray,
smooth_ms: float = 25.0,
rate_hz: float = FEATURE_RATE_HZ,
) -> ProsodyFeatures:
"""Build 200 Hz prosody features from one pen stroke.
`times_s` are event timestamps in seconds, `midi` the pitch under the pen
(fractional MIDI note), `pressure` the stylus pressure in 0..1.
Pitch is read as a literal note and converted to Hz, so the drawing
surface must already sit in the range speech occupies -- see
SPEECH_MIDI_LOW/HIGH and the roll bounds the app is built with.
"""
times_s = np.asarray(times_s, dtype=np.float64)
midi = np.asarray(midi, dtype=np.float64)
pressure = np.asarray(pressure, dtype=np.float64)
if not (times_s.size == midi.size == pressure.size):
raise ValueError("times, midi and pressure must be the same length")
if times_s.size == 0:
return ProsodyFeatures(np.zeros(0, np.float64), np.zeros(0, np.float64))
grid, midi_grid = resample_stroke(times_s, midi, rate_hz)
_, pressure_grid = resample_stroke(times_s, pressure, rate_hz)
# resample_stroke derives its grid from the data, so both calls agree in
# length; guard anyway rather than silently mis-aligning the channels.
n = min(len(midi_grid), len(pressure_grid))
midi_grid, pressure_grid = midi_grid[:n], pressure_grid[:n]
f0_raw = midi_to_hz(midi_grid)
window = max(1, int(round(smooth_ms / 1000.0 * rate_hz)))
f0_hz = smooth(f0_raw, window)
intensity_db = smooth(pressure_to_db(pressure_grid), window)
return ProsodyFeatures(f0_hz=f0_hz, intensity_db=intensity_db)
# Quiet frames around a phrase, so it neither starts nor stops dead.
#
# Measured, not guessed: across the templates the quiet before and after the
# speech has a median of 0 ms and a maximum of 80 ms -- these dataset clips
# are trimmed tight to the utterance, so training really does start at
# speech. An earlier 250 ms on each side was 3x the observed maximum and
# contributed more frames than the drawing itself on a short phrase.
#
# 40 ms keeps the onset from being a step discontinuity without displacing
# content; it sits inside the observed range rather than beyond it.
LEAD_SILENCE_S = 0.04
TAIL_SILENCE_S = 0.04
def pad_with_silence(
features: ProsodyFeatures,
lead_s: float = LEAD_SILENCE_S,
tail_s: float = TAIL_SILENCE_S,
rate_hz: float = FEATURE_RATE_HZ,
) -> ProsodyFeatures:
"""Surround a phrase with quiet frames so it neither starts nor stops dead.
Pitch is held at the adjacent value rather than zeroed, matching how
f0_interp behaves through unvoiced regions, and intensity sits at the
floor.
"""
if features.num_frames == 0:
return features
lead = max(0, int(round(lead_s * rate_hz)))
tail = max(0, int(round(tail_s * rate_hz)))
if not lead and not tail:
return features
f0_parts, db_parts = [], []
if lead:
f0_parts.append(np.full(lead, features.f0_hz[0]))
db_parts.append(np.full(lead, INTENSITY_FLOOR_DB))
f0_parts.append(features.f0_hz)
db_parts.append(features.intensity_db)
if tail:
f0_parts.append(np.full(tail, features.f0_hz[-1]))
db_parts.append(np.full(tail, INTENSITY_FLOOR_DB))
return ProsodyFeatures(
f0_hz=np.concatenate(f0_parts),
intensity_db=np.concatenate(db_parts),
)
def concatenate(
features: list[ProsodyFeatures],
gap_s: float | list[float] = 0.12,
rate_hz: float = FEATURE_RATE_HZ,
lead_s: float = LEAD_SILENCE_S,
tail_s: float = TAIL_SILENCE_S,
) -> ProsodyFeatures:
"""Join strokes into one utterance, with a quiet gap between them.
Each stroke is a "note"; the gaps are what give the model something it can
read as separate syllables rather than one continuous tone.
`gap_s` may be a single duration or one per gap. Per-gap is what the
canvas passes: the pen's actual lift time is recorded, so there is no
reason to substitute a constant for it. A fixed gap also compounds -- at
120 ms across the 3-8 strokes a phrase typically has, the inserted silence
can outweigh the drawing itself.
Across a gap the pitch is **interpolated by pYAAPT's own routine**, not
held and not linearly ramped. The conformer was trained on f0_interp,
which median filters and fits a shape-preserving PCHIP spline through the
unvoiced regions; holding the previous value gives a flat shelf then a
step, and a straight line gives a shape speech never has. Gap frames are
marked unvoiced (f0 = 0) and handed to interpolate_f0, so they are filled
the same way training filled them. Intensity drops to the floor, which is
what makes the gap read as a pause.
"""
usable = [f for f in features if f.num_frames > 0]
if not usable:
return ProsodyFeatures(np.zeros(0, np.float64), np.zeros(0, np.float64))
if isinstance(gap_s, (int, float)):
gaps = [float(gap_s)] * max(0, len(usable) - 1)
else:
gaps = [float(g) for g in gap_s]
if len(gaps) < len(usable) - 1:
gaps += [0.12] * (len(usable) - 1 - len(gaps))
f0_parts: list[np.ndarray] = []
db_parts: list[np.ndarray] = []
for i, feat in enumerate(usable):
gap_frames = max(0, int(round(gaps[i - 1] * rate_hz))) if i > 0 else 0
if i > 0 and gap_frames:
# Zero marks the gap as unvoiced; interpolate_f0 fills it below,
# exactly as yaapt filled unvoiced regions in training.
f0_parts.append(np.zeros(gap_frames))
db_parts.append(np.full(gap_frames, INTENSITY_FLOOR_DB))
f0_parts.append(np.asarray(feat.f0_hz, dtype=np.float64))
db_parts.append(np.asarray(feat.intensity_db, dtype=np.float64))
joined = ProsodyFeatures(
f0_hz=interpolate_f0(np.concatenate(f0_parts)),
intensity_db=np.concatenate(db_parts),
)
return pad_with_silence(joined, lead_s, tail_s, rate_hz)