From 1239e5e0ac0b6301e54e2db285f3bb1103a3fe85 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Fri, 30 Oct 2026 03:22:14 +0100 Subject: [PATCH] feat: basic pitch transcription Co-Authored-By: Claude Sonnet 5.5 --- src/pipeline/transcribe.ts | 219 ++++++++++++++++++++++++++++++ tests/pipeline/transcribe.test.ts | 41 ++++++ 2 files changed, 260 insertions(+) create mode 100644 src/pipeline/transcribe.ts create mode 100644 tests/pipeline/transcribe.test.ts diff --git a/src/pipeline/transcribe.ts b/src/pipeline/transcribe.ts new file mode 100644 index 0000000..eeeb25b --- /dev/null +++ b/src/pipeline/transcribe.ts @@ -0,0 +1,219 @@ +// Basic Pitch transcription via ONNX Runtime. Windowing mirrors @spotify/basic-pitch inference.js and the note +// post-processing (outputToNotesPoly, noteFramesToTime) is vendored from its toMidi.js, because that module imports +// @tonejs/midi (MIDI file writing we do not need). Original code: Copyright 2022 Spotify AB, Apache License 2.0 +// (https://www.apache.org/licenses/LICENSE-2.0), ported here to TypeScript with pitch-range arguments adapted. +// Pure: no fs / Bun APIs; only ORT APIs that also exist in onnxruntime-web (InferenceSession.create, Tensor, run). +import * as ort from "onnxruntime-node"; +import type { Audio, TimedNote } from "../core/types.ts"; +import { toMono, resample } from "../audio/wav.ts"; + +export interface TranscribeOpts { + onsetThresh?: number; // 0.5 + frameThresh?: number; // 0.3 + minNoteMs?: number; // 58 + minPitch?: number; // MIDI, inclusive + maxPitch?: number; // MIDI, inclusive + monophonic?: boolean; +} + +const SR = 22050; +const FFT_HOP = 256; +const FPS = Math.floor(SR / FFT_HOP); // 86 +const N_SAMPLES = SR * 2 - FFT_HOP; // 43844 +const OVERLAP_FRAMES = 30; +const OVERLAP_OVER_2 = OVERLAP_FRAMES / 2; +const OVERLAP_LEN = OVERLAP_FRAMES * FFT_HOP; // 7680 +const HOP_SIZE = N_SAMPLES - OVERLAP_LEN; +const MIDI_OFFSET = 21; +const MAX_FREQ_IDX = 87; +const ANNOT_N_FRAMES = FPS * 2; +const WINDOW_OFFSET = (FFT_HOP / SR) * (ANNOT_N_FRAMES - N_SAMPLES / FFT_HOP) + 0.0018; +const BATCH = 8; +const INPUT = "serving_default_input_2:0"; +// Same order as inference.js: Identity (contours), Identity_1 (frames / note), Identity_2 (onsets). +const OUT_NOTE = "StatefulPartitionedCall:1"; +const OUT_ONSET = "StatefulPartitionedCall:2"; + +interface NoteEvent { startFrame: number; durationFrames: number; pitchMidi: number; amplitude: number } + +const modelFrameToTime = (frame: number): number => + (frame * FFT_HOP) / SR - WINDOW_OFFSET * Math.floor(frame / ANNOT_N_FRAMES); + +function globalMax(a: number[][]): number { + let m = 0; + for (const row of a) for (const v of row) if (v > m) m = v; + return m; +} + +function getInferredOnsets(onsets: number[][], frames: number[][], nDiff = 2): number[][] { + const nF = frames.length, nB = frames[0]!.length; + const frameDiff: number[][] = Array.from({ length: nF }, () => new Array(nB).fill(0)); + for (let r = nDiff; r < nF; r++) { + for (let c = 0; c < nB; c++) { + let mn = Infinity; + for (let n = 1; n <= nDiff; n++) mn = Math.min(mn, frames[r]![c]! - (r - n >= 0 ? frames[r - n]![c]! : 0)); + frameDiff[r]![c] = Math.max(mn, 0); + } + } + const onsetMax = globalMax(onsets), diffMax = globalMax(frameDiff); + return onsets.map((row, r) => row.map((v, c) => Math.max(v, (onsetMax * frameDiff[r]![c]!) / diffMax))); +} + +function argRelMax(a: number[][], order = 1): [number, number][] { + const out: [number, number][] = []; + for (let col = 0; col < a[0]!.length; col++) { + for (let row = 0; row < a.length; row++) { + let ok = true; + for (let k = Math.max(0, row - order); ok && k <= Math.min(a.length - 1, row + order); k++) { + if (k !== row) ok = a[row]![col]! > a[k]![col]!; + } + if (ok) out.push([row, col]); + } + } + return out; +} + +/** Vendored from basic-pitch toMidi.js outputToNotesPoly (melodia trick on, onsets inferred). */ +function outputToNotesPoly( + frames: number[][], onsets: number[][], onsetThresh: number, frameThresh: number, minNoteLen: number, + minIdx: number, maxIdx: number, energyTolerance = 11, +): NoteEvent[] { + const nFrames = frames.length; + for (const m of [onsets, frames]) for (const row of m) { row.fill(0, maxIdx + 1); row.fill(0, 0, minIdx); } + const inferred = getInferredOnsets(onsets, frames); + const peaks = inferred.map((o) => o.map(() => 0)); + for (const [r, c] of argRelMax(inferred)) peaks[r]![c] = inferred[r]![c]!; + const starts: number[] = [], freqs: number[] = []; + for (let i = 0; i < peaks.length; i++) for (let j = 0; j < peaks[i]!.length; j++) if (peaks[i]![j]! > onsetThresh) { starts.push(i); freqs.push(j); } + starts.reverse(); freqs.reverse(); + const remaining = frames.map((f) => f.slice()); + const events: NoteEvent[] = []; + const clearNeighbours = (i: number, f: number): void => { + remaining[i]![f] = 0; + if (f < MAX_FREQ_IDX) remaining[i]![f + 1] = 0; + if (f > 0) remaining[i]![f - 1] = 0; + }; + const meanAmp = (s: number, e: number, f: number): number => { + let sum = 0; + for (let i = s; i < e; i++) sum += frames[i]![f]!; + return sum / (e - s); + }; + starts.forEach((s, idx) => { + const f = freqs[idx]!; + if (s >= nFrames - 1) return; + let i = s + 1, k = 0; + while (i < nFrames - 1 && k < energyTolerance) { + k = remaining[i]![f]! < frameThresh ? k + 1 : 0; + i++; + } + i -= k; + if (i - s <= minNoteLen) return; + for (let j = s; j < i; j++) clearNeighbours(j, f); + events.push({ startFrame: s, durationFrames: i - s, pitchMidi: f + MIDI_OFFSET, amplitude: meanAmp(s, i, f) }); + }); + while (globalMax(remaining) > frameThresh) { + let iMid = 0, f = 0, best = remaining[0]![0]!; + for (let r = 0; r < nFrames; r++) { + const row = remaining[r]!; + for (let c = 0; c < row.length; c++) if (row[c]! > best) { best = row[c]!; iMid = r; f = c; } + } + remaining[iMid]![f] = 0; + let i = iMid + 1, k = 0; + while (i < nFrames - 1 && k < energyTolerance) { + k = remaining[i]![f]! < frameThresh ? k + 1 : 0; + clearNeighbours(i, f); + i++; + } + const iEnd = i - 1 - k; + i = iMid - 1; k = 0; + while (i > 0 && k < energyTolerance) { + k = remaining[i]![f]! < frameThresh ? k + 1 : 0; + clearNeighbours(i, f); + i--; + } + const iStart = i + 1 + k; + if (iEnd - iStart <= minNoteLen) continue; + events.push({ startFrame: iStart, durationFrames: iEnd - iStart, pitchMidi: f + MIDI_OFFSET, amplitude: meanAmp(iStart, iEnd, f) }); + } + return events; +} + +type Piece = TimedNote & { amp: number }; + +/** Keep, at each instant, only the loudest active note; later/quieter overlaps are trimmed away. */ +function makeMonophonic(notes: Piece[], minDur: number): TimedNote[] { + const bounds = [...new Set(notes.flatMap((n) => [n.onset, n.offset]))].sort((a, b) => a - b); + const out: TimedNote[] = []; + let cur: { n: Piece; start: number; end: number } | null = null; + const flush = (): void => { + if (cur && cur.end - cur.start >= minDur) { + out.push({ onset: cur.start, offset: cur.end, pitch: cur.n.pitch, velocity: cur.n.velocity }); + } + cur = null; + }; + for (let b = 0; b + 1 < bounds.length; b++) { + const t0 = bounds[b]!, t1 = bounds[b + 1]!; + let top: Piece | null = null; + for (const n of notes) if (n.onset <= t0 && n.offset >= t1 && (!top || n.amp > top.amp)) top = n; + if (cur && top === cur.n && cur.end === t0) { cur.end = t1; continue; } + flush(); + // continuation of a note that was interrupted by a louder one would be a spurious onset: only start at a note's own onset + if (top && top.onset === t0) cur = { n: top, start: t0, end: t1 }; + } + flush(); + return out; +} + +export async function createTranscriber( + modelBytes: Uint8Array, +): Promise<(audio: Audio, opts?: TranscribeOpts) => Promise> { + const session = await ort.InferenceSession.create(modelBytes, { executionProviders: ["cpu"] }); + + return async (audio, opts = {}) => { + const mono = resample(toMono(audio), audio.sampleRate, SR); + let ss = 0; + for (const v of mono) ss += v * v; + if (mono.length === 0 || Math.sqrt(ss / mono.length) < 1e-4) return []; + + // pad OVERLAP/2 zeros at the front, then frame with pad_end (tf.signal.frame(..., padEnd=true)) + const padded = new Float32Array(OVERLAP_LEN / 2 + mono.length); + padded.set(mono, OVERLAP_LEN / 2); + const nWin = Math.ceil(padded.length / HOP_SIZE); + const frames: number[][] = [], onsets: number[][] = []; + const nOrig = Math.floor(mono.length * (FPS / SR)); + for (let w0 = 0; w0 < nWin && frames.length < nOrig; w0 += BATCH) { + const n = Math.min(BATCH, nWin - w0); + const input = new Float32Array(n * N_SAMPLES); + for (let b = 0; b < n; b++) { + const s = (w0 + b) * HOP_SIZE; + input.set(padded.subarray(s, Math.min(s + N_SAMPLES, padded.length)), b * N_SAMPLES); + } + const res = await session.run({ [INPUT]: new ort.Tensor("float32", input, [n, N_SAMPLES, 1]) }); + const note = res[OUT_NOTE]!, onset = res[OUT_ONSET]!; + const T = note.dims[1]!, B = note.dims[2]!; + const nd = note.data as Float32Array, od = onset.data as Float32Array; + for (let b = 0; b < n && frames.length < nOrig; b++) { + for (let t = OVERLAP_OVER_2; t < T - OVERLAP_OVER_2 && frames.length < nOrig; t++) { + const o = (b * T + t) * B; + frames.push(Array.from(nd.subarray(o, o + B))); + onsets.push(Array.from(od.subarray(o, o + B))); + } + } + } + if (frames.length === 0) return []; + + const minLen = Math.round(((opts.minNoteMs ?? 58) / 1000) * (SR / FFT_HOP)); + const minIdx = Math.max(0, (opts.minPitch ?? MIDI_OFFSET) - MIDI_OFFSET); + const maxIdx = Math.min(MAX_FREQ_IDX, (opts.maxPitch ?? MIDI_OFFSET + MAX_FREQ_IDX) - MIDI_OFFSET); + const events = outputToNotesPoly(frames, onsets, opts.onsetThresh ?? 0.5, opts.frameThresh ?? 0.3, minLen, minIdx, maxIdx); + const notes: Piece[] = events + .map((e) => { + const onset = modelFrameToTime(e.startFrame); + const offset = modelFrameToTime(e.startFrame + e.durationFrames); + return { onset, offset, pitch: e.pitchMidi, velocity: e.amplitude, amp: e.amplitude }; + }) + .sort((a, b) => a.onset - b.onset || a.pitch - b.pitch); + if (opts.monophonic) return makeMonophonic(notes, (opts.minNoteMs ?? 58) / 1000); + return notes.map(({ amp: _a, ...n }) => n); + }; +} diff --git a/tests/pipeline/transcribe.test.ts b/tests/pipeline/transcribe.test.ts new file mode 100644 index 0000000..633e7bd --- /dev/null +++ b/tests/pipeline/transcribe.test.ts @@ -0,0 +1,41 @@ +import { test, expect, beforeAll } from "bun:test"; +import { join } from "node:path"; +import { buildSong, loadSongDir } from "../../src/bench/dataset.ts"; +import { createTranscriber } from "../../src/pipeline/transcribe.ts"; +import { matchNotes } from "../../src/eval/notes.ts"; +import type { Audio, TimedNote } from "../../src/core/types.ts"; + +const MODEL = join(import.meta.dir, "../../models/basic-pitch/nmp.onnx"); +const SLOW = 300_000; + +let transcribe: Awaited>; + +beforeAll(async () => { + transcribe = await createTranscriber(new Uint8Array(await Bun.file(MODEL).arrayBuffer())); +}, SLOW); + +test("oracle guitar stem of t1a/1 scores F1 > 0.5", async () => { + const dir = await buildSong("t1a", 1); + const d = await loadSongDir(dir); + const track = d.song.tracks.find((t) => t.role === "guitar")!; + const ref = d.truth.tracks.find((t) => t.name === track.name)!.notes; + const est = await transcribe(d.stems.get(track.name)!); + const m = matchNotes(ref, est); + console.log(`guitar F1=${m.f1} P=${m.precision} R=${m.recall} n_ref=${ref.length} n_est=${est.length}`); + expect(m.f1!).toBeGreaterThan(0.5); +}, SLOW); + +test("monophonic mode yields non-overlapping notes on a vocals stem", async () => { + const dir = await buildSong("t1c", 1); + const d = await loadSongDir(dir); + const track = d.song.tracks.find((t) => t.role === "vocals")!; + const est: TimedNote[] = await transcribe(d.stems.get(track.name)!, { monophonic: true }); + expect(est.length).toBeGreaterThan(0); + const s = [...est].sort((a, b) => a.onset - b.onset); + for (let i = 1; i < s.length; i++) expect(s[i]!.onset).toBeGreaterThanOrEqual(s[i - 1]!.offset - 1e-9); +}, SLOW); + +test("silence returns no notes", async () => { + const a: Audio = { sampleRate: 44100, channels: [new Float32Array(88200)] }; + expect(await transcribe(a)).toEqual([]); +});