Upload features.py with huggingface_hub

features.py

@@ -0,0 +1,316 @@
+"""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)