diff --git a/src/pipeline/drums.ts b/src/pipeline/drums.ts new file mode 100644 index 0000000..c58f7a7 --- /dev/null +++ b/src/pipeline/drums.ts @@ -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 { + 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); +} diff --git a/tests/pipeline/drums.test.ts b/tests/pipeline/drums.test.ts new file mode 100644 index 0000000..02b252c --- /dev/null +++ b/tests/pipeline/drums.test.ts @@ -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);