feat: basic pitch transcription
Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
b5e6f63881
commit
1239e5e0ac
@@ -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<number>(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<TimedNote[]>> {
|
||||
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);
|
||||
};
|
||||
}
|
||||
@@ -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<ReturnType<typeof createTranscriber>>;
|
||||
|
||||
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([]);
|
||||
});
|
||||
Reference in New Issue
Block a user