diff --git a/autoware_ml/losses/detection2d/__init__.py b/autoware_ml/losses/detection2d/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/autoware_ml/losses/detection2d/losses.py b/autoware_ml/losses/detection2d/losses.py new file mode 100644 index 00000000..7b91491e --- /dev/null +++ b/autoware_ml/losses/detection2d/losses.py @@ -0,0 +1,222 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""2D detection losses for auxiliary image-plane supervision. + +These losses back the auxiliary 2D head of camera 3D detectors +(Focal-PETR-style): quality focal classification, GIoU box regression, and +Gaussian-heatmap centerness supervision. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from autoware_ml.models.detection3d.task_modules.boxes2d import bbox_overlaps + + +class QualityFocalLoss(nn.Module): + """Quality focal loss over sigmoid logits with IoU-quality targets. + + Positive queries are supervised toward the IoU between their predicted + and assigned boxes instead of a hard 1, following Generalized Focal Loss. + """ + + def __init__(self, beta: float = 2.0, loss_weight: float = 1.0) -> None: + """Initialize the quality focal loss. + + Args: + beta: Modulating exponent on the quality gap. + loss_weight: Multiplier applied to the final loss. + """ + super().__init__() + self.beta = beta + self.loss_weight = loss_weight + + def forward( + self, + logits: torch.Tensor, + labels: torch.Tensor, + quality_scores: torch.Tensor, + class_weights: torch.Tensor | None = None, + avg_factor: float | None = None, + ) -> torch.Tensor: + """Compute quality focal loss. + + Args: + logits: Classification logits with shape ``(N, C)``. + labels: Class indices with shape ``(N,)``; ``C`` marks background. + quality_scores: IoU quality targets with shape ``(N,)``, consumed + only at positive rows. + class_weights: Optional per-element weights with shape ``(N, C)`` + (used by partial-ignore to mask class columns). + avg_factor: Optional normalization factor. + + Returns: + Scalar loss value. + """ + num_classes = logits.size(1) + prob = logits.sigmoid() + # All positions start as background: target 0 modulated by p^beta. + loss = F.binary_cross_entropy_with_logits( + logits, torch.zeros_like(logits), reduction="none" + ) * prob.pow(self.beta) + + pos_inds = torch.nonzero((labels >= 0) & (labels < num_classes), as_tuple=False).squeeze(-1) + if pos_inds.numel() > 0: + pos_labels = labels[pos_inds].long() + quality = quality_scores[pos_inds].to(logits.dtype) + gap = (quality - prob[pos_inds, pos_labels]).abs().pow(self.beta) + loss[pos_inds, pos_labels] = ( + F.binary_cross_entropy_with_logits( + logits[pos_inds, pos_labels], quality, reduction="none" + ) + * gap + ) + if class_weights is not None: + loss = loss * class_weights + loss = loss.sum() + if avg_factor is not None: + loss = loss / max(avg_factor, 1.0) + return self.loss_weight * loss + + +class GIoULoss(nn.Module): + """GIoU loss between aligned unnormalized ``(x1, y1, x2, y2)`` boxes.""" + + def __init__(self, loss_weight: float = 1.0) -> None: + """Initialize the GIoU loss. + + Args: + loss_weight: Multiplier applied to the final loss. + """ + super().__init__() + self.loss_weight = loss_weight + + def forward( + self, + pred_boxes: torch.Tensor, + target_boxes: torch.Tensor, + weights: torch.Tensor | None = None, + avg_factor: float | None = None, + ) -> torch.Tensor: + """Compute the GIoU loss on aligned box pairs. + + Args: + pred_boxes: Predicted boxes with shape ``(N, 4)``. + target_boxes: Target boxes with shape ``(N, 4)``. + weights: Optional per-box weights with shape ``(N,)`` or ``(N, 4)``. + avg_factor: Optional normalization factor. + + Returns: + Scalar loss value. + """ + giou = bbox_overlaps(pred_boxes, target_boxes, mode="giou", is_aligned=True) + loss = 1.0 - giou + if weights is not None: + if weights.dim() > 1: + weights = weights.mean(dim=-1) + loss = loss * weights + loss = loss.sum() + if avg_factor is not None: + loss = loss / max(avg_factor, 1.0) + return self.loss_weight * loss + + +class WeightedL1Loss(nn.Module): + """Elementwise-weighted L1 loss normalized by an average factor.""" + + def __init__(self, loss_weight: float = 1.0) -> None: + """Initialize the weighted L1 loss. + + Args: + loss_weight: Multiplier applied to the final loss. + """ + super().__init__() + self.loss_weight = loss_weight + + def forward( + self, + prediction: torch.Tensor, + target: torch.Tensor, + weights: torch.Tensor | None = None, + avg_factor: float | None = None, + ) -> torch.Tensor: + """Compute the weighted L1 loss. + + Args: + prediction: Predicted values. + target: Target values with the same shape. + weights: Optional elementwise weights with the same shape. + avg_factor: Optional normalization factor. + + Returns: + Scalar loss value. + """ + loss = (prediction - target).abs() + if weights is not None: + loss = loss * weights + loss = loss.sum() + if avg_factor is not None: + loss = loss / max(avg_factor, 1.0) + return self.loss_weight * loss + + +class HeatmapGaussianFocalLoss(nn.Module): + """Gaussian focal loss on already-sigmoided probability heatmaps. + + The normalization factor is supplied by the caller (the auxiliary 2D head + shares one positive count across its five losses). + """ + + def __init__(self, alpha: float = 2.0, gamma: float = 4.0, loss_weight: float = 1.0) -> None: + """Initialize the heatmap Gaussian focal loss. + + Args: + alpha: Focusing parameter shared by positive and negative terms. + gamma: Modulating exponent on the negative Gaussian weights. + loss_weight: Multiplier applied to the final loss. + """ + super().__init__() + self.alpha = alpha + self.gamma = gamma + self.loss_weight = loss_weight + + def forward( + self, + probabilities: torch.Tensor, + target: torch.Tensor, + avg_factor: float | None = None, + ) -> torch.Tensor: + """Compute the Gaussian focal loss on probability heatmaps. + + Args: + probabilities: Sigmoid probabilities clipped away from 0 and 1. + target: Gaussian heatmap targets in ``[0, 1]``. + avg_factor: Optional normalization factor. + + Returns: + Scalar loss value. + """ + pos_weights = target.eq(1).to(probabilities.dtype) + # neg_weights is exactly zero wherever target == 1. + neg_weights = (1.0 - target).pow(self.gamma) + pos_loss = -probabilities.log() * (1.0 - probabilities).pow(self.alpha) * pos_weights + neg_loss = -(1.0 - probabilities).log() * probabilities.pow(self.alpha) * neg_weights + loss = pos_loss.sum() + neg_loss.sum() + if avg_factor is not None: + loss = loss / max(avg_factor, 1.0) + return self.loss_weight * loss diff --git a/autoware_ml/losses/detection3d/focal.py b/autoware_ml/losses/detection3d/focal.py index 08be47c5..75f13f2a 100644 --- a/autoware_ml/losses/detection3d/focal.py +++ b/autoware_ml/losses/detection3d/focal.py @@ -41,7 +41,9 @@ def forward( Args: logits: Raw classification logits with shape ``(N, C)``. targets: One-hot classification targets with shape ``(N, C)``. - weights: Optional per-query weights with shape ``(N,)``. + weights: Optional weights, either per-query with shape ``(N,)`` + or per-query-per-class with the same shape as ``logits`` + (used by partial-ignore to mask individual class columns). avg_factor: Optional normalization factor. Returns: @@ -53,7 +55,10 @@ def forward( alpha_factor = self.alpha * targets + (1.0 - self.alpha) * (1.0 - targets) loss = ce * alpha_factor * (1.0 - p_t).pow(self.gamma) if weights is not None: - loss = loss * weights.unsqueeze(-1) + if weights.dim() == loss.dim(): + loss = loss * weights + else: + loss = loss * weights.unsqueeze(-1) loss = loss.sum() if avg_factor is not None: loss = loss / max(avg_factor, 1.0) diff --git a/autoware_ml/models/common/necks/cp_fpn.py b/autoware_ml/models/common/necks/cp_fpn.py new file mode 100644 index 00000000..46538864 --- /dev/null +++ b/autoware_ml/models/common/necks/cp_fpn.py @@ -0,0 +1,108 @@ +# Copyright 2021 megvii-model. All Rights Reserved. +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint-friendly FPN used by the StreamPETR reference recipe. + +This is a native port of the StreamPETR ``CPFPN`` neck: plain 1x1 lateral +convolutions (no norm, no activation), a nearest-neighbor top-down pathway, +and a single 3x3 output convolution on the highest-resolution level only. +Unlike :class:`GeneralizedLSSFPN` it is weight-compatible with reference +StreamPETR checkpoints (parameter names mirror the mm ``ConvModule`` layout). +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class _ConvUnit(nn.Module): + """Hold one convolution under a ``conv`` attribute (mm ConvModule layout).""" + + def __init__(self, conv: nn.Conv2d) -> None: + super().__init__() + self.conv = conv + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply the wrapped convolution.""" + return self.conv(x) + + +class CPFPN(nn.Module): + """StreamPETR feature pyramid with per-level laterals and top-down adds. + + Outputs one feature map per input level; only the first (highest + resolution) level passes through a 3x3 refinement convolution, matching + the reference implementation. + """ + + def __init__(self, in_channels: Sequence[int], out_channels: int) -> None: + """Initialize the CPFPN neck. + + Args: + in_channels: Input channel dimensions ordered from high to low + resolution. + out_channels: Unified output channel dimension. + """ + super().__init__() + self.in_channels = list(in_channels) + self.out_channels = out_channels + self.lateral_convs = nn.ModuleList( + [ + _ConvUnit(nn.Conv2d(channels, out_channels, kernel_size=1)) + for channels in self.in_channels + ] + ) + self.fpn_convs = nn.ModuleList( + [_ConvUnit(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))] + ) + self.init_weights() + + def init_weights(self) -> None: + """Xavier-initialize every convolution (reference init).""" + for module in self.modules(): + if isinstance(module, nn.Conv2d): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + + def forward(self, inputs: Sequence[torch.Tensor]) -> tuple[torch.Tensor, ...]: + """Fuse the feature pyramid top-down and refine the finest level. + + Args: + inputs: Feature maps ordered from high to low resolution. + + Returns: + One fused feature map per input level. + """ + if len(inputs) != len(self.in_channels): + raise ValueError( + f"Expected {len(self.in_channels)} input feature maps, got {len(inputs)}." + ) + laterals = [ + lateral_conv(inputs[level]) for level, lateral_conv in enumerate(self.lateral_convs) + ] + for level in range(len(laterals) - 1, 0, -1): + laterals[level - 1] = laterals[level - 1] + F.interpolate( + laterals[level], size=laterals[level - 1].shape[2:], mode="nearest" + ) + outputs = [ + self.fpn_convs[0](laterals[0]) if level == 0 else laterals[level] + for level in range(len(laterals)) + ] + return tuple(outputs) diff --git a/autoware_ml/models/detection3d/heads/focal2d.py b/autoware_ml/models/detection3d/heads/focal2d.py new file mode 100644 index 00000000..3f8f1853 --- /dev/null +++ b/autoware_ml/models/detection3d/heads/focal2d.py @@ -0,0 +1,460 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Auxiliary 2D detection head for camera 3D detectors (Focal-PETR-style). + +The head predicts per-token class scores, centerness, LTRB boxes, and +projected 3D centers on every camera's neck feature map. Its five losses +shape the image features during training; inference never runs it, so it adds +no deployment cost. The reference FocalHead's top-k token pruning +(``sample_weight`` / ``topk_indexes``) is intentionally not ported: the 3D +head consumes every token, this head is loss-only. Ground truth comes from +projecting the 3D boxes onto each camera +(:class:`autoware_ml.transforms.camera.annotations2d.LoadAnnotations2DFromBoxes3D`). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from autoware_ml.losses.detection2d.losses import ( + GIoULoss, + HeatmapGaussianFocalLoss, + QualityFocalLoss, + WeightedL1Loss, +) +from autoware_ml.models.detection3d.partial_ignore import ( + mask_ignored_columns, + normalize_status_flags, + resolve_partial_ignore_labels, +) +from autoware_ml.models.detection3d.task_modules.assigners2d import HungarianAssigner2D +from autoware_ml.models.detection3d.task_modules.boxes2d import ( + bbox_cxcywh_to_xyxy, + bbox_overlaps, + bbox_xyxy_to_cxcywh, +) +from autoware_ml.models.detection3d.task_modules.heatmap import draw_heatmap_gaussian +from autoware_ml.models.detection3d.task_modules.streaming import ( + inverse_sigmoid, + reduce_mean_count, +) + + +def _token_locations( + feature_height: int, + feature_width: int, + stride: int, + image_height: int, + image_width: int, + device: torch.device, +) -> torch.Tensor: + """Build normalized pixel-center locations for every feature-map token. + + Returns: + Locations with shape ``(feature_height, feature_width, 2)`` as + ``(x, y)`` normalized by the padded image size. + """ + shifts_x = ( + torch.arange(0, stride * feature_width, step=stride, dtype=torch.float32, device=device) + + stride // 2 + ) / image_width + shifts_y = ( + torch.arange(0, stride * feature_height, step=stride, dtype=torch.float32, device=device) + + stride // 2 + ) / image_height + shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x, indexing="ij") + return torch.stack((shift_x, shift_y), dim=-1) + + +def _apply_ltrb(locations: torch.Tensor, pred_ltrb: torch.Tensor) -> torch.Tensor: + """Decode normalized LTRB distances into clamped ``(cx, cy, w, h)`` boxes.""" + pred_boxes = torch.stack( + [ + locations[..., 0] - pred_ltrb[..., 0], + locations[..., 1] - pred_ltrb[..., 1], + locations[..., 0] + pred_ltrb[..., 2], + locations[..., 1] + pred_ltrb[..., 3], + ], + dim=-1, + ).clamp(min=0.0, max=1.0) + return bbox_xyxy_to_cxcywh(pred_boxes) + + +def _apply_center_offset(locations: torch.Tensor, center_offset: torch.Tensor) -> torch.Tensor: + """Decode center offsets in inverse-sigmoid space into normalized centers.""" + return (inverse_sigmoid(locations) + center_offset).sigmoid() + + +@dataclass +class _Targets2D: + """Assignment targets for one camera image.""" + + labels: torch.Tensor + bbox_targets: torch.Tensor + bbox_weights: torch.Tensor + centers2d_targets: torch.Tensor + num_pos: int + + +class FocalHead2D(nn.Module): + """Dense auxiliary 2D head over multiview neck features. + + Ported from the StreamPETR ``FocalHead`` reference: shared 3x3 conv + + GroupNorm towers for classification and regression, with 1x1 prediction + convs for class logits, centerness, LTRB boxes, and projected 3D centers. + """ + + def __init__( + self, + num_classes: int, + in_channels: int = 256, + embed_dims: int = 256, + stride: int = 16, + assigner: HungarianAssigner2D | None = None, + loss_cls_weight: float = 2.0, + loss_bbox_weight: float = 5.0, + loss_iou_weight: float = 2.0, + loss_centers2d_weight: float = 10.0, + loss_centerness_weight: float = 1.0, + class_names: list[str] | None = None, + partial_ignore_classes: list[str] | None = None, + ) -> None: + """Initialize the auxiliary 2D head. + + Args: + num_classes: Number of detector classes. + in_channels: Neck feature channels. + embed_dims: Hidden channels of the shared conv towers. + stride: Feature-map stride relative to the padded image. + assigner: Hungarian assigner for proposal-target matching. + loss_cls_weight: Quality-focal classification loss weight. + loss_bbox_weight: Normalized box L1 loss weight. + loss_iou_weight: GIoU loss weight. + loss_centers2d_weight: Projected-center L1 loss weight. + loss_centerness_weight: Gaussian-heatmap centerness loss weight. + class_names: Ordered detector class names (for partial-ignore). + partial_ignore_classes: Class names that are only partially + annotated across scenes. + """ + super().__init__() + if embed_dims % 32 != 0: + raise ValueError(f"embed_dims must be a multiple of 32 (GroupNorm), got {embed_dims}.") + self.num_classes = num_classes + self.stride = stride + self.assigner = assigner if assigner is not None else HungarianAssigner2D() + self.partial_ignore_labels = resolve_partial_ignore_labels( + class_names, partial_ignore_classes + ) + + self.shared_cls = nn.Sequential( + nn.Conv2d(in_channels, embed_dims, kernel_size=3, padding=1), + nn.GroupNorm(32, num_channels=embed_dims), + nn.ReLU(), + ) + self.shared_reg = nn.Sequential( + nn.Conv2d(in_channels, embed_dims, kernel_size=3, padding=1), + nn.GroupNorm(32, num_channels=embed_dims), + nn.ReLU(), + ) + self.cls = nn.Conv2d(embed_dims, num_classes, kernel_size=1) + self.centerness = nn.Conv2d(embed_dims, 1, kernel_size=1) + self.ltrb = nn.Conv2d(embed_dims, 4, kernel_size=1) + self.center2d = nn.Conv2d(embed_dims, 2, kernel_size=1) + + bias_init = -math.log((1.0 - 0.01) / 0.01) + nn.init.constant_(self.cls.bias, bias_init) + nn.init.constant_(self.centerness.bias, bias_init) + + self.loss_cls2d = QualityFocalLoss(beta=2.0, loss_weight=loss_cls_weight) + self.loss_bbox2d = WeightedL1Loss(loss_weight=loss_bbox_weight) + self.loss_iou2d = GIoULoss(loss_weight=loss_iou_weight) + self.loss_centers2d = WeightedL1Loss(loss_weight=loss_centers2d_weight) + self.loss_centerness = HeatmapGaussianFocalLoss(loss_weight=loss_centerness_weight) + + def forward( + self, + img_features: torch.Tensor, + image_height: int, + image_width: int, + ) -> dict[str, Any]: + """Predict dense 2D outputs for every camera. + + Args: + img_features: Neck features ``(batch, num_cams, C, H, W)``. + image_height: Padded image height in pixels. + image_width: Padded image width in pixels. + + Returns: + Dense per-token predictions flattened to + ``(batch * num_cams, H * W, ...)``. + """ + batch_size, num_cams, _, feature_height, feature_width = img_features.shape + x = img_features.flatten(0, 1) + + cls_feat = self.shared_cls(x) + cls_logits = ( + self.cls(cls_feat) + .permute(0, 2, 3, 1) + .reshape(batch_size * num_cams, -1, self.num_classes) + ) + centerness = ( + self.centerness(cls_feat).permute(0, 2, 3, 1).reshape(batch_size * num_cams, -1, 1) + ) + + reg_feat = self.shared_reg(x) + ltrb = self.ltrb(reg_feat).permute(0, 2, 3, 1).contiguous().sigmoid() + centers2d_offset = self.center2d(reg_feat).permute(0, 2, 3, 1).contiguous() + + locations = _token_locations( + feature_height, feature_width, self.stride, image_height, image_width, x.device + )[None] + pred_bboxes = _apply_ltrb(locations, ltrb).view(batch_size * num_cams, -1, 4) + pred_centers2d = _apply_center_offset(locations, centers2d_offset).view( + batch_size * num_cams, -1, 2 + ) + + return { + "enc_cls_scores": cls_logits, + "enc_bbox_preds": pred_bboxes, + "pred_centers2d": pred_centers2d, + "centerness": centerness, + "pad_shape_2d": (image_height, image_width), + } + + def _get_targets_single( + self, + cls_logits: torch.Tensor, + bbox_pred: torch.Tensor, + pred_centers2d: torch.Tensor, + gt_bboxes: torch.Tensor, + gt_labels: torch.Tensor, + gt_centers2d: torch.Tensor, + image_height: int, + image_width: int, + ) -> _Targets2D: + num_queries = bbox_pred.size(0) + assigned = self.assigner.assign( + bbox_pred, + cls_logits, + pred_centers2d, + gt_bboxes, + gt_labels, + gt_centers2d, + image_height, + image_width, + ) + pos_inds = torch.nonzero(assigned.gt_inds > 0, as_tuple=False).squeeze(-1) + + labels = gt_labels.new_full((num_queries,), self.num_classes, dtype=torch.long) + bbox_targets = bbox_pred.new_zeros((num_queries, 4)) + bbox_weights = bbox_pred.new_zeros((num_queries, 4)) + centers2d_targets = bbox_pred.new_zeros((num_queries, 2)) + if pos_inds.numel() > 0: + matched_gt_inds = assigned.gt_inds[pos_inds] - 1 + labels[pos_inds] = gt_labels[matched_gt_inds].long() + factor = bbox_pred.new_tensor([image_width, image_height, image_width, image_height]) + bbox_targets[pos_inds] = bbox_xyxy_to_cxcywh(gt_bboxes[matched_gt_inds] / factor) + bbox_weights[pos_inds] = 1.0 + centers2d_targets[pos_inds] = gt_centers2d[matched_gt_inds] / factor[:2] + return _Targets2D( + labels=labels, + bbox_targets=bbox_targets, + bbox_weights=bbox_weights, + centers2d_targets=centers2d_targets, + num_pos=int(pos_inds.numel()), + ) + + def _build_heatmap( + self, + gt_centers2d: torch.Tensor, + gt_bboxes: torch.Tensor, + image_height: int, + image_width: int, + device: torch.device, + ) -> torch.Tensor: + heatmap = torch.zeros( + image_height // self.stride, image_width // self.stride, device=device + ) + if gt_centers2d.numel() == 0: + return heatmap + bounds = torch.cat( + [ + gt_centers2d[:, 0:1] - gt_bboxes[:, 0:1], + gt_centers2d[:, 1:2] - gt_bboxes[:, 1:2], + gt_bboxes[:, 2:3] - gt_centers2d[:, 0:1], + gt_bboxes[:, 3:4] - gt_centers2d[:, 1:2], + ], + dim=-1, + ) + radii = torch.ceil(bounds.min(dim=-1).values / self.stride).clamp(min=1.0) + for center, radius in zip(gt_centers2d / self.stride, radii.tolist()): + draw_heatmap_gaussian( + heatmap, (int(center[0].item()), int(center[1].item())), int(radius) + ) + return heatmap + + def loss( + self, + outputs: dict[str, torch.Tensor], + gt_bboxes_2d: list[list[torch.Tensor]], + gt_labels_2d: list[list[torch.Tensor]], + centers_2d: list[list[torch.Tensor]], + annotation_status: list[bool] | torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Compute the five auxiliary 2D losses. + + Args: + outputs: Model outputs holding the dense 2D predictions. + gt_bboxes_2d: Per-sample, per-camera projected 2D boxes in + unnormalized ``(x1, y1, x2, y2)`` pixels. + gt_labels_2d: Per-sample, per-camera class labels. + centers_2d: Per-sample, per-camera projected 3D centers in pixels. + annotation_status: Per-sample annotation-completeness flags for + partial-ignore; required when ``partial_ignore_classes`` is + configured. + + Returns: + Loss dictionary with ``loss_*2d`` keys. + """ + cls_scores = outputs["enc_cls_scores"] + bbox_preds = outputs["enc_bbox_preds"] + pred_centers2d = outputs["pred_centers2d"] + centerness = outputs["centerness"] + image_height, image_width = outputs["pad_shape_2d"] + device = cls_scores.device + + batch_size = len(gt_bboxes_2d) + num_cams = len(gt_bboxes_2d[0]) + num_images = batch_size * num_cams + if cls_scores.size(0) != num_images: + raise ValueError( + f"2D predictions cover {cls_scores.size(0)} images but annotations cover " + f"{num_images}; check the 2D annotation transform and collation." + ) + + def to_tensor(value: object, columns: int) -> torch.Tensor: + tensor = torch.as_tensor(value, dtype=torch.float32, device=device) + return tensor.reshape(-1, columns) + + targets: list[_Targets2D] = [] + heatmaps = [] + image_status: list[bool] = [] + if self.partial_ignore_labels is not None: + annotation_status = normalize_status_flags(annotation_status, batch_size) + else: + annotation_status = [True] * batch_size + for sample_index in range(batch_size): + for camera_index in range(num_cams): + image_index = sample_index * num_cams + camera_index + gt_bboxes = to_tensor(gt_bboxes_2d[sample_index][camera_index], 4) + gt_labels = torch.as_tensor( + gt_labels_2d[sample_index][camera_index], dtype=torch.long, device=device + ).reshape(-1) + gt_centers = to_tensor(centers_2d[sample_index][camera_index], 2) + targets.append( + self._get_targets_single( + cls_scores[image_index], + bbox_preds[image_index], + pred_centers2d[image_index], + gt_bboxes, + gt_labels, + gt_centers, + image_height, + image_width, + ) + ) + heatmaps.append( + self._build_heatmap(gt_centers, gt_bboxes, image_height, image_width, device) + ) + image_status.append(annotation_status[sample_index]) + + labels = torch.cat([target.labels for target in targets], dim=0) + bbox_targets = torch.cat([target.bbox_targets for target in targets], dim=0) + bbox_weights = torch.cat([target.bbox_weights for target in targets], dim=0) + centers2d_targets = torch.cat([target.centers2d_targets for target in targets], dim=0) + # Global (cross-rank) positive count; see reduce_mean_count. Safe + # collective: this method runs unconditionally on every training rank. + local_pos = cls_scores.new_tensor(float(sum(target.num_pos for target in targets))) + num_total_pos = torch.clamp(reduce_mean_count(local_pos), min=1.0).item() + + factor = cls_scores.new_tensor([image_width, image_height, image_width, image_height]) + flat_bbox_preds = bbox_preds.reshape(-1, 4) + pred_xyxy = bbox_cxcywh_to_xyxy(flat_bbox_preds) * factor + target_xyxy = bbox_cxcywh_to_xyxy(bbox_targets) * factor + + loss_iou = self.loss_iou2d( + pred_xyxy.float(), target_xyxy.float(), bbox_weights.float(), avg_factor=num_total_pos + ) + iou_scores = bbox_overlaps( + target_xyxy.float(), pred_xyxy.float(), mode="iou", is_aligned=True + ).detach() + + flat_cls_scores = cls_scores.reshape(-1, self.num_classes) + class_weights = self._partial_ignore_class_weights( + flat_cls_scores, image_status, cls_scores.size(1) + ) + loss_cls = self.loss_cls2d( + flat_cls_scores, + labels, + iou_scores, + class_weights=class_weights, + avg_factor=num_total_pos, + ) + + heatmap_targets = torch.stack(heatmaps, dim=0).view(num_images, -1, 1) + centerness_prob = centerness.sigmoid().clamp(min=1e-4, max=1 - 1e-4) + loss_centerness = self.loss_centerness( + centerness_prob, heatmap_targets, avg_factor=num_total_pos + ) + + loss_bbox = self.loss_bbox2d( + flat_bbox_preds, bbox_targets, bbox_weights, avg_factor=num_total_pos + ) + loss_centers2d = self.loss_centers2d( + pred_centers2d.reshape(-1, 2), + centers2d_targets, + bbox_weights[:, 0:2], + avg_factor=num_total_pos, + ) + + return { + "loss_cls2d": loss_cls, + "loss_bbox2d": loss_bbox, + "loss_iou2d": loss_iou, + "loss_centers2d": loss_centers2d, + "loss_centerness2d": loss_centerness, + } + + def _partial_ignore_class_weights( + self, + flat_cls_scores: torch.Tensor, + image_status: list[bool], + tokens_per_image: int, + ) -> torch.Tensor | None: + """Zero partially annotated class columns on unannotated images.""" + if self.partial_ignore_labels is None or all(image_status): + return None + class_weights = flat_cls_scores.new_ones(flat_cls_scores.shape) + for image_index, status in enumerate(image_status): + if not status: + start = image_index * tokens_per_image + rows = torch.arange(start, start + tokens_per_image, device=flat_cls_scores.device) + mask_ignored_columns(class_weights, rows, self.partial_ignore_labels) + return class_weights diff --git a/autoware_ml/models/detection3d/heads/streampetr.py b/autoware_ml/models/detection3d/heads/streampetr.py index 558b925b..c8ea0f46 100644 --- a/autoware_ml/models/detection3d/heads/streampetr.py +++ b/autoware_ml/models/detection3d/heads/streampetr.py @@ -14,6 +14,11 @@ import torch.nn as nn from autoware_ml.losses.detection3d.focal import SigmoidFocalLoss +from autoware_ml.models.detection3d.partial_ignore import ( + mask_ignored_columns, + normalize_status_flags, + resolve_partial_ignore_labels, +) from autoware_ml.models.detection3d.task_modules.assigners import HungarianAssigner3D from autoware_ml.models.detection3d.task_modules.bbox_coders import ( NMSFreeBBoxCoder3D, @@ -28,6 +33,7 @@ nerf_positional_encoding, pos2posemb1d, pos2posemb3d, + reduce_mean_count, topk_gather, transform_reference_points, ) @@ -110,6 +116,9 @@ class StreamPETRTargets: labels: torch.Tensor bbox_targets: torch.Tensor bbox_weights: torch.Tensor + # Per-query-per-class classification weights ``(num_queries, num_classes)`` + # used by partial-ignore; ``None`` means uniform weights. + label_weights: torch.Tensor | None = None class StreamPETRHead(nn.Module): @@ -146,9 +155,14 @@ def __init__( dn_weight: float = 1.0, split: float = 0.75, use_bottom_center: bool = True, + class_names: list[str] | None = None, + partial_ignore_classes: list[str] | None = None, ) -> None: super().__init__() self.num_classes = num_classes + self.partial_ignore_labels = resolve_partial_ignore_labels( + class_names, partial_ignore_classes + ) self.num_queries = num_queries self.num_decoder_layers = num_decoder_layers self.hidden_dim = hidden_dim @@ -744,11 +758,20 @@ def _get_targets( box_params: torch.Tensor, gt_boxes: list[torch.Tensor], gt_labels: list[torch.Tensor], + annotation_status: list[bool] | None = None, ) -> list[StreamPETRTargets]: + use_classwise_weights = ( + self.partial_ignore_labels is not None + and annotation_status is not None + and not all(annotation_status) + ) targets: list[StreamPETRTargets] = [] - for sample_logits, sample_boxes, sample_gt_boxes, sample_gt_labels in zip( - cls_logits, box_params, gt_boxes, gt_labels - ): + for sample_index, ( + sample_logits, + sample_boxes, + sample_gt_boxes, + sample_gt_labels, + ) in enumerate(zip(cls_logits, box_params, gt_boxes, gt_labels)): num_queries = sample_logits.shape[0] labels = sample_gt_labels.new_full((num_queries,), -1) bbox_targets = sample_boxes.new_zeros((num_queries, 9)) @@ -767,9 +790,26 @@ def _get_targets( labels[pos_inds] = sample_gt_labels[matched_gt_inds] bbox_targets[pos_inds] = sample_gt_boxes[matched_gt_inds] bbox_weights[pos_inds] = 1.0 + + label_weights = None + if use_classwise_weights: + label_weights = sample_logits.new_ones((num_queries, self.num_classes)) + if not annotation_status[sample_index]: + # Zero the ignored class columns on every query of the + # un-annotated frame (same semantics as FocalHead2D): no + # GT of those classes can exist here, so only their + # background terms are dropped. + mask_ignored_columns( + label_weights, + torch.arange(num_queries, device=label_weights.device), + self.partial_ignore_labels, + ) targets.append( StreamPETRTargets( - labels=labels, bbox_targets=bbox_targets, bbox_weights=bbox_weights + labels=labels, + bbox_targets=bbox_targets, + bbox_weights=bbox_weights, + label_weights=label_weights, ) ) return targets @@ -780,8 +820,9 @@ def _loss_single( bbox_preds: torch.Tensor, gt_boxes: list[torch.Tensor], gt_labels: list[torch.Tensor], + annotation_status: list[bool] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - targets = self._get_targets(cls_scores, bbox_preds, gt_boxes, gt_labels) + targets = self._get_targets(cls_scores, bbox_preds, gt_boxes, gt_labels, annotation_status) target_labels = [] positive_counts = [] @@ -792,10 +833,23 @@ def _loss_single( one_hot[pos_mask, sample_targets.labels[pos_mask]] = 1.0 target_labels.append(one_hot) target_labels_tensor = torch.stack(target_labels, dim=0) - # One device sync per decoder layer instead of one per sample. - total_pos = int(torch.stack(positive_counts).sum().item()) + # ``_get_targets`` materializes weights for either every sample or none. + label_weights_tensor = None + if targets and targets[0].label_weights is not None: + label_weights_tensor = torch.stack( + [sample_targets.label_weights for sample_targets in targets], dim=0 + ) + # Global (cross-rank) positive count; see reduce_mean_count for the + # normalization contract. One collective per decoder layer, uniform + # across ranks: this method runs unconditionally for every layer. + total_pos = torch.clamp( + reduce_mean_count(torch.stack(positive_counts).sum()), min=1.0 + ).item() loss_cls = self.loss_cls_weight * self.loss_cls( - cls_scores, target_labels_tensor, avg_factor=max(total_pos, 1) + cls_scores, + target_labels_tensor, + weights=label_weights_tensor, + avg_factor=total_pos, ) # The regression loss runs over every query with 0/1 positive weights so @@ -813,13 +867,13 @@ def _loss_single( target_tensor = torch.stack(target_encodings, dim=0) weight_tensor = torch.stack(positive_weights, dim=0).unsqueeze(-1) per_box = self.loss_bbox(encoded_preds[..., :10], target_tensor) * self.code_weights - bbox_loss = self.loss_bbox_weight * (per_box * weight_tensor).sum() / max(total_pos, 1) + bbox_loss = self.loss_bbox_weight * (per_box * weight_tensor).sum() / total_pos return loss_cls, bbox_loss def prepare_for_loss( self, mask_dict: dict[str, torch.Tensor], - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, torch.Tensor]: """Gather denoising outputs aligned with the replicated GT targets.""" output_known_class, output_known_coord = mask_dict["output_known_lbs_bboxes"] known_labels, known_bboxs = mask_dict["known_lbs_bboxes"] @@ -836,7 +890,45 @@ def prepare_for_loss( (batch_selection, map_known_indice) ].permute(1, 0, 2) num_targets = known_indices.numel() - return known_labels, known_bboxs, output_known_class, output_known_coord, num_targets + return ( + known_labels, + known_bboxs, + output_known_class, + output_known_coord, + num_targets, + batch_selection, + ) + + def _dn_label_weights( + self, + cls_scores: torch.Tensor, + known_labels: torch.Tensor, + known_bids: torch.Tensor | None, + annotation_status: list[bool] | None, + ) -> torch.Tensor | None: + """Build partial-ignore class weights for denoising queries. + + Every denoising query of an un-annotated frame drops the ignored class + columns (same semantics as :meth:`_get_targets`): those classes carry + no ground truth there, so only their background terms are removed. + """ + if ( + self.partial_ignore_labels is None + or known_bids is None + or known_bids.numel() == 0 + or annotation_status is None + or all(annotation_status) + ): + return None + status_tensor = torch.as_tensor( + annotation_status, device=known_labels.device, dtype=torch.bool + ) + ignore_rows = torch.nonzero(~status_tensor[known_bids.long()], as_tuple=False).squeeze(-1) + if ignore_rows.numel() == 0: + return None + label_weights = cls_scores.new_ones((known_labels.shape[0], self.num_classes)) + mask_ignored_columns(label_weights, ignore_rows, self.partial_ignore_labels) + return label_weights def dn_loss_single( self, @@ -844,13 +936,24 @@ def dn_loss_single( bbox_preds: torch.Tensor, known_bboxs: torch.Tensor, known_labels: torch.Tensor, - num_total_pos: int, + num_total_pos: float, + known_bids: torch.Tensor | None = None, + annotation_status: list[bool] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute denoising-query supervision for one decoder layer.""" + """Compute denoising-query supervision for one decoder layer. + + ``num_total_pos`` is the cross-rank mean denoising-target count + (already clamped to >= 1), synchronized once per step by + :meth:`loss` so that ranks without ground truth still participate in + the collective. + """ class_targets = cls_scores.new_zeros((known_labels.shape[0], self.num_classes)) foreground = known_labels < self.num_classes if foreground.any(): class_targets[foreground, known_labels[foreground]] = 1.0 + label_weights = self._dn_label_weights( + cls_scores, known_labels, known_bids, annotation_status + ) # The classification average factor scales the target count by the # expected positive rate of the noised queries (reference recipe): # a ball of radius ``split`` inside the unit noise cube. @@ -858,7 +961,9 @@ def dn_loss_single( loss_cls = ( self.dn_weight * self.loss_cls_weight - * self.loss_cls(cls_scores, class_targets, avg_factor=cls_avg_factor) + * self.loss_cls( + cls_scores, class_targets, weights=label_weights, avg_factor=cls_avg_factor + ) ) target_encoding = normalize_boxes3d(known_bboxs) @@ -868,7 +973,7 @@ def dn_loss_single( * ( self.loss_bbox(bbox_preds[:, :10], target_encoding[:, :10]) * self.code_weights ).sum() - / max(num_total_pos, 1) + / num_total_pos ) return loss_cls, loss_bbox @@ -993,16 +1098,28 @@ def loss( outputs: dict[str, torch.Tensor], gt_boxes: list[torch.Tensor], gt_labels: list[torch.Tensor], + annotation_status: list[bool] | torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: - """Compute multi-layer StreamPETR losses.""" + """Compute multi-layer StreamPETR losses. + + ``annotation_status`` carries one bool per sample (``False`` = the + frame's scene lacks the partially annotated classes); it is required + whenever ``partial_ignore_classes`` is configured. + """ all_cls_scores = outputs["all_cls_scores"] all_bbox_preds = outputs["all_bbox_preds"] gt_boxes = self._gravity_center_boxes(gt_boxes) + if self.partial_ignore_labels is not None: + annotation_status = normalize_status_flags(annotation_status, len(gt_boxes)) + else: + annotation_status = None losses_cls = [] losses_bbox = [] for cls_scores, bbox_preds in zip(all_cls_scores, all_bbox_preds): - loss_cls, loss_bbox = self._loss_single(cls_scores, bbox_preds, gt_boxes, gt_labels) + loss_cls, loss_bbox = self._loss_single( + cls_scores, bbox_preds, gt_boxes, gt_labels, annotation_status + ) losses_cls.append(loss_cls) losses_bbox.append(loss_bbox) @@ -1027,14 +1144,42 @@ def loss( loss_dict[f"d{layer_index}.dn_loss_bbox"] = zero.clone() if outputs["dn_mask_dict"] is not None: - known_labels, known_bboxs, output_known_class, output_known_coord, num_tgt = ( - self.prepare_for_loss(outputs["dn_mask_dict"]) - ) + ( + known_labels, + known_bboxs, + output_known_class, + output_known_coord, + num_tgt, + known_bids, + ) = self.prepare_for_loss(outputs["dn_mask_dict"]) + else: + num_tgt = 0 + + if self.with_dn and self.training: + # Cross-rank mean denoising-target count, synced once per step. + # This collective is gated on flags that are uniform across ranks + # (never on this rank's GT presence): a rank whose batch has no + # ground truth must still participate, contributing count 0 — + # its DN loss is the zero filled above, which is exactly its + # share of the global per-object mean. + dn_avg_factor = torch.clamp( + reduce_mean_count(all_cls_scores.new_tensor(float(num_tgt))), min=1.0 + ).item() + else: + dn_avg_factor = max(float(num_tgt), 1.0) + + if outputs["dn_mask_dict"] is not None: dn_losses_cls = [] dn_losses_bbox = [] for cls_scores, bbox_preds in zip(output_known_class, output_known_coord): dn_loss_cls, dn_loss_bbox = self.dn_loss_single( - cls_scores, bbox_preds, known_bboxs, known_labels, num_tgt + cls_scores, + bbox_preds, + known_bboxs, + known_labels, + dn_avg_factor, + known_bids=known_bids, + annotation_status=annotation_status, ) dn_losses_cls.append(dn_loss_cls) dn_losses_bbox.append(dn_loss_bbox) diff --git a/autoware_ml/models/detection3d/partial_ignore.py b/autoware_ml/models/detection3d/partial_ignore.py new file mode 100644 index 00000000..5ba214f1 --- /dev/null +++ b/autoware_ml/models/detection3d/partial_ignore.py @@ -0,0 +1,114 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Partial-ignore support for partially annotated classes. + +Some dataset scenes are annotated for every class except a subset. Training on +such frames must not punish background predictions of the un-annotated classes +as false positives. The per-frame ``annotation_status`` flag marks whether the +frame's scene carries those annotations; when it is ``False``, classification +weights for the ignored class columns are zeroed on that frame's queries. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +def resolve_partial_ignore_labels( + class_names: Sequence[str] | None, + partial_ignore_classes: Sequence[str] | None, +) -> list[int] | None: + """Map partially annotated class names to label indices. + + Args: + class_names: Ordered detector class names. + partial_ignore_classes: Class names that are only partially annotated. + + Returns: + Label indices of the partially annotated classes, or ``None`` when + partial-ignore is disabled. + """ + if not partial_ignore_classes: + return None + if class_names is None: + raise ValueError("class_names is required when partial_ignore_classes is set.") + name_to_index = {name: index for index, name in enumerate(class_names)} + missing = [name for name in partial_ignore_classes if name not in name_to_index] + if missing: + raise ValueError(f"partial_ignore_classes {missing} not found in class_names.") + return [name_to_index[name] for name in partial_ignore_classes] + + +def normalize_status_flags(value: object, batch_size: int) -> list[bool]: + """Normalize per-sample annotation-status flags to a plain bool list. + + Accepts tensors, nested lists, or scalars produced by different collation + paths. + + Args: + value: Raw ``annotation_status`` value from the batch. + batch_size: Expected number of samples. + + Returns: + One bool per sample. + + Raises: + ValueError: If the flags are missing or their count does not match + ``batch_size`` — silently defaulting would train the ignored + classes' background as false positives. + """ + if value is None: + raise ValueError( + "annotation_status is required when partial_ignore_classes is configured; " + "check the datamodule's annotation_status_field and collation_map." + ) + flattened = _flatten_flags(value) + if len(flattened) != batch_size: + raise ValueError( + f"annotation_status carries {len(flattened)} flags for a batch of " + f"{batch_size} samples; check the collation_map strategy for the key." + ) + return flattened + + +def mask_ignored_columns( + weights: torch.Tensor, + rows: torch.Tensor, + ignore_labels: Sequence[int], +) -> None: + """Zero the ignored class columns on the selected rows, in place. + + Args: + weights: Classification weights ``(num_rows, num_classes)``. + rows: Row indices belonging to un-annotated frames. + ignore_labels: Label indices of the partially annotated classes. + """ + if rows.numel() == 0: + return + columns = torch.as_tensor(ignore_labels, device=weights.device, dtype=torch.long) + weights[rows[:, None], columns] = 0.0 + + +def _flatten_flags(value: object) -> list[bool]: + if torch.is_tensor(value): + return [bool(item) for item in value.detach().cpu().flatten().tolist()] + if isinstance(value, (list, tuple)): + flattened: list[bool] = [] + for item in value: + flattened.extend(_flatten_flags(item)) + return flattened + return [bool(value)] diff --git a/autoware_ml/models/detection3d/streampetr.py b/autoware_ml/models/detection3d/streampetr.py index f3485f5d..7ecd68d9 100644 --- a/autoware_ml/models/detection3d/streampetr.py +++ b/autoware_ml/models/detection3d/streampetr.py @@ -216,10 +216,12 @@ def __init__( img_backbone: nn.Module, img_neck: nn.Module, bbox_head: nn.Module, + img_roi_head: nn.Module | None = None, use_grid_mask: bool = False, optimizer: Callable[..., Optimizer] | None = None, scheduler: Callable[[Optimizer], LRScheduler] | None = None, optimizer_group_overrides: Mapping[str, Mapping[str, Any]] | None = None, + scheduler_config: Mapping[str, Any] | None = None, metrics: Sequence[MetricSuite] | None = None, ) -> None: """Initialize StreamPETR. @@ -228,22 +230,28 @@ def __init__( img_backbone: Image backbone. img_neck: Image neck. bbox_head: Detection head. + img_roi_head: Optional auxiliary 2D detection head supervising the + image features during training (Focal-PETR-style). Inference + never runs it. use_grid_mask: Whether to apply grid-mask image augmentation during training. optimizer: Optimizer factory. scheduler: Scheduler factory. optimizer_group_overrides: Per-submodule optimizer overrides. + scheduler_config: Lightning scheduler metadata such as ``interval``. metrics: Detection metrics accumulated during validation and test. """ super().__init__( optimizer=optimizer, scheduler=scheduler, optimizer_group_overrides=optimizer_group_overrides, + scheduler_config=scheduler_config, metrics=metrics, ) self.img_backbone = img_backbone self.img_neck = img_neck self.bbox_head = bbox_head + self.img_roi_head = img_roi_head self.use_grid_mask = use_grid_mask self.grid_mask = GridMask() self.image_feature_extractor = MultiviewImageFeatureExtractor( @@ -336,7 +344,7 @@ def forward( batch_size, num_cams, *image_batch.shape[2:] ) img_features = self._extract_img_features(image_batch) - return self.bbox_head( + outputs = self.bbox_head( img_features=img_features, img=image_batch, camera_intrinsics=self._stack_optional_tensor(camera_intrinsics), @@ -349,6 +357,17 @@ def forward( gt_boxes=gt_boxes, gt_labels=gt_labels, ) + # The auxiliary 2D head only shapes image features during training; + # inference and deployment never execute it. + if self.img_roi_head is not None and self.training: + outputs.update( + self.img_roi_head( + img_features, + image_height=int(image_batch.shape[-2]), + image_width=int(image_batch.shape[-1]), + ) + ) + return outputs def _stack_optional_tensor( self, @@ -377,9 +396,25 @@ def compute_metrics( Returns: Loss dictionary produced by the detection head. """ - return self.bbox_head.loss( - outputs, batch_inputs_dict["gt_boxes"], batch_inputs_dict["gt_labels"] + losses = self.bbox_head.loss( + outputs, + batch_inputs_dict["gt_boxes"], + batch_inputs_dict["gt_labels"], + annotation_status=batch_inputs_dict.get("annotation_status"), ) + if self.img_roi_head is not None and self.training: + roi_losses = self.img_roi_head.loss( + outputs, + gt_bboxes_2d=batch_inputs_dict["gt_bboxes_2d"], + gt_labels_2d=batch_inputs_dict["gt_labels_2d"], + centers_2d=batch_inputs_dict["centers_2d"], + annotation_status=batch_inputs_dict.get("annotation_status"), + ) + losses.update(roi_losses) + losses["loss"] = losses["loss"] + sum( + value for key, value in roi_losses.items() if key.startswith("loss_") + ) + return losses def predict_outputs( self, batch_inputs_dict: dict[str, Any], outputs: dict[str, torch.Tensor] diff --git a/autoware_ml/models/detection3d/task_modules/assigners2d.py b/autoware_ml/models/detection3d/task_modules/assigners2d.py new file mode 100644 index 00000000..cb500442 --- /dev/null +++ b/autoware_ml/models/detection3d/task_modules/assigners2d.py @@ -0,0 +1,157 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hungarian assignment for the auxiliary 2D detection head. + +The assigner matches dense image-plane proposals against projected 2D ground +truth with the DETR-style weighted cost: classification, box L1, GIoU, and +projected-center L1 (the Focal-PETR extension). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +from scipy.optimize import linear_sum_assignment + +from autoware_ml.models.detection3d.task_modules.assigners import AssignResult +from autoware_ml.models.detection3d.task_modules.boxes2d import ( + bbox_cxcywh_to_xyxy, + bbox_overlaps, + bbox_xyxy_to_cxcywh, +) +from autoware_ml.models.detection3d.task_modules.match_costs import ClassificationCost + + +@dataclass(frozen=True) +class BBoxL1Cost2D: + """Pairwise L1 cost between normalized 2D boxes. + + Predictions are ``(cx, cy, w, h)`` in [0, 1]; ground truth arrives as + normalized ``(x1, y1, x2, y2)`` and is converted according to + ``box_format``. + """ + + weight: float = 1.0 + box_format: str = "xywh" + + def __call__(self, bbox_pred: torch.Tensor, gt_bboxes: torch.Tensor) -> torch.Tensor: + """Compute the pairwise box L1 cost.""" + if self.box_format == "xywh": + gt_bboxes = bbox_xyxy_to_cxcywh(gt_bboxes) + elif self.box_format == "xyxy": + bbox_pred = bbox_cxcywh_to_xyxy(bbox_pred) + else: + raise ValueError(f"Unsupported box_format '{self.box_format}'.") + return torch.cdist(bbox_pred.float(), gt_bboxes.float(), p=1) * self.weight + + +@dataclass(frozen=True) +class IoUCost2D: + """Pairwise negative-IoU (or GIoU) cost between unnormalized 2D boxes.""" + + weight: float = 1.0 + iou_mode: str = "giou" + + def __call__(self, bboxes: torch.Tensor, gt_bboxes: torch.Tensor) -> torch.Tensor: + """Compute the pairwise overlap cost.""" + overlaps = bbox_overlaps(bboxes, gt_bboxes, mode=self.iou_mode, is_aligned=False) + return -overlaps * self.weight + + +@dataclass(frozen=True) +class Center2DL1Cost: + """Pairwise L1 cost between normalized projected 2D centers.""" + + weight: float = 1.0 + + def __call__(self, center_pred: torch.Tensor, gt_centers: torch.Tensor) -> torch.Tensor: + """Compute the pairwise center L1 cost.""" + return torch.cdist(center_pred.float(), gt_centers.float(), p=1) * self.weight + + +@dataclass +class HungarianAssigner2D: + """One-to-one assignment between dense 2D proposals and ground truth. + + The default cost weights mirror the reference Focal-PETR recipe + (2 / 5 / 2 / 10), matching :class:`FocalHead2D`'s default loss weights. + """ + + cls_cost: ClassificationCost = field(default_factory=lambda: ClassificationCost(weight=2.0)) + reg_cost: BBoxL1Cost2D = field(default_factory=lambda: BBoxL1Cost2D(weight=5.0)) + iou_cost: IoUCost2D = field(default_factory=lambda: IoUCost2D(weight=2.0)) + centers2d_cost: Center2DL1Cost = field(default_factory=lambda: Center2DL1Cost(weight=10.0)) + + def assign( + self, + bbox_pred: torch.Tensor, + cls_pred: torch.Tensor, + pred_centers2d: torch.Tensor, + gt_bboxes: torch.Tensor, + gt_labels: torch.Tensor, + gt_centers2d: torch.Tensor, + image_height: int, + image_width: int, + ) -> AssignResult: + """Assign each proposal to background (0) or a 1-based ground-truth index. + + Args: + bbox_pred: Proposal boxes ``(num_queries, 4)`` as normalized + ``(cx, cy, w, h)``. + cls_pred: Proposal classification logits ``(num_queries, C)``. + pred_centers2d: Proposal centers ``(num_queries, 2)`` normalized. + gt_bboxes: Ground-truth boxes ``(num_gts, 4)`` in unnormalized + ``(x1, y1, x2, y2)`` pixels. + gt_labels: Ground-truth labels ``(num_gts,)``. + gt_centers2d: Ground-truth projected centers ``(num_gts, 2)`` in + unnormalized pixels. + image_height: Padded image height in pixels. + image_width: Padded image width in pixels. + + Returns: + Assignment result over the proposals. + """ + num_gts = gt_bboxes.size(0) + num_queries = bbox_pred.size(0) + assigned_gt_inds = bbox_pred.new_zeros((num_queries,), dtype=torch.long) + assigned_labels = bbox_pred.new_full((num_queries,), -1, dtype=torch.long) + if num_gts == 0 or num_queries == 0: + return AssignResult( + num_gts=num_gts, + gt_inds=assigned_gt_inds, + max_overlaps=None, + labels=assigned_labels, + ) + + factor = gt_bboxes.new_tensor([image_width, image_height, image_width, image_height]) + cost = ( + self.cls_cost(cls_pred, gt_labels) + + self.reg_cost(bbox_pred, gt_bboxes / factor) + + self.iou_cost(bbox_cxcywh_to_xyxy(bbox_pred) * factor, gt_bboxes) + + self.centers2d_cost(pred_centers2d, gt_centers2d / factor[:2]) + ) + cost = torch.nan_to_num(cost, nan=100.0, posinf=100.0, neginf=-100.0) + matched_rows, matched_cols = linear_sum_assignment(cost.detach().cpu().numpy()) + matched_rows = torch.as_tensor(matched_rows, device=bbox_pred.device, dtype=torch.long) + matched_cols = torch.as_tensor(matched_cols, device=bbox_pred.device, dtype=torch.long) + assigned_gt_inds[matched_rows] = matched_cols + 1 + assigned_labels[matched_rows] = gt_labels[matched_cols].long() + return AssignResult( + num_gts=num_gts, + gt_inds=assigned_gt_inds, + max_overlaps=None, + labels=assigned_labels, + ) diff --git a/autoware_ml/models/detection3d/task_modules/boxes2d.py b/autoware_ml/models/detection3d/task_modules/boxes2d.py new file mode 100644 index 00000000..3cdd7255 --- /dev/null +++ b/autoware_ml/models/detection3d/task_modules/boxes2d.py @@ -0,0 +1,85 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""2D bounding-box utilities for auxiliary image-plane supervision. + +These helpers back the auxiliary 2D detection head used by camera 3D +detectors (Focal-PETR-style): box format conversion and IoU/GIoU overlaps. +""" + +from __future__ import annotations + +import torch + + +def bbox_cxcywh_to_xyxy(boxes: torch.Tensor) -> torch.Tensor: + """Convert ``(cx, cy, w, h)`` boxes to ``(x1, y1, x2, y2)``.""" + cx, cy, w, h = boxes.unbind(-1) + return torch.stack([cx - 0.5 * w, cy - 0.5 * h, cx + 0.5 * w, cy + 0.5 * h], dim=-1) + + +def bbox_xyxy_to_cxcywh(boxes: torch.Tensor) -> torch.Tensor: + """Convert ``(x1, y1, x2, y2)`` boxes to ``(cx, cy, w, h)``.""" + x1, y1, x2, y2 = boxes.unbind(-1) + return torch.stack([(x1 + x2) * 0.5, (y1 + y2) * 0.5, x2 - x1, y2 - y1], dim=-1) + + +def bbox_overlaps( + boxes1: torch.Tensor, + boxes2: torch.Tensor, + mode: str = "iou", + is_aligned: bool = False, + eps: float = 1e-6, +) -> torch.Tensor: + """Compute IoU or GIoU between two sets of ``(x1, y1, x2, y2)`` boxes. + + Args: + boxes1: Boxes with shape ``(N, 4)``. + boxes2: Boxes with shape ``(M, 4)`` (or ``(N, 4)`` when aligned). + mode: ``"iou"`` or ``"giou"``. + is_aligned: Compute element-wise overlaps instead of pairwise. + eps: Numerical stability constant. + + Returns: + Overlap tensor with shape ``(N,)`` when aligned, else ``(N, M)``. + """ + if mode not in ("iou", "giou"): + raise ValueError(f"Unsupported overlap mode '{mode}'.") + area1 = (boxes1[..., 2] - boxes1[..., 0]).clamp(min=0) * ( + boxes1[..., 3] - boxes1[..., 1] + ).clamp(min=0) + area2 = (boxes2[..., 2] - boxes2[..., 0]).clamp(min=0) * ( + boxes2[..., 3] - boxes2[..., 1] + ).clamp(min=0) + + if not is_aligned: + boxes1 = boxes1[:, None, :] + boxes2 = boxes2[None, :, :] + area1 = area1[:, None] + area2 = area2[None, :] + + lt = torch.maximum(boxes1[..., :2], boxes2[..., :2]) + rb = torch.minimum(boxes1[..., 2:], boxes2[..., 2:]) + wh = (rb - lt).clamp(min=0) + intersection = wh[..., 0] * wh[..., 1] + union = (area1 + area2 - intersection).clamp(min=eps) + ious = intersection / union + if mode == "iou": + return ious + + enclose_lt = torch.minimum(boxes1[..., :2], boxes2[..., :2]) + enclose_rb = torch.maximum(boxes1[..., 2:], boxes2[..., 2:]) + enclose_wh = (enclose_rb - enclose_lt).clamp(min=0) + enclose_area = (enclose_wh[..., 0] * enclose_wh[..., 1]).clamp(min=eps) + return ious - (enclose_area - union) / enclose_area diff --git a/autoware_ml/models/detection3d/task_modules/streaming.py b/autoware_ml/models/detection3d/task_modules/streaming.py index 50e7f6a0..2adff572 100644 --- a/autoware_ml/models/detection3d/task_modules/streaming.py +++ b/autoware_ml/models/detection3d/task_modules/streaming.py @@ -13,6 +13,22 @@ import torch.nn as nn +def reduce_mean_count(value: torch.Tensor) -> torch.Tensor: + """Average a scalar sample count across DDP ranks (mmdetection's ``reduce_mean``). + + DDP averages gradients, so normalizing every rank's ``sum`` loss by this + global mean count yields the true per-object mean over the whole effective + batch. Value-preserving float cast when distributed is unavailable or + uninitialized. Collective: every rank in the default group must call this + the same number of times per step. + """ + if torch.distributed.is_available() and torch.distributed.is_initialized(): + value = value.float().clone() + torch.distributed.all_reduce(value.div_(torch.distributed.get_world_size())) + return value + return value.float() + + def inverse_sigmoid(x: torch.Tensor, eps: float = 1e-4) -> torch.Tensor: """Apply the inverse sigmoid transform with clamping for stability.""" dtype = x.dtype diff --git a/autoware_ml/tests/models/test_cp_fpn.py b/autoware_ml/tests/models/test_cp_fpn.py new file mode 100644 index 00000000..2ffda30a --- /dev/null +++ b/autoware_ml/tests/models/test_cp_fpn.py @@ -0,0 +1,17 @@ +"""Tests for the CPFPN neck.""" + +from __future__ import annotations + +import torch + +from autoware_ml.models.common.necks.cp_fpn import CPFPN + + +def test_cpfpn_outputs_match_input_levels_and_channels() -> None: + neck = CPFPN(in_channels=[24, 40], out_channels=16) + high = torch.randn(2, 24, 12, 20) + low = torch.randn(2, 40, 6, 10) + outputs = neck((high, low)) + assert len(outputs) == 2 + assert outputs[0].shape == (2, 16, 12, 20) + assert outputs[1].shape == (2, 16, 6, 10) diff --git a/autoware_ml/tests/models/test_streampetr_loss_distributed.py b/autoware_ml/tests/models/test_streampetr_loss_distributed.py new file mode 100644 index 00000000..eccf6d62 --- /dev/null +++ b/autoware_ml/tests/models/test_streampetr_loss_distributed.py @@ -0,0 +1,246 @@ +"""Distributed (multi-rank) tests for the StreamPETR loss normalization. + +These spawn real CPU process groups with the gloo backend so the actual +``reduce_mean_count`` collectives run. + +The correctness property checked is the invariant that global (cross-rank) +normalization provides and rank-local normalization violates: **sharding a +batch across ranks must produce the same DDP-averaged loss as computing the +whole batch on one rank**. Rank-local counting fails this whenever ranks hold +different object counts (each GPU gets an equal vote regardless of load, an +upward-biased mean of per-rank means), which also made results silently depend +on the GPU count. + +The DN test additionally pins the collective-uniformity contract: a rank whose +batch has no ground truth must still reach the DN count collective, otherwise +mixed-GT steps deadlock. Workers are joined with a timeout so a regression +shows up as a test failure rather than a hung pytest. +""" + +from __future__ import annotations + +import os +import tempfile +import time + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from autoware_ml.models.detection3d.heads.streampetr import StreamPETRHead +from autoware_ml.models.detection3d.task_modules.assigners import HungarianAssigner3D +from autoware_ml.models.detection3d.task_modules.bbox_coders import NMSFreeBBoxCoder3D +from autoware_ml.models.detection3d.task_modules.match_costs import ( + BBox3DL1Cost, + ClassificationCost, + IoU3DCost, +) +from autoware_ml.models.detection3d.task_modules.streaming import reduce_mean_count + +pytestmark = pytest.mark.skipif( + not (dist.is_available() and dist.is_gloo_available()), + reason="gloo backend unavailable", +) + +_WORLD_SIZE = 2 +_NUM_CLASSES = 3 +_NUM_QUERIES = 32 + + +def _build_head() -> StreamPETRHead: + torch.manual_seed(0) + return StreamPETRHead( + num_classes=_NUM_CLASSES, + in_channels=128, + hidden_dim=128, + num_queries=_NUM_QUERIES, + num_decoder_layers=3, + num_heads=4, + feedforward_channels=256, + memory_len=32, + topk_proposals=8, + num_propagated=8, + with_dn=True, + with_ego_pos=True, + depth_num=8, + LID=True, + position_range=[-12.0, -12.0, -6.0, 12.0, 12.0, 6.0], + scalar=2, + noise_scale=0.5, + dn_weight=1.0, + split=0.5, + use_bottom_center=True, + bbox_coder=NMSFreeBBoxCoder3D( + pc_range=[-10.0, -10.0, -5.0, 10.0, 10.0, 5.0], + post_center_range=[-12.0, -12.0, -6.0, 12.0, 12.0, 6.0], + score_threshold=0.01, + max_num=16, + ), + assigner=HungarianAssigner3D( + cls_cost=ClassificationCost(weight=2.0), + reg_cost=BBox3DL1Cost( + weight=0.25, code_weights=(2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0) + ), + iou_cost=IoU3DCost(weight=0.0), + ), + point_cloud_range=[-10.0, -10.0, -5.0, 10.0, 10.0, 5.0], + code_weights=[2.0, 2.0] + [1.0] * 8, + ) + + +def _sample( + num_boxes: int, seed: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Deterministic (cls_scores, bbox_preds, gt_boxes, gt_labels) for one frame.""" + generator = torch.Generator().manual_seed(seed) + cls_scores = torch.randn(_NUM_QUERIES, _NUM_CLASSES, generator=generator) + bbox_preds = torch.randn(_NUM_QUERIES, 10, generator=generator) + gt_boxes = torch.rand(num_boxes, 9, generator=generator) * 4.0 - 2.0 + gt_boxes[:, 3:6] = gt_boxes[:, 3:6].abs() + 1.0 # positive dims + gt_labels = torch.arange(num_boxes, dtype=torch.long) % _NUM_CLASSES + return cls_scores, bbox_preds, gt_boxes, gt_labels + + +def _spawn(worker, world_size: int = _WORLD_SIZE, timeout: float = 120.0) -> None: + with tempfile.TemporaryDirectory() as directory: + context = mp.spawn( + worker, + args=(world_size, os.path.join(directory, "store")), + nprocs=world_size, + join=False, + ) + # A finite join converts a collective-uniformity regression (deadlock) + # into a test failure instead of a hung pytest process. join(timeout) + # returns after a single wait pass, so poll until the deadline. + deadline = time.monotonic() + timeout + while not context.join(timeout=5.0): + if time.monotonic() > deadline: + for process in context.processes: + process.terminate() + raise AssertionError(f"distributed workers did not finish within {timeout}s") + + +def _reduce_mean_count_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size + ) + try: + # Rank counts 2 and 8 (the section-6.4 worked example): mean must be 5. + local_count = torch.tensor(float(2 if rank == 0 else 8)) + reduced = reduce_mean_count(local_count) + assert reduced.item() == pytest.approx(5.0) + # The input must not be mutated in place. + assert local_count.item() == pytest.approx(2.0 if rank == 0 else 8.0) + finally: + dist.destroy_process_group() + + +def _sharding_invariance_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size + ) + try: + head = _build_head().train() + # Rank 0 holds the sparse frame (2 boxes), rank 1 the dense one (8). + frames = [_sample(2, seed=10), _sample(8, seed=20)] + + # Sharded pass: each rank sees only its own frame. + cls_scores, bbox_preds, gt_boxes, gt_labels = frames[rank] + loss_cls, loss_bbox = head._loss_single( + cls_scores.unsqueeze(0), bbox_preds.unsqueeze(0), [gt_boxes], [gt_labels] + ) + sharded = torch.stack([loss_cls.detach(), loss_bbox.detach()]) + gathered = [torch.zeros_like(sharded) for _ in range(world_size)] + dist.all_gather(gathered, sharded) + ddp_averaged = torch.stack(gathered).mean(dim=0) + + # Combined pass: the same two frames as one batch on every rank. + # reduce_mean_count then averages identical counts, so this equals the + # true per-object mean over the union. + combined_cls, combined_bbox = head._loss_single( + torch.stack([frames[0][0], frames[1][0]]), + torch.stack([frames[0][1], frames[1][1]]), + [frames[0][2], frames[1][2]], + [frames[0][3], frames[1][3]], + ) + + torch.testing.assert_close(ddp_averaged[0], combined_cls.detach(), rtol=1e-5, atol=1e-6) + torch.testing.assert_close(ddp_averaged[1], combined_bbox.detach(), rtol=1e-5, atol=1e-6) + finally: + dist.destroy_process_group() + + +def _dn_mixed_gt_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size + ) + try: + head = _build_head().train() + if rank == 0: + _, _, gt_boxes, gt_labels = _sample(4, seed=30) + else: + gt_boxes = torch.zeros(0, 9) + gt_labels = torch.zeros(0, dtype=torch.long) + + # Minimal outputs contract for StreamPETRHead.loss: per-layer scores / + # box params plus the DN bookkeeping. The GT-less rank gets + # dn_mask_dict=None (what prepare_for_dn returns without boxes) and + # must still reach the DN count collective inside loss(). + num_layers = head.num_decoder_layers + generator = torch.Generator().manual_seed(40) + all_cls_scores = torch.randn(num_layers, 1, _NUM_QUERIES, _NUM_CLASSES, generator=generator) + all_bbox_preds = torch.randn(num_layers, 1, _NUM_QUERIES, 10, generator=generator) + if rank == 0: + pad = 8 + mask_dict = { + "known_lbs_bboxes": ( + gt_labels.repeat(head.scalar), + gt_boxes.repeat(head.scalar, 1), + ), + "known_indices": torch.arange(gt_labels.numel() * head.scalar), + "batch_idx": torch.zeros(gt_labels.numel() * head.scalar, dtype=torch.long), + "map_known_indice": torch.arange(gt_labels.numel() * head.scalar), + "pad_size": pad, + "output_known_lbs_bboxes": ( + torch.randn(num_layers, 1, pad, _NUM_CLASSES, generator=generator), + torch.randn(num_layers, 1, pad, 10, generator=generator), + ), + } + else: + mask_dict = None + outputs = { + "all_cls_scores": all_cls_scores, + "all_bbox_preds": all_bbox_preds, + "dn_mask_dict": mask_dict, + } + + losses = head.loss(outputs, [gt_boxes], [gt_labels]) + + # Both ranks completed (no deadlock) with a uniform key set. + assert "dn_loss_cls" in losses and "dn_loss_bbox" in losses + if rank == 1: + assert losses["dn_loss_cls"].item() == 0.0 + assert losses["dn_loss_bbox"].item() == 0.0 + finally: + dist.destroy_process_group() + + +def test_reduce_mean_count_is_identity_without_process_group() -> None: + value = torch.tensor(3.0) + assert reduce_mean_count(value) is value + # Integer counts are cast to float either way, so the single-GPU dtype + # matches the distributed path. + assert reduce_mean_count(torch.tensor(3)).dtype == torch.float32 + + +def test_reduce_mean_count_averages_counts_across_ranks() -> None: + _spawn(_reduce_mean_count_worker) + + +def test_sharded_loss_matches_single_batch_equivalent() -> None: + _spawn(_sharding_invariance_worker) + + +def test_dn_collective_survives_rank_without_ground_truth() -> None: + _spawn(_dn_mixed_gt_worker) diff --git a/autoware_ml/tests/models/test_streampetr_partial_ignore.py b/autoware_ml/tests/models/test_streampetr_partial_ignore.py new file mode 100644 index 00000000..7003c766 --- /dev/null +++ b/autoware_ml/tests/models/test_streampetr_partial_ignore.py @@ -0,0 +1,258 @@ +"""Tests for StreamPETR partial-ignore and the auxiliary 2D head.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from autoware_ml.losses.detection3d.focal import SigmoidFocalLoss +from autoware_ml.models.detection3d.heads.focal2d import FocalHead2D +from autoware_ml.models.detection3d.heads.streampetr import StreamPETRHead +from autoware_ml.models.detection3d.partial_ignore import ( + mask_ignored_columns, + normalize_status_flags, + resolve_partial_ignore_labels, +) +from autoware_ml.models.detection3d.task_modules.assigners import HungarianAssigner3D +from autoware_ml.models.detection3d.task_modules.bbox_coders import NMSFreeBBoxCoder3D +from autoware_ml.models.detection3d.task_modules.match_costs import ( + BBox3DL1Cost, + ClassificationCost, + IoU3DCost, +) + +CLASS_NAMES = ["car", "truck", "bus", "bicycle", "pedestrian", "traffic_cone", "barrier"] + + +def _build_head(partial_ignore: bool) -> StreamPETRHead: + return StreamPETRHead( + num_classes=7, + in_channels=32, + hidden_dim=32, + num_queries=16, + num_decoder_layers=2, + num_heads=4, + feedforward_channels=64, + memory_len=16, + topk_proposals=4, + num_propagated=4, + with_dn=True, + with_ego_pos=True, + depth_num=8, + LID=True, + position_range=[-12.0, -12.0, -6.0, 12.0, 12.0, 6.0], + scalar=2, + noise_scale=0.5, + dn_weight=1.0, + split=0.5, + use_bottom_center=True, + bbox_coder=NMSFreeBBoxCoder3D( + pc_range=[-10.0, -10.0, -5.0, 10.0, 10.0, 5.0], + post_center_range=[-12.0, -12.0, -6.0, 12.0, 12.0, 6.0], + score_threshold=0.01, + max_num=8, + ), + assigner=HungarianAssigner3D( + cls_cost=ClassificationCost(weight=2.0), + reg_cost=BBox3DL1Cost(weight=0.25), + iou_cost=IoU3DCost(weight=0.0), + ), + point_cloud_range=[-10.0, -10.0, -5.0, 10.0, 10.0, 5.0], + code_weights=[2.0, 2.0] + [1.0] * 8, + class_names=CLASS_NAMES if partial_ignore else None, + partial_ignore_classes=["traffic_cone", "barrier"] if partial_ignore else None, + ) + + +def test_resolve_partial_ignore_labels() -> None: + assert resolve_partial_ignore_labels(CLASS_NAMES, ["traffic_cone", "barrier"]) == [5, 6] + assert resolve_partial_ignore_labels(CLASS_NAMES, None) is None + + +def test_normalize_status_flags_handles_tensors_and_lists() -> None: + assert normalize_status_flags([True, False], 2) == [True, False] + assert normalize_status_flags(torch.tensor([1.0, 0.0]), 2) == [True, False] + assert normalize_status_flags([torch.tensor(False), torch.tensor(True)], 2) == [False, True] + + +def test_normalize_status_flags_rejects_missing_and_mismatched_flags() -> None: + with pytest.raises(ValueError, match="annotation_status is required"): + normalize_status_flags(None, 2) + with pytest.raises(ValueError, match="2 flags for a batch of 3"): + normalize_status_flags([True, False], 3) + + +def test_mask_ignored_columns_zeroes_selected_rows_only() -> None: + weights = torch.ones(4, 7) + mask_ignored_columns(weights, torch.tensor([1, 3]), [5, 6]) + assert torch.all(weights[[1, 3]][:, [5, 6]] == 0.0) + assert torch.all(weights[[0, 2]] == 1.0) + assert torch.all(weights[:, :5] == 1.0) + # No rows selected: a no-op. + mask_ignored_columns(weights, torch.zeros(0, dtype=torch.long), [5, 6]) + + +def test_sigmoid_focal_loss_classwise_weights_zero_masked_columns() -> None: + loss_fn = SigmoidFocalLoss() + logits = torch.randn(4, 7) + targets = torch.zeros(4, 7) + ones = torch.ones(4, 7) + baseline = loss_fn(logits, targets, avg_factor=1.0) + weighted = loss_fn(logits, targets, weights=ones, avg_factor=1.0) + assert torch.allclose(baseline, weighted) + + masked_weights = ones.clone() + masked_weights[:, [5, 6]] = 0.0 + masked = loss_fn(logits, targets, weights=masked_weights, avg_factor=1.0) + manual = loss_fn(logits[:, :5], targets[:, :5], avg_factor=1.0) + assert torch.allclose(masked, manual, atol=1e-6) + + +def test_get_targets_zeroes_ignore_columns_on_every_query_of_unannotated_frames() -> None: + head = _build_head(partial_ignore=True) + num_queries = 16 + cls_logits = torch.randn(1, num_queries, 7) + box_params = torch.randn(1, num_queries, 10) + box_params[..., :3] = 0.0 + gt_boxes = [torch.tensor([[0.0, 0.0, 0.0, 4.0, 2.0, 1.5, 0.1, 0.0, 0.0]], dtype=torch.float32)] + gt_labels = [torch.tensor([0], dtype=torch.long)] + + targets = head._get_targets(cls_logits, box_params, gt_boxes, gt_labels, [False])[0] + assert targets.label_weights is not None + assert targets.label_weights.shape == (num_queries, 7) + + # Every query of the un-annotated frame loses only the cone/barrier columns. + assert torch.all(targets.label_weights[:, [5, 6]] == 0.0) + assert torch.all(targets.label_weights[:, :5] == 1.0) + + # A fully annotated batch skips the classwise-weight tensors entirely. + annotated = head._get_targets(cls_logits, box_params, gt_boxes, gt_labels, [True])[0] + assert annotated.label_weights is None + + # In a mixed batch every sample gets a weights tensor; the annotated one + # stays all-ones so the stacked tensor is uniform. + mixed = head._get_targets( + cls_logits.repeat(2, 1, 1), + box_params.repeat(2, 1, 1), + gt_boxes * 2, + gt_labels * 2, + [True, False], + ) + assert torch.all(mixed[0].label_weights == 1.0) + assert torch.all(mixed[1].label_weights[:, [5, 6]] == 0.0) + + +def test_dn_label_weights_mask_all_rows_of_unannotated_samples() -> None: + head = _build_head(partial_ignore=True) + cls_scores = torch.randn(6, 7) + known_labels = torch.tensor([0, 7, 7, 1, 7, 5]) + known_bids = torch.tensor([0, 0, 1, 1, 1, 0]) + weights = head._dn_label_weights(cls_scores, known_labels, known_bids, [True, False]) + assert weights is not None + # Every row of sample 1 (indices 2, 3, 4) is masked on columns 5/6. + assert torch.all(weights[[2, 3, 4]][:, [5, 6]] == 0.0) + assert torch.all(weights[[2, 3, 4]][:, :5] == 1.0) + # Rows of the annotated sample 0 stay fully weighted. + assert torch.all(weights[[0, 1, 5]] == 1.0) + + assert head._dn_label_weights(cls_scores, known_labels, known_bids, [True, True]) is None + + +def test_head_loss_wiring_applies_annotation_status() -> None: + """head.loss must thread the status flags through to the focal loss.""" + torch.manual_seed(0) + head = _build_head(partial_ignore=True).eval() + outputs = { + "all_cls_scores": torch.randn(2, 1, 16, 7), + "all_bbox_preds": torch.rand(2, 1, 16, 10), + "dn_mask_dict": None, + } + gt_boxes = [torch.tensor([[0.0, 0.0, 0.0, 4.0, 2.0, 1.5, 0.1, 0.0, 0.0]], dtype=torch.float32)] + gt_labels = [torch.tensor([0], dtype=torch.long)] + + annotated = head.loss(outputs, gt_boxes, gt_labels, annotation_status=[True]) + ignored = head.loss(outputs, gt_boxes, gt_labels, annotation_status=[False]) + # Masking removes non-negative background terms from the cls loss only. + assert ignored["loss_cls"] < annotated["loss_cls"] + assert torch.allclose(ignored["loss_bbox"], annotated["loss_bbox"]) + + with pytest.raises(ValueError, match="annotation_status is required"): + head.loss(outputs, gt_boxes, gt_labels) + + # Heads without partial-ignore never require the flags. + plain_head = _build_head(partial_ignore=False).eval() + plain_losses = plain_head.loss(outputs, gt_boxes, gt_labels) + assert torch.isfinite(plain_losses["loss"]) + + +def test_focal_head_2d_forward_and_loss_with_partial_ignore() -> None: + torch.manual_seed(0) + head = FocalHead2D( + num_classes=7, + in_channels=32, + embed_dims=32, + stride=16, + class_names=CLASS_NAMES, + partial_ignore_classes=["traffic_cone", "barrier"], + ) + batch_size, num_cams = 2, 2 + img_features = torch.randn(batch_size, num_cams, 32, 6, 10) + outputs = head(img_features, image_height=96, image_width=160) + assert outputs["enc_cls_scores"].shape == (4, 60, 7) + assert outputs["enc_bbox_preds"].shape == (4, 60, 4) + assert outputs["pred_centers2d"].shape == (4, 60, 2) + assert outputs["centerness"].shape == (4, 60, 1) + + gt_boxes = np.array([[10.0, 10.0, 60.0, 60.0]], dtype=np.float32) + gt_centers = np.array([[35.0, 35.0]], dtype=np.float32) + gt_labels = np.array([5], dtype=np.int64) + empty_boxes = np.zeros((0, 4), dtype=np.float32) + empty_centers = np.zeros((0, 2), dtype=np.float32) + empty_labels = np.zeros((0,), dtype=np.int64) + annotations = { + "gt_bboxes_2d": [[gt_boxes, empty_boxes], [empty_boxes, empty_boxes]], + "gt_labels_2d": [[gt_labels, empty_labels], [empty_labels, empty_labels]], + "centers_2d": [[gt_centers, empty_centers], [empty_centers, empty_centers]], + } + + losses = head.loss(outputs, annotation_status=[True, False], **annotations) + for key in ( + "loss_cls2d", + "loss_bbox2d", + "loss_iou2d", + "loss_centers2d", + "loss_centerness2d", + ): + assert torch.isfinite(losses[key]), key + + # Masking removes strictly positive background terms of sample 1's images + # and touches nothing else. + unmasked = head.loss(outputs, annotation_status=[True, True], **annotations) + assert losses["loss_cls2d"] < unmasked["loss_cls2d"] + for key in ("loss_bbox2d", "loss_iou2d", "loss_centers2d", "loss_centerness2d"): + assert torch.allclose(losses[key], unmasked[key]), key + + with pytest.raises(ValueError, match="annotation_status is required"): + head.loss(outputs, **annotations) + + +def test_focal_head_2d_partial_ignore_class_weights_select_exact_rows() -> None: + head = FocalHead2D( + num_classes=7, + in_channels=32, + embed_dims=32, + class_names=CLASS_NAMES, + partial_ignore_classes=["traffic_cone", "barrier"], + ) + tokens_per_image = 3 + flat_scores = torch.randn(4 * tokens_per_image, 7) + weights = head._partial_ignore_class_weights( + flat_scores, [True, False, False, True], tokens_per_image + ) + assert weights is not None + expected = torch.ones_like(flat_scores) + expected[3:9, 5:7] = 0.0 + assert torch.equal(weights, expected) + + assert head._partial_ignore_class_weights(flat_scores, [True] * 4, tokens_per_image) is None