60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
"""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
|