| """Standalone architecture extracted from the original lidarx implementation.""" |
| import torch |
| from torch import nn |
|
|
|
|
| def coarse_projection_features(x: torch.Tensor) -> torch.Tensor: |
| """Return `projection_coarse_32x8x8` features for `[B, 1, D, H, W]` tensors.""" |
| if x.ndim != 5: |
| msg = f"expected input with shape [batch, channel, depth, height, width], got {x.shape}" |
| raise ValueError(msg) |
| if x.shape[1] != 1: |
| msg = f"expected one histogram channel, got {x.shape[1]}" |
| raise ValueError(msg) |
|
|
| x = x.squeeze(1) |
| depth = x.shape[1] |
| height = x.shape[2] |
| width = x.shape[3] |
| if depth % 32 != 0 or height % 8 != 0 or width % 8 != 0: |
| msg = ( |
| "hist coarse projection requires depth divisible by 32 and " |
| f"height/width divisible by 8, got {(depth, height, width)}" |
| ) |
| raise ValueError(msg) |
|
|
| global_features = torch.cat( |
| [ |
| x.mean(dim=(1, 2, 3), keepdim=False).unsqueeze(1), |
| x.std(dim=(1, 2, 3), unbiased=False, keepdim=False).unsqueeze(1), |
| x.amin(dim=(1, 2, 3), keepdim=False).unsqueeze(1), |
| x.amax(dim=(1, 2, 3), keepdim=False).unsqueeze(1), |
| ], |
| dim=1, |
| ) |
| depth_features = torch.cat( |
| [ |
| x.mean(dim=(2, 3)), |
| x.std(dim=(2, 3), unbiased=False), |
| ], |
| dim=1, |
| ) |
| row_features = x.mean(dim=(1, 3)) |
| col_features = x.mean(dim=(1, 2)) |
| coarse = ( |
| x.reshape( |
| x.shape[0], |
| 32, |
| depth // 32, |
| 8, |
| height // 8, |
| 8, |
| width // 8, |
| ) |
| .mean(dim=(2, 4, 6)) |
| .reshape(x.shape[0], -1) |
| ) |
| return torch.cat( |
| [global_features, depth_features, row_features, col_features, coarse], |
| dim=1, |
| ) |
|
|
|
|
| def _coarse_grid_features( |
| x: torch.Tensor, |
| *, |
| depth_bins: int, |
| row_bins: int, |
| col_bins: int, |
| ) -> torch.Tensor: |
| depth = x.shape[1] |
| height = x.shape[2] |
| width = x.shape[3] |
| if depth % depth_bins != 0 or height % row_bins != 0 or width % col_bins != 0: |
| msg = ( |
| "hist coarse grid requires dimensions divisible by requested bins, " |
| f"got volume {(depth, height, width)} and bins {(depth_bins, row_bins, col_bins)}" |
| ) |
| raise ValueError(msg) |
|
|
| return ( |
| x.reshape( |
| x.shape[0], |
| depth_bins, |
| depth // depth_bins, |
| row_bins, |
| height // row_bins, |
| col_bins, |
| width // col_bins, |
| ) |
| .mean(dim=(2, 4, 6)) |
| .reshape(x.shape[0], -1) |
| ) |
|
|
|
|
| def multiscale_coarse_projection_features(x: torch.Tensor) -> torch.Tensor: |
| """Return coarse histogram features with added low-resolution context views.""" |
| base = coarse_projection_features(x) |
| volume = x.squeeze(1) |
| depth_context = _coarse_grid_features(volume, depth_bins=16, row_bins=8, col_bins=8) |
| spatial_context = _coarse_grid_features(volume, depth_bins=32, row_bins=4, col_bins=4) |
| return torch.cat([base, depth_context, spatial_context], dim=1) |
|
|
|
|
| class HistCoarseMultiScaleMLP(nn.Module): |
| """Small supervised MLP over multi-scale coarse histogram projections. |
| |
| This keeps the strong non-backbone inductive bias of the coarse-only model, |
| but gives the classifier both the best 32x8x8 block view and two coarser |
| context views before a lightweight gated residual head. |
| """ |
|
|
| def __init__(self, hidden_features: int = 128, dropout: float = 0.2) -> None: |
| super().__init__() |
| bottleneck_features = max(1, hidden_features // 2) |
| self.input_norm = nn.LazyBatchNorm1d() |
| self.input_projection = nn.Sequential( |
| nn.LazyLinear(hidden_features), |
| nn.SiLU(inplace=True), |
| nn.Dropout(dropout), |
| ) |
| self.residual = nn.Sequential( |
| nn.Linear(hidden_features, hidden_features), |
| nn.SiLU(inplace=True), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_features, hidden_features), |
| ) |
| self.gate = nn.Sequential( |
| nn.Linear(hidden_features, hidden_features), |
| nn.Sigmoid(), |
| ) |
| self.head = nn.Sequential( |
| nn.LayerNorm(hidden_features), |
| nn.SiLU(inplace=True), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_features, bottleneck_features), |
| nn.SiLU(inplace=True), |
| nn.Dropout(dropout), |
| nn.Linear(bottleneck_features, 1), |
| ) |
|
|
| def forward_features(self, x: torch.Tensor) -> torch.Tensor: |
| """Return the learned 128-D representation immediately before the classifier.""" |
| features = self.input_norm(multiscale_coarse_projection_features(x)) |
| embedding = self.input_projection(features) |
| return embedding + self.gate(embedding) * self.residual(embedding) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.head(self.forward_features(x)) |
|
|
