/
/
/
1"""Pure numpy DBN postprocessor for beat/downbeat tracking."""
2
3from __future__ import annotations
4
5from typing import TYPE_CHECKING, Any
6
7import numpy as np
8
9if TYPE_CHECKING:
10 from numpy.typing import NDArray
11
12
13class DBNDownBeatTracker:
14 """
15 Pure numpy DBN postprocessor for beat and downbeat tracking.
16
17 Drop-in replacement for madmom's DBNDownBeatTrackingProcessor.
18
19 :param beats_per_bar: List of time signatures to try (e.g. [3, 4]).
20 :param min_bpm: Minimum tempo in BPM.
21 :param max_bpm: Maximum tempo in BPM.
22 :param fps: Frames per second of the input activations.
23 :param transition_lambda: Tempo change penalty (higher = more stable).
24 :param observation_lambda: Beat subdivision granularity.
25 :param threshold: Minimum activation to consider.
26 :param correct: Whether to snap beats to activation peaks.
27 """
28
29 def __init__(
30 self,
31 beats_per_bar: list[int] | None = None,
32 min_bpm: float = 55.0,
33 max_bpm: float = 215.0,
34 fps: int = 50,
35 transition_lambda: float = 100.0,
36 observation_lambda: int = 16,
37 threshold: float = 0.05,
38 correct: bool = True,
39 ) -> None:
40 """Initialize the DBN tracker with the given parameters."""
41 if beats_per_bar is None:
42 beats_per_bar = [3, 4]
43 self.fps = fps
44 self.threshold = threshold
45 self.correct = correct
46 self.observation_lambda = observation_lambda
47
48 self._hmms: list[dict[str, Any]] = []
49 min_interval = round(60.0 * fps / max_bpm)
50 max_interval = round(60.0 * fps / min_bpm)
51
52 for num_beats in beats_per_bar:
53 positions, intervals = self._build_bar_state_space(
54 num_beats, min_interval, max_interval
55 )
56 tm_states, tm_pointers, tm_log_probs = self._build_transition_model(
57 positions, intervals, num_beats, transition_lambda
58 )
59 om_pointers = self._build_observation_model(positions, observation_lambda)
60 self._hmms.append(
61 {
62 "num_beats": num_beats,
63 "positions": positions,
64 "om_pointers": om_pointers,
65 "tm_states": tm_states,
66 "tm_pointers": tm_pointers,
67 "tm_log_probs": tm_log_probs,
68 }
69 )
70
71 def __call__(self, activations: NDArray[np.float64]) -> tuple[NDArray[np.float64], int]:
72 """
73 Run DBN decoding on beat/downbeat activations.
74
75 :param activations: Shape (T, 2), columns [beat_act, downbeat_act].
76 Values should be probabilities in (0, 1).
77 :return: Tuple of (beats, num_beats).
78 beats: Shape (M, 2) array of [time_seconds, beat_position].
79 beat_position is 1 for downbeats, 2..num_beats for other beats.
80 num_beats: Beats per bar of the winning meter hypothesis (0 if no beats found).
81 """
82 # Threshold: trim leading/trailing silence
83 first = 0
84 if self.threshold:
85 above = np.nonzero(activations.max(axis=1) >= self.threshold)[0]
86 if len(above) > 0:
87 first = int(above[0])
88 last = int(above[-1]) + 1
89 activations = activations[first:last]
90
91 log_dens = self._compute_log_densities(activations, self.observation_lambda)
92
93 # Run Viterbi for each meter hypothesis, pick best
94 best_path = None
95 best_log_prob = -np.inf
96 best_hmm = self._hmms[0]
97
98 for hmm in self._hmms:
99 path, log_prob = self._viterbi(
100 log_dens,
101 hmm["om_pointers"],
102 hmm["tm_states"],
103 hmm["tm_pointers"],
104 hmm["tm_log_probs"],
105 )
106 if log_prob > best_log_prob:
107 best_log_prob = log_prob
108 best_path = path
109 best_hmm = hmm
110
111 assert best_path is not None
112 positions = best_hmm["positions"][best_path]
113 beat_numbers = positions.astype(int) + 1
114
115 if self.correct:
116 beats = self._correct_beats(best_path, best_hmm, activations, beat_numbers)
117 else:
118 beats = np.nonzero(np.diff(beat_numbers))[0] + 1
119
120 if len(beats) == 0:
121 return np.empty((0, 2), dtype=np.float64), 0
122
123 beat_times = (beats + first) / self.fps
124 beat_positions = beat_numbers[beats]
125
126 return np.column_stack([beat_times, beat_positions]), int(best_hmm["num_beats"])
127
128 @staticmethod
129 def _build_bar_state_space(
130 num_beats: int,
131 min_interval: int,
132 max_interval: int,
133 ) -> tuple[NDArray[np.float32], NDArray[np.int32]]:
134 """
135 Build a bar-pointer state space for the given tempo range.
136
137 Each state represents a position within a bar at a specific tempo.
138 For each tempo interval i (in frames), there are i discrete positions
139 per beat. A bar with B beats has B x i states for interval i.
140
141 :param num_beats: Number of beats per bar (e.g. 4 for 4/4 time).
142 :param min_interval: Minimum beat interval in frames.
143 :param max_interval: Maximum beat interval in frames.
144 :return: Tuple of (state_positions, state_intervals).
145 state_positions: float64 array, position in bar [0, num_beats).
146 state_intervals: int32 array, tempo interval for each state.
147 """
148 intervals_range = np.arange(min_interval, max_interval + 1)
149 counts = intervals_range
150 one_beat_intervals = np.repeat(intervals_range, counts)
151 offsets = np.concatenate([np.arange(i, dtype=np.float32) / i for i in intervals_range])
152
153 all_positions = np.concatenate([beat + offsets for beat in range(num_beats)]).astype(
154 np.float32
155 )
156 all_intervals = np.tile(one_beat_intervals, num_beats)
157
158 return all_positions, all_intervals.astype(np.int32)
159
160 @staticmethod
161 def _build_transition_model(
162 positions: NDArray[np.float32],
163 intervals: NDArray[np.int32],
164 num_beats: int,
165 transition_lambda: float,
166 ) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:
167 """
168 Build a sparse CSR transition model for the bar state space.
169
170 Within a beat, each state transitions deterministically to the next
171 (constant tempo). At beat boundaries, transitions follow an exponential
172 tempo-change distribution controlled by transition_lambda.
173
174 :param positions: State positions from _build_bar_state_space.
175 :param intervals: State intervals from _build_bar_state_space.
176 :param num_beats: Number of beats per bar.
177 :param transition_lambda: Penalty for tempo changes (higher = less change).
178 :return: Tuple of (tm_states, tm_pointers, tm_log_probs) in CSR format.
179 For state s, predecessors are tm_states[tm_pointers[s]:tm_pointers[s+1]]
180 with log probabilities tm_log_probs[tm_pointers[s]:tm_pointers[s+1]].
181 """
182 num_states = len(positions)
183
184 # Identify first and last states per (beat, interval)
185 unique_intervals = np.arange(intervals.min(), intervals.max() + 1)
186 num_intervals = len(unique_intervals)
187 states_per_beat = num_states // num_beats
188
189 # Build first/last state indices per beat
190 first_states_per_beat: list[NDArray[np.int32]] = []
191 last_states_per_beat: list[NDArray[np.int32]] = []
192 for beat in range(num_beats):
193 offset = beat * states_per_beat
194 firsts = []
195 lasts = []
196 idx = 0
197 for interval in unique_intervals:
198 firsts.append(offset + idx)
199 idx += interval
200 lasts.append(offset + idx - 1)
201 first_states_per_beat.append(np.array(firsts, dtype=np.int32))
202 last_states_per_beat.append(np.array(lasts, dtype=np.int32))
203
204 # Compute exponential transition probabilities between tempi
205 from_intervals = unique_intervals.astype(np.float64)
206 to_intervals = unique_intervals.astype(np.float64)
207 ratio = to_intervals[np.newaxis, :] / from_intervals[:, np.newaxis]
208 trans_prob = np.exp(-transition_lambda * np.abs(ratio - 1.0))
209 # Normalize each row
210 row_sums = trans_prob.sum(axis=1, keepdims=True)
211 trans_prob = trans_prob / row_sums
212 trans_log_prob = np.log(trans_prob).astype(np.float32)
213
214 # Vectorized CSR construction: separate within-beat from boundary states
215 frac = positions - np.floor(positions)
216 within_mask = frac > 0
217 within_states = np.nonzero(within_mask)[0]
218 boundary_states = np.nonzero(~within_mask)[0]
219
220 # Within-beat: predecessor is state - 1, log_prob = 0
221 sources_w = within_states - 1
222 dests_w = within_states
223 logprobs_w = np.zeros(len(within_states), dtype=np.float32)
224
225 # Boundary states: predecessors from last states of previous beat
226 sources_b: list[int] = []
227 dests_b: list[int] = []
228 logprobs_b: list[float] = []
229 min_interval_val = int(intervals.min())
230
231 for state in boundary_states:
232 beat = int(positions[state])
233 prev_beat = (beat - 1) % num_beats
234 prev_lasts = last_states_per_beat[prev_beat]
235 cur_interval_idx = int(intervals[state]) - min_interval_val
236
237 for from_idx in range(num_intervals):
238 lp = trans_log_prob[from_idx, cur_interval_idx]
239 if lp > -50:
240 sources_b.append(int(prev_lasts[from_idx]))
241 dests_b.append(int(state))
242 logprobs_b.append(float(lp))
243
244 sources_all = np.concatenate([sources_w, np.array(sources_b, dtype=np.int32)]).astype(
245 np.int32
246 )
247 dests_all = np.concatenate([dests_w, np.array(dests_b, dtype=np.int32)]).astype(np.int32)
248 logprobs_all = np.concatenate([logprobs_w, np.array(logprobs_b, dtype=np.float32)]).astype(
249 np.float32
250 )
251
252 return DBNDownBeatTracker._csr_from_arrays(sources_all, dests_all, logprobs_all, num_states)
253
254 @staticmethod
255 def _csr_from_arrays(
256 sources_arr: NDArray[np.int32],
257 dests_arr: NDArray[np.int32],
258 log_probs_arr: NDArray[np.float32],
259 num_states: int,
260 ) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:
261 """Convert transition arrays to CSR format indexed by destination."""
262 order = np.argsort(dests_arr, kind="stable")
263 sources_arr = sources_arr[order]
264 dests_arr = dests_arr[order]
265 log_probs_arr = log_probs_arr[order]
266
267 tm_pointers = np.zeros(num_states + 1, dtype=np.int32)
268 tm_pointers[1:] = np.cumsum(np.bincount(dests_arr, minlength=num_states))
269
270 return sources_arr, tm_pointers, log_probs_arr
271
272 @staticmethod
273 def _build_observation_model(
274 positions: NDArray[np.float32],
275 observation_lambda: int = 16,
276 ) -> NDArray[np.int32]:
277 """
278 Build observation model pointers mapping states to observation classes.
279
280 Class 0 = no-beat, class 1 = beat, class 2 = downbeat.
281
282 :param positions: State positions from _build_bar_state_space.
283 :param observation_lambda: Beat subdivision granularity.
284 :return: Array of observation class indices, one per state.
285 """
286 border = 1.0 / observation_lambda
287 pointers = np.zeros(len(positions), dtype=np.int32)
288 # Beat states: fractional position within beat < border
289 frac = positions % 1.0
290 pointers[frac < border] = 1
291 # Downbeat states: absolute position < border (first beat of bar)
292 pointers[positions < border] = 2
293 return pointers
294
295 @staticmethod
296 def _compute_log_densities(
297 activations: NDArray[np.float64],
298 observation_lambda: int = 16,
299 ) -> NDArray[np.float32]:
300 """
301 Compute log observation densities from beat/downbeat activations.
302
303 :param activations: Shape (T, 2), columns [beat_act, downbeat_act].
304 :param observation_lambda: Beat subdivision granularity.
305 :return: Shape (T, 3) log densities for [no-beat, beat, downbeat].
306 """
307 act = activations.astype(np.float32, copy=False)
308 beat_act = act[:, 0]
309 downbeat_act = act[:, 1]
310 no_beat_act = np.maximum(1.0 - beat_act - downbeat_act, 1e-7)
311
312 log_dens = np.empty((len(activations), 3), dtype=np.float32)
313 log_dens[:, 0] = np.log(no_beat_act / (observation_lambda - 1))
314 log_dens[:, 1] = np.log(np.maximum(beat_act, 1e-7))
315 log_dens[:, 2] = np.log(np.maximum(downbeat_act, 1e-7))
316 return log_dens
317
318 @staticmethod
319 def _viterbi(
320 log_densities: NDArray[np.float32],
321 om_pointers: NDArray[np.int32],
322 tm_states: NDArray[np.int32],
323 tm_pointers: NDArray[np.int32],
324 tm_log_probs: NDArray[np.float32],
325 ) -> tuple[NDArray[np.int32], float]:
326 """
327 Run Viterbi decoding on the HMM.
328
329 :param log_densities: Shape (T, 3), per-frame log observation densities.
330 :param om_pointers: Shape (S,), observation class per state.
331 :param tm_states: CSR source states for transitions.
332 :param tm_pointers: CSR pointer array for transitions.
333 :param tm_log_probs: CSR log probabilities for transitions.
334 :return: Tuple of (path, log_probability).
335 path: shape (T,) int32 array of state indices.
336 log_probability: float, log prob of best path.
337 """
338 num_frames = len(log_densities)
339 num_states = len(om_pointers)
340
341 # Classify states: single-predecessor (within-beat) vs multi-predecessor (beat boundary)
342 num_preds = np.diff(tm_pointers)
343 single_mask = num_preds == 1
344 multi_mask = ~single_mask
345 single_states = np.nonzero(single_mask)[0]
346 multi_states = np.nonzero(multi_mask)[0]
347
348 # For single-predecessor states, precompute the single source
349 single_sources = np.empty(num_states, dtype=np.int32)
350 single_sources[single_states] = tm_states[tm_pointers[single_states]]
351
352 # For multi-predecessor states, build padded arrays for vectorized max
353 if len(multi_states) > 0:
354 max_preds = num_preds[multi_states].max()
355 multi_source_pad = np.zeros((len(multi_states), max_preds), dtype=np.int32)
356 multi_logprob_pad = np.full((len(multi_states), max_preds), -np.inf, dtype=np.float32)
357 for i, s in enumerate(multi_states):
358 start = tm_pointers[s]
359 end = tm_pointers[s + 1]
360 n = end - start
361 multi_source_pad[i, :n] = tm_states[start:end]
362 multi_logprob_pad[i, :n] = tm_log_probs[start:end]
363
364 # Initialize: uniform over all states, ping-pong buffers
365 buf_a = np.empty(num_states, dtype=np.float32)
366 buf_b = np.empty(num_states, dtype=np.float32)
367 prev_v = buf_a
368 prev_v[:] = np.float32(-np.log(num_states))
369 bt = np.empty((num_frames, len(multi_states)), dtype=np.int32)
370
371 # Pre-compute constant index arrays used every frame
372 single_src_idx = single_sources[single_states]
373 has_multi = len(multi_states) > 0
374 if has_multi:
375 arange_multi = np.arange(len(multi_states))
376
377 for t in range(num_frames):
378 cur_v = buf_b if prev_v is buf_a else buf_a
379 cur_v[:] = -np.inf
380 obs = log_densities[t, om_pointers]
381
382 # Single-predecessor states: direct assignment (log_prob = 0)
383 cur_v[single_states] = prev_v[single_src_idx] + obs[single_states]
384
385 # Multi-predecessor states: max over predecessors
386 if has_multi:
387 scores = prev_v[multi_source_pad] + multi_logprob_pad
388 best_idx = scores.argmax(axis=1)
389 cur_v[multi_states] = scores[arange_multi, best_idx] + obs[multi_states]
390 bt[t] = best_idx
391
392 prev_v = cur_v
393
394 # Backtrack â O(1) lookup instead of searchsorted
395 multi_lookup = np.full(num_states, -1, dtype=np.int32)
396 multi_lookup[multi_states] = np.arange(len(multi_states), dtype=np.int32)
397
398 path = np.empty(num_frames, dtype=np.int32)
399 best_state = int(np.argmax(prev_v))
400 log_prob = float(prev_v[best_state])
401
402 for t in range(num_frames - 1, -1, -1):
403 path[t] = best_state
404 mi = multi_lookup[best_state]
405 if mi >= 0:
406 best_state = int(multi_source_pad[mi, bt[t, mi]])
407 else:
408 best_state = int(single_sources[best_state])
409
410 return path, log_prob
411
412 def _correct_beats(
413 self,
414 path: NDArray[np.int32],
415 hmm: dict[str, Any],
416 activations: NDArray[np.float64],
417 beat_numbers: NDArray[np.int64],
418 ) -> NDArray[np.int64]:
419 """Snap detected beats to activation peaks within beat regions."""
420 om_pointers = hmm["om_pointers"]
421 beat_range = om_pointers[path] >= 1
422
423 transitions = np.diff(beat_range.astype(np.int32))
424 idx = np.nonzero(transitions)[0] + 1
425
426 if beat_range[0]:
427 idx = np.concatenate([[0], idx])
428 if beat_range[-1]:
429 idx = np.concatenate([idx, [len(beat_range)]])
430
431 if len(idx) % 2 != 0:
432 idx = idx[:-1]
433
434 beats = []
435 for i in range(0, len(idx), 2):
436 left, right = idx[i], idx[i + 1]
437 # Find peak activation in region (sum both columns)
438 region_act = activations[left:right].sum(axis=1)
439 peak = int(np.argmax(region_act)) + left
440 beats.append(peak)
441
442 return np.array(beats, dtype=np.int64)
443