feat: heuristic drum 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
1d490f6897
commit
9bad6490eb
@@ -0,0 +1,121 @@
|
||||
// Heuristic drum transcriber ("heuristic baseline": no licensed open drum model runs in the TS/ONNX stack yet).
|
||||
// Band-split spectral flux -> adaptive peak picking per class. Pure DSP, browser-portable.
|
||||
import type { Audio, TimedNote } from "../core/types.ts";
|
||||
import { stft } from "../dsp/fft.ts";
|
||||
import { toMono } from "../dsp/mono.ts";
|
||||
|
||||
const N_FFT = 1024;
|
||||
const HOP = 256;
|
||||
const MIN_IOI = 0.03; // s, per class
|
||||
const K_MAD = 6; // threshold = median + K_MAD * MAD
|
||||
const ABS_FLOOR = 0.4; // absolute flux floor (log-compressed units); keeps silence/noise floor silent
|
||||
const CYMBAL_DECAY = 0.3; // s: hat-band energy that persists this long marks a cymbal
|
||||
|
||||
const REL = 0.3; // also >= REL * 99th-percentile flux;
|
||||
const CK = 1, CS = 1; // log-compression gain on peak-normalised magnitudes (low gain keeps transient leakage out of the kick band)
|
||||
const KICK_BLEED = 0.15;
|
||||
const KICK = 36, SNARE = 38, HAT = 42, CYMBAL = 49;
|
||||
|
||||
/** Log-compressed half-wave-rectified flux summed over bins [lo, hi). */
|
||||
function bandFlux(mags: Float32Array[], lo: number, hi: number, c = 1000): Float64Array {
|
||||
const out = new Float64Array(mags.length);
|
||||
for (let f = 1; f < mags.length; f++) {
|
||||
let s = 0;
|
||||
const a = mags[f]!, b = mags[f - 1]!;
|
||||
for (let k = lo; k < hi; k++) {
|
||||
const d = Math.log1p(c * a[k]!) - Math.log1p(c * b[k]!);
|
||||
if (d > 0) s += d;
|
||||
}
|
||||
out[f] = s;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function median(xs: ArrayLike<number>): number {
|
||||
const a = Array.from(xs).sort((x, y) => x - y);
|
||||
return a.length ? a[a.length >> 1]! : 0;
|
||||
}
|
||||
|
||||
/** Local maxima above median + k*MAD (and an absolute floor), at least `minGap` frames apart (higher peak wins). */
|
||||
function pickPeaks(env: Float64Array, k: number, floor: number, minGap: number): number[] {
|
||||
const med = median(env);
|
||||
const mad = median(env.map((v) => Math.abs(v - med)));
|
||||
const sorted = Array.from(env).sort((a, b) => a - b);
|
||||
const p99 = sorted[Math.floor(sorted.length * 0.99)] ?? 0;
|
||||
const thr = Math.max(med + k * mad, floor, REL * p99);
|
||||
const peaks: number[] = [];
|
||||
const w = 2;
|
||||
for (let i = 0; i < env.length; i++) {
|
||||
const v = env[i]!;
|
||||
if (v < thr) continue;
|
||||
let isMax = true;
|
||||
for (let j = Math.max(0, i - w); j <= Math.min(env.length - 1, i + w); j++) {
|
||||
if (env[j]! > v || (env[j]! === v && j < i)) { isMax = false; break; }
|
||||
}
|
||||
if (!isMax) continue;
|
||||
const last = peaks[peaks.length - 1];
|
||||
if (last !== undefined && i - last < minGap) {
|
||||
if (v > env[last]!) peaks[peaks.length - 1] = i;
|
||||
continue;
|
||||
}
|
||||
peaks.push(i);
|
||||
}
|
||||
return peaks;
|
||||
}
|
||||
|
||||
export function transcribeDrums(audio: Audio): TimedNote[] {
|
||||
const x = toMono(audio);
|
||||
let peak = 0;
|
||||
for (let i = 0; i < x.length; i++) peak = Math.max(peak, Math.abs(x[i]!));
|
||||
if (peak < 1e-4) return [];
|
||||
for (let i = 0; i < x.length; i++) x[i]! /= peak; // level-invariant thresholds
|
||||
if (x.length <= N_FFT) return [];
|
||||
const sr = audio.sampleRate;
|
||||
const spec = stft(x, N_FFT, HOP);
|
||||
const bins = N_FFT / 2 + 1;
|
||||
const mags = spec.re.map((re, f) => {
|
||||
const im = spec.im[f]!, m = new Float32Array(bins);
|
||||
for (let k = 0; k < bins; k++) m[k] = Math.hypot(re[k]!, im[k]!);
|
||||
return m;
|
||||
});
|
||||
const bin = (hz: number) => Math.min(bins, Math.round((hz * N_FFT) / sr));
|
||||
const dt = HOP / sr;
|
||||
const minGap = Math.max(1, Math.round(MIN_IOI / dt));
|
||||
|
||||
const kickEnv = bandFlux(mags, 1, bin(150), CK);
|
||||
const snareLow = bandFlux(mags, bin(150), bin(2000), CS);
|
||||
const snareHigh = bandFlux(mags, bin(2000), bin(7000), CS);
|
||||
const hatEnv = bandFlux(mags, bin(7000), bins);
|
||||
const snareEnv = snareLow.map((v, i) => Math.min(v, 2 * snareHigh[i]!)); // needs broadband noise, not just low tonal energy
|
||||
|
||||
// hat-band level (linear energy) for decay measurement
|
||||
const hatLevel = new Float64Array(mags.length);
|
||||
for (let f = 0; f < mags.length; f++) {
|
||||
let s = 0;
|
||||
for (let k = bin(7000); k < bins; k++) s += mags[f]![k]! ** 2;
|
||||
hatLevel[f] = Math.sqrt(s);
|
||||
}
|
||||
|
||||
const notes: TimedNote[] = [];
|
||||
const add = (frames: number[], pitch: number, dur: number) => {
|
||||
for (const f of frames) notes.push({ onset: f * dt, offset: f * dt + dur, pitch, velocity: 100 });
|
||||
};
|
||||
const kicks = pickPeaks(kickEnv, K_MAD, ABS_FLOOR, minGap);
|
||||
const nearMax = (e: Float64Array, f: number) => Math.max(e[f - 1] ?? 0, e[f]!, e[f + 1] ?? 0);
|
||||
const snares = pickPeaks(snareEnv, K_MAD, ABS_FLOOR, minGap).filter((f) => nearMax(kickEnv, f) < KICK_BLEED * nearMax(snareEnv, f)); // kick clicks are broadband too: drop snare candidates carried by low-band flux
|
||||
const hatPeaks = pickPeaks(hatEnv, K_MAD, ABS_FLOOR, minGap);
|
||||
|
||||
const hats: number[] = [], cymbals: number[] = [];
|
||||
const decayFrames = Math.round(CYMBAL_DECAY / dt);
|
||||
for (const f of hatPeaks) {
|
||||
let peak = 0;
|
||||
for (let j = f; j <= Math.min(mags.length - 1, f + 3); j++) peak = Math.max(peak, hatLevel[j]!);
|
||||
const after = hatLevel[Math.min(mags.length - 1, f + decayFrames)]!;
|
||||
(peak > 0 && after / peak > 0.25 ? cymbals : hats).push(f);
|
||||
}
|
||||
add(kicks, KICK, 0.1);
|
||||
add(snares, SNARE, 0.1);
|
||||
add(hats, HAT, 0.05);
|
||||
add(cymbals, CYMBAL, 0.5);
|
||||
return notes.sort((a, b) => a.onset - b.onset || a.pitch - b.pitch);
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
import { test, expect, beforeAll } from "bun:test";
|
||||
import { join } from "node:path";
|
||||
import { createSeparator } from "../../src/pipeline/separate.ts";
|
||||
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
|
||||
import { transcribeDrums } from "../../src/pipeline/drums.ts";
|
||||
import { matchNotes } from "../../src/eval/notes.ts";
|
||||
import { DRUM_CLASSES } from "../../src/eval/score.ts";
|
||||
import type { Audio, TimedNote } from "../../src/core/types.ts";
|
||||
import { ticksToSeconds } from "../../src/core/time.ts";
|
||||
|
||||
const SLOW = 300_000;
|
||||
const SEEDS = [1, 2, 3];
|
||||
const CLASSES = ["kick", "snare", "hihat", "cymbal"] as const;
|
||||
|
||||
const classOf = (p: number) => CLASSES.find((c) => DRUM_CLASSES[c].includes(p));
|
||||
const only = (notes: TimedNote[], c: string) => notes.filter((n) => classOf(n.pitch) === c).map((n) => ({ ...n, pitch: 0 }));
|
||||
|
||||
const songs: { ref: TimedNote[]; stem: Audio }[] = [];
|
||||
|
||||
beforeAll(async () => {
|
||||
for (const seed of SEEDS) {
|
||||
const d = await loadSongDir(await buildSong("t1b", seed));
|
||||
const t = d.song.tracks.find((x) => x.role === "drums")!;
|
||||
songs.push({ ref: t.notes.map((n) => ({ onset: ticksToSeconds(n.tick, d.song.bpm), offset: ticksToSeconds(n.tick + n.dur, d.song.bpm), pitch: n.pitch })), stem: d.stems.get(t.name)! });
|
||||
}
|
||||
}, SLOW);
|
||||
|
||||
function f1(c: string): number {
|
||||
let tp = 0, fp = 0, fn = 0;
|
||||
for (const s of songs) {
|
||||
const m = matchNotes(only(s.ref, c), only(transcribeDrums(s.stem), c));
|
||||
tp += m.tp; fp += m.fp; fn += m.fn;
|
||||
}
|
||||
const f = tp + fp + fn === 0 ? 0 : (2 * tp) / (2 * tp + fp + fn);
|
||||
console.log(`heuristic drums oracle t1b seeds 1-3 ${c}: F1=${f.toFixed(3)} tp=${tp} fp=${fp} fn=${fn}`);
|
||||
return f;
|
||||
}
|
||||
|
||||
test("oracle drum stems: kick F1 > 0.7", () => { expect(f1("kick")).toBeGreaterThan(0.7); }, SLOW);
|
||||
test("oracle drum stems: snare F1 > 0.6", () => { expect(f1("snare")).toBeGreaterThan(0.6); }, SLOW);
|
||||
test("reports hihat and cymbal F1", () => { f1("hihat"); f1("cymbal"); }, SLOW);
|
||||
|
||||
test("silence yields no notes", () => {
|
||||
expect(transcribeDrums({ sampleRate: 44100, channels: [new Float32Array(44100 * 3)] })).toEqual([]);
|
||||
});
|
||||
|
||||
test("demucs drum stem of t1b/1 (report only)", async () => {
|
||||
const model = Bun.file(join(import.meta.dir, "../../models/demucs/htdemucs_6s.onnx"));
|
||||
if (!(await model.exists())) return;
|
||||
const separate = await createSeparator(new Uint8Array(await model.arrayBuffer()), { intraOpNumThreads: 8 });
|
||||
const d = await loadSongDir(await buildSong("t1b", 1));
|
||||
const est = transcribeDrums((await separate(d.mix)).drums);
|
||||
for (const c of CLASSES) {
|
||||
const m = matchNotes(songs[0]!.ref.filter((n) => classOf(n.pitch) === c).map((n) => ({ ...n, pitch: 0 })), only(est, c));
|
||||
console.log(`heuristic drums demucs t1b/1 ${c}: F1=${(m.f1 ?? 0).toFixed(3)} tp=${m.tp} fp=${m.fp} fn=${m.fn}`);
|
||||
}
|
||||
}, SLOW);
|
||||
Reference in New Issue
Block a user