Use classifier head for match verdict, show proposed masks on no-match

app.py

@@ -18,7 +18,7 @@ from huggingface_hub import hf_hub_download
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
-from model import CNNSampleDetector, SSLAMSampleDetector, SampleDetector
+from model import CNNSampleDetector, SSLAMSampleDetector, SampleDetector, pair_summary_features
 
 
 SAMPLE_RATE = int(os.environ.get("APP_SAMPLE_RATE", "16000"))
@@ -282,24 +282,33 @@ def _encode(model, mels: torch.Tensor, batch_size: int) -> torch.Tensor:
     return torch.cat(embs, dim=0)
 
 
-def _score_pairs(model, track_mels: torch.Tensor, source_mels: torch.Tensor, batch_size: int) -> torch.Tensor:
+def _score_pairs(
+    model,
+    track_mels: torch.Tensor,
+    source_mels: torch.Tensor,
+    batch_size: int,
+    pair_head_loaded: bool,
+) -> torch.Tensor:
     track_emb = _encode(model, track_mels, batch_size)
     source_emb = _encode(model, source_mels, batch_size)
     n_track, n_source = track_emb.shape[0], source_emb.shape[0]
-    scores = []
+    scores = torch.zeros(n_track, n_source, device=track_emb.device)
 
-    pair_indices = [(i, j) for i in range(n_track) for j in range(n_source)]
-    for start in range(0, len(pair_indices), batch_size):
-        chunk = pair_indices[start:start + batch_size]
-        ti = torch.tensor([p[0] for p in chunk], device=track_emb.device)
-        sj = torch.tensor([p[1] for p in chunk], device=track_emb.device)
-        t = track_emb.index_select(0, ti)
-        s = source_emb.index_select(0, sj)
-        combined = torch.cat([t, s, torch.abs(t - s), t * s], dim=-1)
-        logits = model.head(combined)
-        scores.append(torch.softmax(logits, dim=-1)[:, 1])
+    for i in range(n_track):
+        for j in range(n_source):
+            t = track_emb[i:i + 1]
+            s = source_emb[j:j + 1]
+            if pair_head_loaded:
+                pair_feat = pair_summary_features(
+                    model.pair_mask_head(track_mels[i:i + 1], source_mels[j:j + 1])
+                )
+                combined = torch.cat([t, s, torch.abs(t - s), t * s, pair_feat], dim=-1)
+            else:
+                combined = torch.cat([t, s, torch.abs(t - s), t * s], dim=-1)
+            logits = model.head(combined)
+            scores[i, j] = torch.softmax(logits, dim=-1)[0, 1]
 
-    return torch.cat(scores).reshape(n_track, n_source)
+    return scores
 
 
 def _intervals_from_mask(mask: np.ndarray, window: BeatWindow, max_end: float) -> list[tuple[float, float]]:
