first commit

models/head.py

@@ -0,0 +1,52 @@
+"""Per-query prediction heads: lane/background classification + cubic curve regression.
+
+Phase 1 note: curve coefficients are regressed directly (unconstrained) rather than
+with LSTR's full numerical-reparameterization trick -- kept simple for the first
+trainable pass; revisit if training shows y^3-term gradient instability.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int = 2):
+        super().__init__()
+        dims = [in_dim] + [hidden_dim] * (num_layers - 1) + [out_dim]
+        layers = []
+        for i in range(len(dims) - 1):
+            layers.append(nn.Linear(dims[i], dims[i + 1]))
+            if i < len(dims) - 2:
+                layers.append(nn.ReLU(inplace=True))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class PredictionHeads(nn.Module):
+    """x = k*y^3 + m*y^2 + n*y + b, plus (y_start, y_end), plus lane/background logits."""
+
+    def __init__(self, d_model: int, hidden_dim: int = 128):
+        super().__init__()
+        self.cls_head = MLP(d_model, hidden_dim, 2, num_layers=2)
+        self.curve_head = MLP(d_model, hidden_dim, 4, num_layers=3)     # k, m, n, b
+        self.endpoint_head = MLP(d_model, hidden_dim, 2, num_layers=2)  # y_start, y_end (pre-sigmoid)
+
+    def forward(self, queries: torch.Tensor) -> dict:
+        """queries: (B, N_lanes, d_model)"""
+        cls_logits = self.cls_head(queries)              # (B, N, 2)
+        curve_coeffs = self.curve_head(queries)           # (B, N, 4)
+        endpoints = torch.sigmoid(self.endpoint_head(queries))  # (B, N, 2) in [0,1]
+        return {"cls_logits": cls_logits, "curve_coeffs": curve_coeffs, "endpoints": endpoints}
+
+
+def eval_curve(curve_coeffs: torch.Tensor, ys: torch.Tensor) -> torch.Tensor:
+    """curve_coeffs: (..., 4) = [k,m,n,b]; ys: (...,S) or (S,) broadcastable normalized y.
+    Returns x(y) with shape (..., S).
+    """
+    k, m, n, b = curve_coeffs.unbind(dim=-1)  # each (...,)
+    k, m, n, b = k.unsqueeze(-1), m.unsqueeze(-1), n.unsqueeze(-1), b.unsqueeze(-1)
+    y = ys
+    return k * y**3 + m * y**2 + n * y + b