Select contiguous beat runs via diagonal Kadane search on pair similarity matrix

app.py

@@ -311,6 +311,61 @@ def _intervals_from_mask(mask: np.ndarray, window: BeatWindow, max_end: float) -
     return _merge_intervals(intervals)
 
 
+def _find_contiguous_beats(pair_probs: np.ndarray, min_beats: int = 2) -> tuple[np.ndarray, np.ndarray]:
+    """Find the best contiguous diagonal run in the beat similarity matrix.
+
+    Searches every diagonal offset (track_beat - source_beat) and uses
+    Kadane's algorithm to find the highest-scoring contiguous segment along
+    each diagonal. Returns boolean masks over track and source beats.
+    """
+    n_track, n_source = pair_probs.shape
+    best_score = -np.inf
+    best_track_mask = np.zeros(n_track, dtype=bool)
+    best_source_mask = np.zeros(n_source, dtype=bool)
+
+    for d in range(-(n_source - 1), n_track):
+        # diagonal: track[i], source[i - d] for valid i
+        i0 = max(0, d)
+        j0 = max(0, -d)
+        length = min(n_track - i0, n_source - j0)
+        if length < min_beats:
+            continue
+
+        diag = pair_probs[i0:i0 + length, j0:j0 + length].diagonal()
+
+        # Kadane's max-subarray on the diagonal values
+        curr_sum = 0.0
+        curr_start = 0
+        best_sum = -np.inf
+        seg_start = seg_end = 0
+
+        for k, val in enumerate(diag):
+            curr_sum += val
+            if curr_sum > best_sum:
+                best_sum = curr_sum
+                seg_start = curr_start
+                seg_end = k
+            if curr_sum < 0:
+                curr_sum = 0.0
+                curr_start = k + 1
+
+        seg_len = seg_end - seg_start + 1
+        if seg_len < min_beats:
+            continue
+        avg_score = best_sum / seg_len
+
+        if avg_score > best_score:
+            best_score = avg_score
+            track_mask = np.zeros(n_track, dtype=bool)
+            source_mask = np.zeros(n_source, dtype=bool)
+            track_mask[i0 + seg_start: i0 + seg_end + 1] = True
+            source_mask[j0 + seg_start: j0 + seg_end + 1] = True
+            best_track_mask = track_mask
+            best_source_mask = source_mask
+
+    return best_track_mask, best_source_mask
+
+
 def _localize_match(
     model,
     track_mel: torch.Tensor,
@@ -332,15 +387,17 @@ def _localize_match(
     with torch.inference_mode():
         pair_probs = torch.sigmoid(model.pair_mask_head(track_mel, source_mel))[0].detach().cpu().numpy()
 
-    selected = pair_probs >= float(threshold)
-    if not selected.any():
+    track_mask, source_mask = _find_contiguous_beats(pair_probs, min_beats=2)
+
+    # Fall back to top-k individual beats if no contiguous run was found
+    if not track_mask.any():
         top_k = min(6, pair_probs.size)
         flat = np.argpartition(pair_probs.reshape(-1), -top_k)[-top_k:]
         selected = np.zeros_like(pair_probs, dtype=bool)
         selected.reshape(-1)[flat] = True
+        track_mask = selected.any(axis=1)
+        source_mask = selected.any(axis=0)
 
-    track_mask = selected.any(axis=1)
-    source_mask = selected.any(axis=0)
     track_regions = _intervals_from_mask(
         track_mask,
         track_window,