lucasw89/STSheets
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: drum transcriber noise gate, kick+snare, test floors

Browse Source 
Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
2026-10-30 04:55:06 +01:00
co-authored by Claude Sonnet 5.5
parent 9bad6490eb
commit eab70b24f6
2 changed files with 75 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pipeline/drums.ts
+23 -11
View File
@@ -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);
tests/pipeline/drums.test.ts
+52 -11
View File
@@ -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<Item> {
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 });