first commit

models/matcher.py

@@ -0,0 +1,59 @@
+"""Hungarian bipartite matching between predicted queries and GT lanes.
+
+Replaces the paper's vague "distance-based positive/negative" anchor rule. This is
+load-bearing for training (not an optional improvement): DETR-style architectures
+with fixed slot assignment typically fail to train due to the permutation problem.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+from models.head import eval_curve
+
+
+@torch.no_grad()
+def hungarian_match(
+    cls_logits: torch.Tensor,       # (B, N, 2)
+    curve_coeffs: torch.Tensor,     # (B, N, 4)
+    target_xs: torch.Tensor,        # (B, N, S)
+    target_valid_mask: torch.Tensor,  # (B, N, S) bool
+    target_lane_valid: torch.Tensor,  # (B, N) bool
+    sample_ys: torch.Tensor,        # (S,)
+    cls_weight: float = 1.0,
+    reg_weight: float = 5.0,
+) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    """Returns, per batch item, (pred_indices, gt_indices) of the matched pairs."""
+    B, N, _ = cls_logits.shape
+    probs = F.softmax(cls_logits, dim=-1)[..., 1]  # (B, N) prob of "lane"
+
+    pred_xs = eval_curve(curve_coeffs, sample_ys.view(1, 1, -1).expand(B, N, -1))  # (B, N, S)
+
+    results = []
+    for b in range(B):
+        valid_gt = target_lane_valid[b].nonzero(as_tuple=True)[0]  # (M,)
+        if valid_gt.numel() == 0:
+            results.append((torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long)))
+            continue
+
+        gt_xs = target_xs[b, valid_gt]            # (M, S)
+        gt_mask = target_valid_mask[b, valid_gt]  # (M, S)
+
+        # (N, M) cost matrix
+        diff = (pred_xs[b].unsqueeze(1) - gt_xs.unsqueeze(0)).abs()  # (N, M, S)
+        mask = gt_mask.unsqueeze(0).float()                          # (1, M, S)
+        denom = mask.sum(dim=-1).clamp(min=1.0)                      # (N broadcast, M)
+        reg_cost = (diff * mask).sum(dim=-1) / denom                 # (N, M)
+
+        cls_cost = -probs[b].unsqueeze(1).expand(-1, valid_gt.numel())  # (N, M)
+
+        cost = cls_weight * cls_cost + reg_weight * reg_cost
+        cost_np = cost.cpu().numpy()
+
+        pred_idx, gt_idx_local = linear_sum_assignment(cost_np)
+        gt_idx = valid_gt[gt_idx_local]
+
+        results.append((torch.as_tensor(pred_idx, dtype=torch.long),
+                         torch.as_tensor(gt_idx, dtype=torch.long)))
+    return results