@@ -390,10 +399,10 @@ def _draw_mel(ax, clip: AudioClip, regions: list[tuple[float, float]], color: st
     ax.set_ylabel("Frequency (Hz)")
     ax.set_xlim(t_start, t_end)
 
-    if matched and regions:
+    if regions:
         for start, end in regions:
-            ax.axvspan(start, end, color=color, alpha=0.38, linewidth=0)
-    elif not matched:
+            ax.axvspan(start, end, color=color, alpha=0.38 if matched else 0.22, linewidth=0)
+    if not matched:
         ax.text(
             0.5, 0.5, "No Match",
             transform=ax.transAxes,
@@ -495,7 +504,7 @@ def verify(
     source_mels = torch.stack([_to_mel(w.waveform, source_bpm, args) for w in source_windows]).to(device)
 
     with torch.inference_mode():
-        score_matrix = _score_pairs(model, track_mels, source_mels, batch_size)
+        score_matrix = _score_pairs(model, track_mels, source_mels, batch_size, loaded["pair_head_loaded"])
     best_flat = int(torch.argmax(score_matrix).item())
     best_track = best_flat // score_matrix.shape[1]
     best_source = best_flat % score_matrix.shape[1]
@@ -514,19 +523,16 @@ def verify(
         loaded["pair_head_loaded"],
     )
 
-    highlight_track = track_regions if matched else []
-    highlight_source = source_regions if matched else []
-
-    wfig = _plot_waveforms(track_clip, source_clip, highlight_track, highlight_source, best_score, matched)
-    mfig = _plot_mels(track_clip, source_clip, highlight_track, highlight_source, matched)
+    wfig = _plot_waveforms(track_clip, source_clip, track_regions, source_regions, best_score, matched)
+    mfig = _plot_mels(track_clip, source_clip, track_regions, source_regions, matched)
 
-    verdict = "Likely match" if matched else "No confident match"
+    verdict = "Likely match" if matched else "No match"
     details = [
         f"**{verdict}**",
         f"Score: `{best_score:.3f}` with threshold `{float(match_threshold):.2f}`.",
         f"Estimated BPM: track `{track_bpm:.1f}`, source `{source_bpm:.1f}`.",
-        f"Highlighted track section(s): {_format_intervals(highlight_track)}.",
-        f"Highlighted source section(s): {_format_intervals(highlight_source)}.",
+        f"{'Matched' if matched else 'Proposed'} track section(s): {_format_intervals(track_regions)}.",
+        f"{'Matched' if matched else 'Proposed'} source section(s): {_format_intervals(source_regions)}.",
         f"Model: `{args.get('backbone', 'ast')}` checkpoint epoch `{loaded['epoch']}` on `{device}`.",
     ]
     if note:

model.py

@@ -14,6 +14,7 @@ AST_FREQ_DIM = 128
 SSLAM_HF_REPO = os.environ["SSLAM_MODEL"]
 SSLAM_TIME_DIM = 1024
 SSLAM_FREQ_DIM = 128
+PAIR_SUMMARY_DIM = 8
 
 
 class ASTEncoder(nn.Module):
@@ -61,11 +62,21 @@ class ASTEncoder(nn.Module):
 class PairMaskHead(nn.Module):
     """Beat-by-beat pair matching head over two mel spectrograms."""
 
-    def __init__(self, beats_per_window: int, n_mels: int, beat_dim: int = 64):
+    def __init__(self, beats_per_window: int, n_mels: int, beat_dim: int = 64, frames_per_beat: int = 8):
         super().__init__()
-        self.pool = nn.AdaptiveAvgPool2d((beats_per_window, n_mels))
-        self.beat_proj = nn.Sequential(
-            nn.Linear(n_mels, beat_dim),
+        self.beats_per_window = beats_per_window
+        self.frames_per_beat = frames_per_beat
+        self.pool = nn.AdaptiveAvgPool2d((beats_per_window * frames_per_beat, n_mels))
+        self.patch_encoder = nn.Sequential(
+            nn.Conv2d(1, 16, kernel_size=(3, 5), padding=(1, 2), bias=False),
+            nn.GroupNorm(4, 16),
+            nn.GELU(),
+            nn.Conv2d(16, 32, kernel_size=(3, 5), stride=(2, 2), padding=(1, 2), bias=False),
+            nn.GroupNorm(8, 32),
+            nn.GELU(),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+            nn.Linear(32, beat_dim),
             nn.GELU(),
             nn.Linear(beat_dim, beat_dim),
         )
@@ -73,9 +84,14 @@ class PairMaskHead(nn.Module):
         self.bias = nn.Parameter(torch.zeros(()))
 
     def _beats(self, mel: torch.Tensor) -> torch.Tensor:
-        # mel: [B, 1, T, F] -> [B, beats, F] -> [B, beats, beat_dim]
-        x = self.pool(mel).squeeze(1)
-        return torch.nn.functional.normalize(self.beat_proj(x), dim=-1)
+        # mel: [B, 1, T, F] -> [B * beats, 1, frames_per_beat, F]
+        bsz = mel.shape[0]
+        x = self.pool(mel)
+        x = x.view(bsz, 1, self.beats_per_window, self.frames_per_beat, x.shape[-1])
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+        x = x.view(bsz * self.beats_per_window, 1, self.frames_per_beat, x.shape[-1])
+        x = self.patch_encoder(x).view(bsz, self.beats_per_window, -1)
+        return torch.nn.functional.normalize(x, dim=-1)
 
     def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
         t = self._beats(track_mel)
@@ -83,6 +99,29 @@ class PairMaskHead(nn.Module):
         return torch.einsum("bih,bjh->bij", t, o) * self.logit_scale.exp() + self.bias
 
 
+def pair_summary_features(pair_logits: torch.Tensor) -> torch.Tensor:
+    probs = torch.sigmoid(pair_logits)
+    flat = probs.flatten(1)
+    row_max = probs.max(dim=2).values
+    col_max = probs.max(dim=1).values
+    diag = torch.diagonal(probs, dim1=1, dim2=2)
+    top_k = min(8, flat.shape[1])
+    topk_mean = flat.topk(top_k, dim=1).values.mean(dim=1)
+    return torch.stack(
+        [
+            flat.mean(dim=1),
+            flat.max(dim=1).values,
+            flat.std(dim=1, unbiased=False),
+            topk_mean,
+            row_max.mean(dim=1),
+            row_max.max(dim=1).values,
+            col_max.mean(dim=1),
+            diag.mean(dim=1),
+        ],
+        dim=-1,
+    )
+
+
 class SampleDetector(nn.Module):
     """Siamese AST encoder + interaction head for binary sample detection."""
 
@@ -97,9 +136,10 @@ class SampleDetector(nn.Module):
         super().__init__()
         self.encoder = ASTEncoder(model_name, freeze=freeze_encoder)
         H = self.encoder.ast.config.hidden_size
+        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
         self.head = nn.Sequential(
-            nn.LayerNorm(4 * H),
-            nn.Linear(4 * H, 512),
+            nn.LayerNorm(4 * H + PAIR_SUMMARY_DIM),
+            nn.Linear(4 * H + PAIR_SUMMARY_DIM, 512),
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(512, 128),
@@ -107,7 +147,6 @@ class SampleDetector(nn.Module):
             nn.Dropout(dropout),
             nn.Linear(128, 2),
         )
-        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
 
     def unfreeze_encoder(self, n_blocks: int = 2):
         self.encoder.unfreeze_last_n(n_blocks)
@@ -115,8 +154,9 @@ class SampleDetector(nn.Module):
     def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
         t = self.encoder(track_mel)
         o = self.encoder(orig_mel)
+        pair_features = pair_summary_features(self.pair_mask_head(track_mel, orig_mel))
         # print(f"embeddings: t={t.shape}, o={o.shape}")
-        combined = torch.cat([t, o, torch.abs(t - o), t * o], dim=-1)
+        combined = torch.cat([t, o, torch.abs(t - o), t * o, pair_features], dim=-1)
         # print(f"combined shape: {combined.shape}")
         return self.head(combined)
 
@@ -160,9 +200,10 @@ class CNNSampleDetector(nn.Module):
     def __init__(self, embed_dim: int = 256, dropout: float = 0.3, beats_per_window: int = 16, n_mels: int = 128):
         super().__init__()
         self.encoder = CNNEncoder(embed_dim)
+        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
         self.head = nn.Sequential(
-            nn.LayerNorm(4 * embed_dim),
-            nn.Linear(4 * embed_dim, 256),
+            nn.LayerNorm(4 * embed_dim + PAIR_SUMMARY_DIM),
+            nn.Linear(4 * embed_dim + PAIR_SUMMARY_DIM, 256),
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(256, 64),
@@ -170,12 +211,12 @@ class CNNSampleDetector(nn.Module):
             nn.Dropout(dropout),
             nn.Linear(64, 2),
         )
-        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
 
     def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
         t = self.encoder(track_mel)
         o = self.encoder(orig_mel)
-        combined = torch.cat([t, o, torch.abs(t - o), t * o], dim=-1)
+        pair_features = pair_summary_features(self.pair_mask_head(track_mel, orig_mel))
+        combined = torch.cat([t, o, torch.abs(t - o), t * o, pair_features], dim=-1)
         return self.head(combined)
 
 
@@ -257,9 +298,10 @@ class SSLAMSampleDetector(nn.Module):
         super().__init__()
         self.encoder = SSLAMEncoder(freeze=freeze_encoder)
         H = self.encoder.hidden_size
+        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
         self.head = nn.Sequential(
-            nn.LayerNorm(4 * H),
-            nn.Linear(4 * H, 512),
+            nn.LayerNorm(4 * H + PAIR_SUMMARY_DIM),
+            nn.Linear(4 * H + PAIR_SUMMARY_DIM, 512),
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(512, 128),
@@ -267,7 +309,6 @@ class SSLAMSampleDetector(nn.Module):
             nn.Dropout(dropout),
             nn.Linear(128, 2),
         )
-        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
 
     def unfreeze_encoder(self, n_blocks: int):
         self.encoder.unfreeze_last_n(n_blocks)
@@ -275,7 +316,8 @@ class SSLAMSampleDetector(nn.Module):
     def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
         t = self.encoder(track_mel)
         o = self.encoder(orig_mel)
-        combined = torch.cat([t, o, torch.abs(t - o), t * o], dim=-1)
+        pair_features = pair_summary_features(self.pair_mask_head(track_mel, orig_mel))
+        combined = torch.cat([t, o, torch.abs(t - o), t * o, pair_features], dim=-1)
         return self.head(combined)