99 lines
4.5 KiB
TypeScript
99 lines
4.5 KiB
TypeScript
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 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 }));
|
|
|
|
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) songs.push(await load(seed));
|
|
heldOut = await load(HELD_OUT);
|
|
}, SLOW);
|
|
|
|
function f1(c: string, set: Item[] = songs, label = "seeds 1-3"): number {
|
|
let tp = 0, fp = 0, fn = 0;
|
|
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 ${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("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.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 });
|
|
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);
|