feat: dataset builder
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
a1e9ce86a0
commit
ca9f3bbb83
+19
-5
@@ -2,27 +2,41 @@ import type { Audio } from "../core/types.ts";
|
||||
|
||||
const PEAK_TARGET = Math.pow(10, -1 / 20);
|
||||
|
||||
/** Sum stems (with per-stem gain in dB), then peak-normalise to -1 dBFS. */
|
||||
export function mixStems(stems: { audio: Audio; gainDb: number }[]): Audio {
|
||||
/**
|
||||
* Sum stems (with per-stem gain in dB), then peak-normalise to -1 dBFS.
|
||||
* Also returns each input stem as it contributed to the mix (gain and normalisation applied, padded to the
|
||||
* mix's channel count and length), so that sum(scaled) === mix, and the applied normalisation factor.
|
||||
*/
|
||||
export function mixStemsScaled(stems: { audio: Audio; gainDb: number }[]): { mix: Audio; scaled: Audio[]; scale: number } {
|
||||
if (stems.length === 0) throw new Error("mixStems: no stems");
|
||||
const sr = stems[0]!.audio.sampleRate;
|
||||
const nch = Math.max(...stems.map((s) => s.audio.channels.length));
|
||||
const n = Math.max(...stems.map((s) => s.audio.channels[0]?.length ?? 0));
|
||||
const out = Array.from({ length: nch }, () => new Float32Array(n));
|
||||
const scaled: Audio[] = [];
|
||||
for (const { audio, gainDb } of stems) {
|
||||
if (audio.sampleRate !== sr) throw new Error("mixStems: sample rate mismatch");
|
||||
const g = Math.pow(10, gainDb / 20);
|
||||
const own = Array.from({ length: nch }, () => new Float32Array(n));
|
||||
for (let c = 0; c < nch; c++) {
|
||||
const src = audio.channels[c] ?? audio.channels[0]!;
|
||||
const dst = out[c]!;
|
||||
for (let i = 0; i < src.length; i++) dst[i]! += src[i]! * g;
|
||||
const o = own[c]!;
|
||||
for (let i = 0; i < src.length; i++) { const v = src[i]! * g; dst[i]! += v; o[i] = v; }
|
||||
}
|
||||
scaled.push({ sampleRate: sr, channels: own });
|
||||
}
|
||||
let peak = 0;
|
||||
for (const ch of out) for (let i = 0; i < n; i++) peak = Math.max(peak, Math.abs(ch[i]!));
|
||||
const k = peak > 0 ? PEAK_TARGET / peak : 1;
|
||||
if (peak > 0) {
|
||||
const k = PEAK_TARGET / peak;
|
||||
for (const ch of out) for (let i = 0; i < n; i++) ch[i]! *= k;
|
||||
for (const s of scaled) for (const ch of s.channels) for (let i = 0; i < n; i++) ch[i]! *= k;
|
||||
}
|
||||
return { sampleRate: sr, channels: out };
|
||||
return { mix: { sampleRate: sr, channels: out }, scaled, scale: k };
|
||||
}
|
||||
|
||||
/** Sum stems (with per-stem gain in dB), then peak-normalise to -1 dBFS. */
|
||||
export function mixStems(stems: { audio: Audio; gainDb: number }[]): Audio {
|
||||
return mixStemsScaled(stems).mix;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
import { buildSong } from "./dataset.ts";
|
||||
import type { TierId } from "../gen/tiers.ts";
|
||||
import { TIERS } from "../gen/tiers.ts";
|
||||
|
||||
function parseSeeds(s: string): number[] {
|
||||
const out: number[] = [];
|
||||
for (const part of s.split(",")) {
|
||||
const m = /^(\d+)-(\d+)$/.exec(part);
|
||||
if (m) for (let i = Number(m[1]); i <= Number(m[2]); i++) out.push(i);
|
||||
else if (/^\d+$/.test(part)) out.push(Number(part));
|
||||
else throw new Error(`bad seeds: ${s}`);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function arg(args: string[], name: string, def: string): string {
|
||||
const i = args.indexOf(`--${name}`);
|
||||
return i >= 0 && args[i + 1] ? args[i + 1]! : def;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
const [cmd, ...rest] = process.argv.slice(2);
|
||||
switch (cmd) {
|
||||
case "gen": {
|
||||
const tiers = arg(rest, "tiers", "t1a,t1b,t1c,t1d").split(",") as TierId[];
|
||||
for (const t of tiers) if (!(t in TIERS)) throw new Error(`unknown tier: ${t}`);
|
||||
const seeds = parseSeeds(arg(rest, "seeds", "1-5"));
|
||||
const root = arg(rest, "root", "data");
|
||||
for (const t of tiers) for (const s of seeds) {
|
||||
const t0 = performance.now();
|
||||
const dir = await buildSong(t, s, root);
|
||||
console.log(`${dir} (${((performance.now() - t0) / 1000).toFixed(1)}s)`);
|
||||
}
|
||||
break;
|
||||
}
|
||||
default:
|
||||
console.error("usage: bun run bench gen [--tiers t1a,t1b] [--seeds 1-5] [--root data]");
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
await main();
|
||||
@@ -0,0 +1,55 @@
|
||||
import { join } from "node:path";
|
||||
import { mkdir } from "node:fs/promises";
|
||||
import type { Audio, Song } from "../core/types.ts";
|
||||
import type { TierId } from "../gen/tiers.ts";
|
||||
import { generateSong } from "../gen/generate.ts";
|
||||
import { songToGp, truthFromGp } from "../score/gp.ts";
|
||||
import { renderStems } from "../audio/render.ts";
|
||||
import { applyTrackFx, pan } from "../audio/fx.ts";
|
||||
import { mixStemsScaled } from "../audio/mix.ts";
|
||||
import { readWav, writeWav } from "../audio/wav.ts";
|
||||
import { encodeMp3 } from "../audio/mp3.ts";
|
||||
|
||||
const SF2_PATH = "node_modules/@coderline/alphatab/dist/soundfont/sonivox.sf2";
|
||||
const SR = 44100;
|
||||
|
||||
/**
|
||||
* Build one song into `<root>/<tier>-<seed>/`. Idempotent: skipped when mix.mp3 exists and song.json matches.
|
||||
* stems/<track>.wav are exactly each track's contribution to mix.wav (post-fx, post-pan, gain, normalisation),
|
||||
* so the stems sum to the mix up to PCM16 rounding.
|
||||
*/
|
||||
export async function buildSong(tier: TierId, seed: number, root = "data"): Promise<string> {
|
||||
const dir = join(root, `${tier}-${seed}`);
|
||||
const song = generateSong(tier, seed);
|
||||
const songJson = JSON.stringify(song, null, 1);
|
||||
const jsonFile = Bun.file(join(dir, "song.json"));
|
||||
if ((await Bun.file(join(dir, "mix.mp3")).exists()) && (await jsonFile.exists()) && (await jsonFile.text()) === songJson) {
|
||||
return dir;
|
||||
}
|
||||
await mkdir(join(dir, "stems"), { recursive: true });
|
||||
const sf2 = new Uint8Array(await Bun.file(SF2_PATH).arrayBuffer());
|
||||
const dry = renderStems(song, sf2, SR);
|
||||
const inputs = song.tracks.map((t) => {
|
||||
const raw = dry.get(t.name);
|
||||
if (!raw) throw new Error(`buildSong: no rendered stem for ${t.name}`);
|
||||
return { audio: pan(applyTrackFx(raw, t.role, t.fx, SR), t.pan), gainDb: t.gainDb };
|
||||
});
|
||||
const { mix, scaled } = mixStemsScaled(inputs);
|
||||
await Bun.write(join(dir, "truth.gp"), songToGp(song));
|
||||
for (let i = 0; i < song.tracks.length; i++) await writeWav(join(dir, "stems", `${song.tracks[i]!.name}.wav`), scaled[i]!);
|
||||
await writeWav(join(dir, "mix.wav"), mix);
|
||||
await encodeMp3(join(dir, "mix.wav"), join(dir, "mix.mp3"));
|
||||
await Bun.write(join(dir, "song.json"), songJson); // written last: marks the build complete
|
||||
return dir;
|
||||
}
|
||||
|
||||
export async function loadSongDir(dir: string): Promise<{
|
||||
song: Song; truth: ReturnType<typeof truthFromGp>; mix: Audio; stems: Map<string, Audio>;
|
||||
}> {
|
||||
const song = (await Bun.file(join(dir, "song.json")).json()) as Song;
|
||||
const truth = truthFromGp(new Uint8Array(await Bun.file(join(dir, "truth.gp")).arrayBuffer()));
|
||||
const mix = await readWav(join(dir, "mix.wav"));
|
||||
const stems = new Map<string, Audio>();
|
||||
for (const t of song.tracks) stems.set(t.name, await readWav(join(dir, "stems", `${t.name}.wav`)));
|
||||
return { song, truth, mix, stems };
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
import { test, expect } from "bun:test";
|
||||
import { mkdtemp, rm } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
|
||||
|
||||
test("buildSong writes all artefacts; loadSongDir reads them; stems sum to mix; idempotent", async () => {
|
||||
const root = await mkdtemp(join(tmpdir(), "stsheets-ds-"));
|
||||
try {
|
||||
const dir = await buildSong("t1a", 1, root);
|
||||
const files = ["song.json", "truth.gp", "mix.wav", "mix.mp3"];
|
||||
for (const f of files) expect(await Bun.file(join(dir, f)).exists()).toBe(true);
|
||||
const { song, truth, mix, stems } = await loadSongDir(dir);
|
||||
for (const t of song.tracks) expect(await Bun.file(join(dir, "stems", `${t.name}.wav`)).exists()).toBe(true);
|
||||
expect(truth.tracks[0]!.notes.length).toBe(song.tracks[0]!.notes.length);
|
||||
expect(mix.channels.length).toBe(2);
|
||||
|
||||
// Oracle stems are the exact mix contributions: they sum to the mix (PCM16 rounding only).
|
||||
expect(stems.size).toBe(song.tracks.length);
|
||||
let maxDiff = 0;
|
||||
for (let c = 0; c < mix.channels.length; c++) {
|
||||
const m = mix.channels[c]!;
|
||||
for (let i = 0; i < m.length; i++) {
|
||||
let s = 0;
|
||||
for (const a of stems.values()) s += a.channels[c]![i] ?? 0;
|
||||
maxDiff = Math.max(maxDiff, Math.abs(s - m[i]!));
|
||||
}
|
||||
}
|
||||
expect(maxDiff).toBeLessThan(2e-3);
|
||||
|
||||
const before = (await Bun.file(join(dir, "mix.mp3")).stat()).mtimeMs;
|
||||
expect(await buildSong("t1a", 1, root)).toBe(dir);
|
||||
expect((await Bun.file(join(dir, "mix.mp3")).stat()).mtimeMs).toBe(before);
|
||||
} finally {
|
||||
await rm(root, { recursive: true, force: true });
|
||||
}
|
||||
}, 120_000);
|
||||
Reference in New Issue
Block a user