Add overlap pursuit classifier and update composer - proof-of-concept/instrawave-compose

proof-of-concept / instrawave-compose

Add overlap pursuit classifier and update composer

Jeremy Magland <jmagland@flatironinstitute.org> committed 20 days ago commit 7d7cc56b172d parent b11b069 Browse files

6 changed files+336−15

index.htmlmodified+1−0View file

@@ -110,6 +110,7 @@

110110 <label>classifier

111111 <select id="clf-method">

112112 <option value="builtin">silence_gap_knn (numbl)</option>

113+ <option value="pursuit">overlap_pursuit (numbl)</option>

113114 <option value="custom">custom (numbl)</option>

114115 <option value="legacy">feature kNN (built-in)</option>

115116 </select>

scripts/test-node.tsmodified+27−1View file

@@ -134,7 +134,7 @@ function burst(x: Float32Array, t0: number, freq: number, decay: number, amp: nu

134134 const { buildVoiceTraining } = await import('../src/numbl/snippets.ts');

135135 const { runAlgorithm } = await import('../src/numbl/execute.ts');

136136 const { downsampleBox } = await import('../src/numbl/dsp.ts');

137-const { DEFAULT_CLASSIFIER_SOURCE } = await import('../src/numbl/index.ts');

137+const { DEFAULT_CLASSIFIER_SOURCE, BUILTIN_CLASSIFIERS } = await import('../src/numbl/index.ts');

138138 const { CLF_FS } = await import('../src/numbl/types.ts');

139139

140140 {

@@ -189,6 +189,32 @@ const { CLF_FS } = await import('../src/numbl/types.ts');

189189 });

190190 check(timesOk, 'numbl classify: every hit within 30 ms of the truth');

191191 check(labelsOk, 'numbl classify: thump/tick labels all correct');

192+

193+ // overlap_pursuit on the same well-separated take: the ring-outs run

194+ // longer than one snippet, so this checks the pursuit does not invent

195+ // hits out of decay tails

196+ const pursuitSrc = BUILTIN_CLASSIFIERS.pursuit.source;

197+ {

198+ const det = runAlgorithm(pursuitSrc, x16, CLF_FS, training).detections;

199+ const ok = det.length === takeTimes.length

200+ && det.every((d, i) => Math.abs(d.t - takeTimes[i]) <= 0.03 && d.typeIndex === i % 2);

201+ check(ok, `overlap_pursuit: ${det.length} of ${takeTimes.length} separated hits, no phantoms`);

202+ }

203+

204+ // an overlapping pair 80 ms apart: the two regions merge, which the

205+ // default classifier cannot pull apart but the pursuit should

206+ {

207+ const pair = new Float32Array(Math.floor(1.5 * SR));

208+ noiseInto(pair, 5e-4, 555);

209+ burst(pair, 0.5, 200, 0.05, 0.4, 41);

210+ burst(pair, 0.58, 1500, 0.015, 0.4, 42);

211+ const p16 = downsampleBox(pair, SR, CLF_FS);

212+ const det = runAlgorithm(pursuitSrc, p16, CLF_FS, training).detections;

213+ const hasThump = det.some((d) => d.typeIndex === 0 && Math.abs(d.t - 0.5) <= 0.03);

214+ const hasTick = det.some((d) => d.typeIndex === 1 && Math.abs(d.t - 0.58) <= 0.03);

215+ check(det.length === 2 && hasThump && hasTick,

216+ `overlap_pursuit: overlapping thump+tick both recovered (got ${det.map((d) => `${d.typeIndex}@${d.t.toFixed(3)}`).join(', ') || 'none'})`);

217+ }

192218 }

193219

194220 console.log(failures ? `\n${failures} failure(s)` : '\nall checks passed');

src/main.tsmodified+13−7View file

@@ -26,7 +26,7 @@ import {

2626 saveClassifier, saveIwOverrides, saveTakeHits, saveUi, saveVoices, trainPool,

2727 type TakeHit, type TrainClip, type Voice,

2828 } from './store.ts';

