From eab70b24f66ff2896c1a082558fb10c5ce3cdfe3 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Fri, 30 Oct 2026 04:55:06 +0100 Subject: [PATCH] fix: drum transcriber noise gate, kick+snare, test floors Co-Authored-By: Claude Sonnet 5.5 --- src/pipeline/drums.ts | 34 ++++++++++++------- tests/pipeline/drums.test.ts | 63 +++++++++++++++++++++++++++++------- 2 files changed, 75 insertions(+), 22 deletions(-) diff --git a/src/pipeline/drums.ts b/src/pipeline/drums.ts index c58f7a7..df24938 100644 --- a/src/pipeline/drums.ts +++ b/src/pipeline/drums.ts @@ -11,9 +11,14 @@ 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 REL = 0.3; // threshold also >= REL * 99th-percentile flux +const MIN_OVER_MED = 3; +// log-compression gain on peak-normalised magnitudes (low gain keeps transient leakage out of the kick/snare bands) +const GAIN = 1; +const KICK_NEAR = 0.08; // s +const SNARE_RATIO = 0.82; +const SNARE_TAIL = 0.12; // s +const MIN_PEAK = 3e-3; // absolute gate (~-50 dBFS): quieter stems are treated as silence const KICK = 36, SNARE = 38, HAT = 42, CYMBAL = 49; /** Log-compressed half-wave-rectified flux summed over bins [lo, hi). */ @@ -42,10 +47,10 @@ function pickPeaks(env: Float64Array, k: number, floor: number, minGap: number): 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 thr = Math.max(med + k * mad, floor, REL * p99, MIN_OVER_MED * med); // stationary noise has max ~2x median const peaks: number[] = []; const w = 2; - for (let i = 0; i < env.length; i++) { + for (let i = 2; i < env.length; i++) { // frames 0-1 only see the reflect-padded edge: flux there is an artefact const v = env[i]!; if (v < thr) continue; let isMax = true; @@ -64,12 +69,12 @@ function pickPeaks(env: Float64Array, k: number, floor: number, minGap: number): } export function transcribeDrums(audio: Audio): TimedNote[] { + if ((audio.channels[0]?.length ?? 0) <= N_FFT) return []; 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 []; + if (peak < MIN_PEAK) 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; @@ -82,9 +87,9 @@ export function transcribeDrums(audio: Audio): TimedNote[] { 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 kickEnv = bandFlux(mags, 1, bin(150), GAIN); + const snareLow = bandFlux(mags, bin(150), bin(2000), GAIN); + const snareHigh = bandFlux(mags, bin(2000), bin(7000), GAIN); 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 @@ -102,7 +107,6 @@ export function transcribeDrums(audio: Audio): TimedNote[] { }; 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[] = []; @@ -113,6 +117,14 @@ export function transcribeDrums(audio: Audio): TimedNote[] { const after = hatLevel[Math.min(mags.length - 1, f + decayFrames)]!; (peak > 0 && after / peak > 0.25 ? cymbals : hats).push(f); } + // Snare candidates near a kick hit (kick clicks and kick+hat are broadband too) must be mid-band dominated like a snare + // (low/high flux ratio), not carried by 2-7 kHz hat energy; an isolated candidate needs no such check. + const kickNear = (f: number) => kicks.some((k) => Math.abs(k - f) * dt <= KICK_NEAR); + const cands = pickPeaks(snareEnv, K_MAD, ABS_FLOOR, minGap).filter( + (f) => !kickNear(f) || (snareLow[f]! >= SNARE_RATIO * snareHigh[f]! && !cymbals.some((c) => Math.abs(c - f) * dt <= 0.03)), // a sustained crash under a kick is not a snare + ); + // drop weak tails: a candidate within SNARE_TAIL of a much stronger one is ringing, not a new hit + const snares = cands.filter((f) => !cands.some((g) => g < f && (f - g) * dt <= SNARE_TAIL && snareEnv[f]! < 0.3 * snareEnv[g]!)); add(kicks, KICK, 0.1); add(snares, SNARE, 0.1); add(hats, HAT, 0.05); diff --git a/tests/pipeline/drums.test.ts b/tests/pipeline/drums.test.ts index 02b252c..7d87c70 100644 --- a/tests/pipeline/drums.test.ts +++ b/tests/pipeline/drums.test.ts @@ -10,41 +10,82 @@ import { ticksToSeconds } from "../../src/core/time.ts"; const SLOW = 300_000; const SEEDS = [1, 2, 3]; +const HELD_OUT = 4; 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 }[] = []; +type Item = { ref: TimedNote[]; stem: Audio }; +const songs: Item[] = []; +let heldOut: Item; + +async function load(seed: number): Promise { + const d = await loadSongDir(await buildSong("t1b", seed)); + const t = d.song.tracks.find((x) => x.role === "drums")!; + return { 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)! }; +} 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)! }); - } + for (const seed of SEEDS) songs.push(await load(seed)); + heldOut = await load(HELD_OUT); }, SLOW); -function f1(c: string): number { +function f1(c: string, set: Item[] = songs, label = "seeds 1-3"): number { let tp = 0, fp = 0, fn = 0; - for (const s of songs) { + for (const s of set) { 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}`); + console.log(`heuristic drums oracle t1b ${label} ${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("oracle drum stems: hihat F1 > 0.9, cymbal F1 > 0.6", () => { + expect(f1("hihat")).toBeGreaterThan(0.9); + expect(f1("cymbal")).toBeGreaterThan(0.6); +}, SLOW); + +test("held-out seed 4: kick > 0.6, snare > 0.5, hihat > 0.8", () => { + const h = [heldOut]; + expect(f1("kick", h, "seed 4")).toBeGreaterThan(0.6); + expect(f1("snare", h, "seed 4")).toBeGreaterThan(0.5); + expect(f1("hihat", h, "seed 4")).toBeGreaterThan(0.8); + f1("cymbal", h, "seed 4"); +}, SLOW); + +test("quiet noise-only stem (1e-3 uniform, 5 s) yields no notes", () => { + let seed = 7; + const rnd = () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 2 ** 32) * 2 - 1; + const ch = new Float32Array(44100 * 5).map(() => rnd() * 1e-3); + expect(transcribeDrums({ sampleRate: 44100, channels: [ch] })).toEqual([]); +}); + +test("simultaneous kick+snare at 1/2/3 s: one kick and one snare each, no extra snares", () => { + const sr = 44100, x = new Float32Array(sr * 5); + let seed = 1; + const rnd = () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 2 ** 32) * 2 - 1; + for (const t of [1, 2, 3]) { + const o = t * sr; + for (let i = 0; i < sr * 0.3; i++) x[o + i]! += 0.6 * Math.exp(-i / (sr * 0.06)) * Math.sin((2 * Math.PI * 55 * i) / sr); + for (let i = 0; i < sr * 0.25; i++) x[o + i]! += 0.4 * Math.exp(-i / (sr * 0.05)) * (rnd() * 0.8 + Math.sin((2 * Math.PI * 190 * i) / sr)); + } + const est = transcribeDrums({ sampleRate: sr, channels: [x] }); + for (const p of [36, 38]) { + const times = est.filter((n) => n.pitch === p).map((n) => n.onset); + expect(times.length).toBe(3); + [1, 2, 3].forEach((t, i) => expect(Math.abs(times[i]! - t)).toBeLessThan(0.03)); + } +}); 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 () => { +test.skipIf(!process.env.STSHEETS_SLOW)("demucs drum stem of t1b/1 (report only; STSHEETS_SLOW=1)", 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 });