29-import { CUSTOM_CLASSIFIER_TEMPLATE, DEFAULT_CLASSIFIER_NAME, DEFAULT_CLASSIFIER_SOURCE } from './numbl/index.ts';

29+import { BUILTIN_CLASSIFIERS, CUSTOM_CLASSIFIER_TEMPLATE, type BuiltinMethod } from './numbl/index.ts';

3030 import { buildVoiceTraining } from './numbl/snippets.ts';

3131 import { classifyInWorker } from './numbl/client.ts';

3232 import { downsampleBox } from './numbl/dsp.ts';

@@ -685,7 +685,8 @@ const timeline = new Timeline($('timeline') as HTMLCanvasElement);

685685 const clf = loadClassifier();

686686 const activeClfSource = (): string =>

687687 clf.method === 'custom' ? (clf.customSrc ?? CUSTOM_CLASSIFIER_TEMPLATE)

688- : (clf.builtinSrc ?? DEFAULT_CLASSIFIER_SOURCE);

688+ : clf.method === 'pursuit' ? (clf.pursuitSrc ?? BUILTIN_CLASSIFIERS.pursuit.source)

689+ : (clf.builtinSrc ?? BUILTIN_CLASSIFIERS.builtin.source);

689690

690691 async function processTake(): Promise<void> {

691692 if (!take.clip) return;

@@ -720,7 +721,8 @@ async function processTake(): Promise<void> {

720721 take.hits = hits;

721722 saveTakeHits(take.hits);

722723 refreshTimeline();

723- const algoName = clf.method === 'builtin' ? DEFAULT_CLASSIFIER_NAME : 'the custom algorithm';

724+ const algoName = clf.method === 'custom'

725+ ? 'the custom algorithm' : BUILTIN_CLASSIFIERS[clf.method as BuiltinMethod].name;

724726 const skippedNote = skipped.length ? `; no training audio for ${skipped.join(', ')}` : '';

725727 status(`${hits.length} hits classified by ${algoName} across ${voiceIds.length} trained voice${voiceIds.length === 1 ? '' : 's'}${skippedNote}`);

726728 } catch (err) {

@@ -817,9 +819,11 @@ function syncClfBar(): void {

817819

818820 function syncClfEditor(): void {

819821 if ($('clf-editor').hidden) return;

820- $('clf-title').textContent = clf.method === 'custom' ? 'custom classifier' : DEFAULT_CLASSIFIER_NAME;

822+ $('clf-title').textContent = clf.method === 'custom'

823+ ? 'custom classifier' : BUILTIN_CLASSIFIERS[clf.method as BuiltinMethod].name;

821824 ($('clf-source') as HTMLTextAreaElement).value = activeClfSource();

822- ($('clf-reset') as HTMLButtonElement).hidden = clf.method !== 'builtin' || clf.builtinSrc === null;

825+ ($('clf-reset') as HTMLButtonElement).hidden = clf.method === 'custom' || clf.method === 'legacy'

826+ || (clf.method === 'builtin' ? clf.builtinSrc : clf.pursuitSrc) === null;

823827 }

824828

825829 clfMethodSel.addEventListener('change', () => {

@@ -839,12 +843,14 @@ $('clf-close').addEventListener('click', () => { $('clf-editor').hidden = true;

839843 ($('clf-source') as HTMLTextAreaElement).addEventListener('change', () => {

840844 const src = ($('clf-source') as HTMLTextAreaElement).value;

841845 if (clf.method === 'custom') clf.customSrc = src;

842- else clf.builtinSrc = src === DEFAULT_CLASSIFIER_SOURCE ? null : src;

846+ else if (clf.method === 'pursuit') clf.pursuitSrc = src === BUILTIN_CLASSIFIERS.pursuit.source ? null : src;

847+ else clf.builtinSrc = src === BUILTIN_CLASSIFIERS.builtin.source ? null : src;

843848 saveClassifier(clf);

844849 syncClfEditor();

845850 });

846851 $('clf-reset').addEventListener('click', () => {

847- clf.builtinSrc = null;

852+ if (clf.method === 'pursuit') clf.pursuitSrc = null;

853+ else clf.builtinSrc = null;

848854 saveClassifier(clf);

849855 syncClfEditor();

850856 });

src/numbl/index.tsmodified+13−3View file

@@ -1,8 +1,18 @@

1-/** The shipped classifier source and the starting point for a custom one. */

1+/** The shipped classifier sources and the starting point for a custom one. */

22 import silenceGapKnnSrc from './silence_gap_knn.m?raw';

3+import overlapPursuitSrc from './overlap_pursuit.m?raw';

34

4-export const DEFAULT_CLASSIFIER_NAME = 'silence_gap_knn';

5-export const DEFAULT_CLASSIFIER_SOURCE = silenceGapKnnSrc;

5+/** The built-in numbl classifiers, keyed by their ClassifierState method

6+ * value. silence_gap_knn is the default; overlap_pursuit adds matching

7+ * pursuit over merged regions for takes with overlapping hits. */

8+export const BUILTIN_CLASSIFIERS = {

9+ builtin: { name: 'silence_gap_knn', source: silenceGapKnnSrc },

10+ pursuit: { name: 'overlap_pursuit', source: overlapPursuitSrc },

11+} as const;

12+export type BuiltinMethod = keyof typeof BUILTIN_CLASSIFIERS;

13+

14+export const DEFAULT_CLASSIFIER_NAME = BUILTIN_CLASSIFIERS.builtin.name;

15+export const DEFAULT_CLASSIFIER_SOURCE = BUILTIN_CLASSIFIERS.builtin.source;

616

717 export const CUSTOM_CLASSIFIER_TEMPLATE = `function [times, labels] = my_classifier(x, fs, training)

818 % x the take, column vector at fs Hz (16 kHz)

src/numbl/overlap_pursuit.madded+273−0View file

@@ -0,0 +1,273 @@

1+function [times, labels] = overlap_pursuit(x, fs, training)

2+% Quiet-gap segmentation plus per-block matching pursuit for overlaps.

3+%

4+% Ported from the audio-spike-sorting benchmark. Activity above a noise

5+% threshold splits into regions; a region about one snippet long holds

6+% a single hit, classified by the spectral kNN of silence_gap_knn. A

7+% longer region holds overlapping hits and is decomposed by matching

8+% pursuit: the region start is an anchor, fitted by least squares

9+% against every training snippet (the best instance and gain, refined

10+% by coordinate-descent sweeps that add each hit back and refit it

11+% against the residual), and further hits are found by scanning the

12+% residual with FFT cross-correlation against the whole bank. A

13+% candidate is accepted while its gain and energy reduction are

14+% plausible, unless a hit of the same voice already sits within a few

15+% hundred samples, since instance mismatch of an already-fitted hit is

16+% the usual source of such candidates; that veto also merges genuine

17+% same-voice hits closer than ~16 ms into one.

18+%

19+% Adaptations for live microphone takes. On the benchmark, a region's

20+% extent is the exact union of snippet spans, which justifies a second

21+% anchor at the region end minus one snippet and a coverage rule that

22+% forces an event onto every sample. On a live take a hit's ring-out

23+% ends wherever it sinks into the noise floor, so both inferences would

24+% invent hits out of decay tails; this port keeps only the region-start

25+% anchor and finds every further hit by the threshold-gated correlation

26+% pursuit. Concretely: the segmentation threshold is raised to a robust

27+% noise floor estimate (a multiple of the median absolute sample; on

28+% digitally silent audio the median is zero and the original threshold

29+% takes over); a region up to 5 ms longer than one snippet is a single

30+% hit rather than a near-coincident pair; each merged region is padded

31+% with one snippet length of trailing context so a hit near the region

32+% end keeps a full span, with onsets confined to the region proper; and

33+% the pursuit's energy-reduction floor is a fraction of the weakest

34+% training snippet's energy rather than an absolute number, so it does

35+% not depend on the recording level.

36+nTypes = numel(training.snips);

37+L = size(training.snips{1}, 1);

38+pre = training.pre;

39+n = numel(x);

40+

41+% tuned constants

42+gmin = 0.35; % smallest plausible fitted gain for a pursuit hit

43+prox = round(0.0156 * fs); % same-voice exclusion radius, ~250 samples

44+maxDup = 6; % pursuit iterations per merged region

45+

46+% ---- training bank ------------------------------------------------------

47+nTrain = 0;

48+for k = 1:nTypes

49+ nTrain = nTrain + size(training.snips{k}, 2);

50+end

51+B = zeros(L, nTrain);

52+bt = zeros(nTrain, 1);

53+c = 0;

54+for k = 1:nTypes

55+ S = training.snips{k};

56+ for j = 1:size(S, 2)

57+ c = c + 1;

58+ B(:, c) = S(:, j);

59+ bt(c) = k;

60+ end

61+end

62+n2 = zeros(nTrain, 1);

63+for j = 1:nTrain

64+ n2(j) = sum(B(:, j) .* B(:, j));

65+end

66+% smallest energy reduction for a pursuit hit, relative to the weakest

67+% training snippet so the recording level cancels

68+redmin = 0.03 * min(n2);

69+

70+% ---- kNN training features for single-event blocks ----------------------

71+F = zeros(nTrain, 8);

72+for j = 1:nTrain

73+ F(j, :) = spectral_features(B(pre+1:end, j), fs);

74+end

75+mu = mean(F, 1);

76+sd = std(F, 0, 1) + 1e-9;

77+Fz = (F - repmat(mu, nTrain, 1)) ./ repmat(sd, nTrain, 1);

78+

79+% ---- silence segmentation ------------------------------------------------

80+eps0 = max(1e-4 * max(abs(x)), 6 * median(abs(x)));

81+gapClose = round(0.002 * fs);

82+minDur = round(0.010 * fs);

83+bs = zeros(0, 1);

84+be = zeros(0, 1);

85+i = 1;

86+while i <= n

87+ if abs(x(i)) > eps0

88+ s = i;

89+ e = i;

90+ j = i + 1;

91+ quiet = 0;

92+ while j <= n && quiet <= gapClose

93+ if abs(x(j)) > eps0

94+ e = j;

95+ quiet = 0;

96+ else

97+ quiet = quiet + 1;

98+ end

99+ j = j + 1;

100+ end

101+ if e - s + 1 >= minDur

102+ bs(end+1, 1) = s;

103+ be(end+1, 1) = e;

104+ end

105+ i = e + gapClose + 1;

106+ else

107+ i = i + 1;

108+ end

109+end

110+

111+% ---- per-block detection and classification ------------------------------

112+times = zeros(0, 1);

113+labels = zeros(0, 1);

114+for bi = 1:numel(bs)

115+ s = bs(bi);

116+ D = be(bi) - s + 1;

117+ if D <= L + round(0.005 * fs)

118+ % single hit: onset at the block start, spectral kNN label

119+ seg = x(s+pre : min(n, s+L-1));

120+ f = (spectral_features(seg, fs) - mu) ./ sd;

121+ d = sqrt(sum((Fz - repmat(f, nTrain, 1)).^2, 2));

122+ [ds, di] = sort(d);

123+ kNN = min(5, nTrain);

124+ votes = zeros(nTypes, 1);

125+ for v = 1:kNN

126+ votes(bt(di(v))) = votes(bt(di(v))) + 1 / (ds(v) + 1e-6);

127+ end

128+ [~, lab] = max(votes);

129+ times(end+1, 1) = (s - 1 + pre) / fs;

130+ labels(end+1, 1) = lab;

131+ else

132+ % pad with trailing ring context so a hit near the region end keeps

133+ % a full snippet span; onsets stay confined to the region proper

134+ blk = x(s : min(n, s + D - 1 + L));

135+ [evOn, evType] = decompose(blk, B, bt, n2, L, gmin, redmin, prox, maxDup, D);

136+ for q = 1:numel(evOn)

137+ times(end+1, 1) = (s - 1 + evOn(q) - 1 + pre) / fs;

138+ labels(end+1, 1) = evType(q);

139+ end

140+ end

141+end

142+end

143+

144+function [evOn, evType] = decompose(blk, B, bt, n2, L, gmin, redmin, prox, maxDup, nOn)

145+% Matching pursuit over one merged region. blk is the region plus up to

146+% one snippet length of trailing context; nOn is the length of the

147+% region proper, the only stretch where an onset may be placed.

148+D = numel(blk);

149+nl = D - L + 1;

150+nTrain = numel(bt);

151+

152+% conjugate template spectra for the FFT correlation scans

153+nfft = 2^nextpow2(D);

154+BF = zeros(nfft, nTrain) + 0i;

155+for j = 1:nTrain

156+ BF(:, j) = conj(fft(B(:, j), nfft));

157+end

158+

159+% events as parallel arrays, seeded with the region-start anchor

160+evOn = [1];

161+evType = zeros(1, 1);

162+evIdx = zeros(1, 1);

163+evGain = zeros(1, 1);

164+residual = blk;

165+

166+[residual, evType, evIdx, evGain] = sweep_fit(residual, evOn, evType, evIdx, evGain, B, bt, n2, L, 2);

167+

168+% correlation pursuit for further hits in the region

169+mask = zeros(nl, 1);

170+mask(1 : min(nOn, nl)) = 1;

171+for it = 1:maxDup

172+ [o, kk, idx, g, red] = scan_best(residual, BF, bt, n2, L, mask);

173+ if idx == 0 || red < redmin || g < gmin

174+ break;

175+ end

176+ clash = 0;

177+ for q = 1:numel(evOn)

178+ if evOn(q) == o || (abs(evOn(q) - o) < prox && evType(q) == kk)

179+ clash = 1;

180+ end

181+ end

182+ if clash

183+ break;

184+ end

185+ residual(o:o+L-1) = residual(o:o+L-1) - g * B(:, idx);

186+ evOn(end+1, 1) = o;

187+ evType(end+1, 1) = kk;

188+ evIdx(end+1, 1) = idx;

189+ evGain(end+1, 1) = g;

190+ [residual, evType, evIdx, evGain] = sweep_fit(residual, evOn, evType, evIdx, evGain, B, bt, n2, L, 2);

191+end

192+end

193+

194+function [residual, evType, evIdx, evGain] = sweep_fit(residual, evOn, evType, evIdx, evGain, B, bt, n2, L, nsw)

195+% Coordinate descent: add each event back and refit its type/instance/gain.

196+[~, ord] = sort(evOn);

197+for it = 1:nsw

198+ for z = 1:numel(ord)

199+ q = ord(z);

200+ o = evOn(q);

201+ if evIdx(q) > 0

202+ residual(o:o+L-1) = residual(o:o+L-1) + evGain(q) * B(:, evIdx(q));

203+ end

204+ seg = residual(o:o+L-1);

205+ bestRed = 0;

206+ bestJ = 0;

207+ bestG = 0;

208+ for j = 1:numel(bt)

209+ cj = sum(seg .* B(:, j));

210+ if cj > 0

211+ red = cj * cj / n2(j);

212+ if red > bestRed

213+ bestRed = red;

214+ bestJ = j;

215+ bestG = cj / n2(j);

216+ end

217+ end

218+ end

219+ if bestJ > 0

220+ residual(o:o+L-1) = residual(o:o+L-1) - bestG * B(:, bestJ);

221+ evType(q) = bt(bestJ);

222+ evIdx(q) = bestJ;

223+ evGain(q) = bestG;

224+ end

225+ end

226+end

227+end

228+

229+function [o, kk, idx, g, red] = scan_best(residual, BF, bt, n2, L, mask)

230+% Best (onset, template) over masked lags by FFT cross-correlation;

231+% reduction is the residual energy removed by the least-squares gain.

232+nfft = size(BF, 1);

233+nl = numel(mask);

234+R = fft(residual, nfft);

235+o = 0; kk = 0; idx = 0; g = 0; red = 0;

236+for j = 1:numel(bt)

237+ cf = real(ifft(R .* BF(:, j)));

238+ cl = cf(1:nl);

239+ rv = (cl .* cl) / n2(j);

240+ rv = rv .* (cl > 0) .* mask;

241+ [m, mi] = max(rv);

242+ if m > red

243+ red = m;

244+ o = mi;

245+ kk = bt(j);

246+ idx = j;

247+ g = cl(mi) / n2(j);

248+ end

249+end

250+end

251+

252+function f = spectral_features(seg, fs)

253+% Seven log band-energy ratios (mean removed, level cancels) and the

254+% log spectral centroid; the extra low bands split the two bass types.

255+nfft = 2048;

256+m = min(numel(seg), nfft);

257+s2 = seg(1:m);

258+win = 0.5 - 0.5 * cos(2 * pi * (0:m-1)' / max(m - 1, 1));

259+S = fft(s2 .* win, nfft);

260+p = abs(S(1:nfft/2+1)).^2;

261+hz = fs / nfft;

262+edges = [0 100 200 400 800 1600 3200 6400];

263+b = zeros(1, 7);

264+for q = 1:7

265+ k0 = max(2, floor(edges(q) / hz) + 1);

266+ k1 = min(nfft/2 + 1, floor(edges(q+1) / hz));

267+ b(q) = log10(sum(p(k0:k1)) + 1e-12);

268+end

269+b = b - mean(b);

270+kv = (0:numel(p)-1)' * hz;

271+cen = sum(p .* kv) / (sum(p) + 1e-12);

272+f = [b, log10(cen + 1)];

273+end

src/store.tsmodified+9−4View file

@@ -103,17 +103,22 @@ export function saveUi(ui: UiState): void {

103103 localStorage.setItem(KEY_UI, JSON.stringify(ui));

104104 }

105105

106-/** How takes are classified: the shipped numbl algorithm (default, its

107- * source editable), a custom numbl algorithm, or the original TS

106+/** How takes are classified: one of the shipped numbl algorithms

107+ * (silence_gap_knn under 'builtin', overlap_pursuit under 'pursuit',

108+ * both source-editable), a custom numbl algorithm, or the original TS

108109 * feature-kNN. Null sources mean "the shipped default" / "the template". */

109110 export interface ClassifierState {

110- method: 'builtin' | 'custom' | 'legacy';

111+ method: 'builtin' | 'pursuit' | 'custom' | 'legacy';

111112 builtinSrc: string | null;

113+ pursuitSrc: string | null;

112114 customSrc: string | null;

113115 }

114116 const KEY_CLF = 'instrawave-compose/classifier';

115117 export function loadClassifier(): ClassifierState {

116- return read(KEY_CLF) ?? { method: 'builtin', builtinSrc: null, customSrc: null };

118+ // pursuitSrc was added later; fill it in for state saved before that

119+ const c = read<Omit<ClassifierState, 'pursuitSrc'> & { pursuitSrc?: string | null }>(KEY_CLF)

120+ ?? { method: 'builtin' as const, builtinSrc: null, customSrc: null };

121+ return { pursuitSrc: null, ...c };

117122 }

118123 export function saveClassifier(c: ClassifierState): void {

119124 localStorage.setItem(KEY_CLF, JSON.stringify(c));