From 2bdd538cf71648fcd20f4d224cad2f673fd31964 Mon Sep 17 00:00:00 2001 From: vividf Date: Wed, 2 Sep 2026 17:15:31 +0900 Subject: [PATCH] feat(streampetr): project 2D annotations and fold augmentations into ego poses LoadAnnotations2DFromBoxes3D projects the augmented 3D boxes onto every camera for the auxiliary 2D head; the camera geometric augmentations fold their transforms into the ego poses so the temporal memory warp stays consistent; train-time camera-order shuffling follows the reference recipe. Signed-off-by: vividf --- autoware_ml/tests/transforms/test_camera.py | 208 ++++++++++++++++++ .../tests/transforms/test_geometry3d.py | 62 ++++++ .../transforms/camera/annotations2d.py | 191 ++++++++++++++++ autoware_ml/transforms/camera/geometry.py | 8 +- autoware_ml/transforms/camera/loading.py | 20 +- .../transforms/camera_lidar/geometry.py | 8 +- autoware_ml/transforms/geometry3d.py | 27 +++ 7 files changed, 518 insertions(+), 6 deletions(-) create mode 100644 autoware_ml/transforms/camera/annotations2d.py diff --git a/autoware_ml/tests/transforms/test_camera.py b/autoware_ml/tests/transforms/test_camera.py index 320b7f5d..0de32453 100644 --- a/autoware_ml/tests/transforms/test_camera.py +++ b/autoware_ml/tests/transforms/test_camera.py @@ -18,6 +18,7 @@ import numpy.typing as npt import pytest +from autoware_ml.transforms.camera.annotations2d import LoadAnnotations2DFromBoxes3D from autoware_ml.transforms.camera.distortion import UndistortImage from autoware_ml.transforms.camera.masking import GridMask from autoware_ml.transforms.camera.normalize import NormalizeMultiviewImage @@ -281,3 +282,210 @@ def test_grid_mask_handles_chw_stack() -> None: assert output["img"].shape == images.shape assert (output["img"] == 0).any() + + +def test_load_annotations_2d_treats_box_z_as_gravity_center() -> None: + """The projected 2D box and center must straddle z, not sit above it. + + ``gt_boxes`` stores the gravity center. Treating that z as the bottom face + (building corners from z upward, or adding dz/2 before projecting the + center) silently lifts every 2D auxiliary target by half an object height, + which is invisible in the loss but wrong everywhere. + """ + # Pinhole camera: fx = fy = 100, principal point at (50, 50). + cam2img = np.eye(4) + cam2img[0, 0] = cam2img[1, 1] = 100.0 + cam2img[0, 2] = cam2img[1, 2] = 50.0 + # Lidar (x forward, y left, z up) -> camera (x right, y down, z forward). + lidar2cam = np.array( + [ + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + # One 2 m cube centered on the optical axis, 10 m ahead. + gt_boxes = np.array([[10.0, 0.0, 0.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0]], dtype=np.float32) + + output = LoadAnnotations2DFromBoxes3D()( + { + "img": np.zeros((1, 3, 100, 100), dtype=np.float32), + "gt_boxes": gt_boxes, + "gt_labels": np.array([0]), + "lidar2cam": [lidar2cam], + "camera_intrinsics": [cam2img], + } + ) + + center = output["centers_2d"][0][0] + x1, y1, x2, y2 = output["gt_bboxes_2d"][0][0] + # z = 0 is on the optical axis, so the center projects to the principal + # point. The bottom-face reading would put it at y = 40. + assert center == pytest.approx([50.0, 50.0], abs=1e-3) + # The cube spans 9-11 m in depth, so its near face projects largest and + # sets the extent: 100 * 1 / 9 either side of the principal point. + extent = 100.0 / 9.0 + assert (y1, y2) == pytest.approx((50.0 - extent, 50.0 + extent), abs=1e-3) + assert (x1, x2) == pytest.approx((50.0 - extent, 50.0 + extent), abs=1e-3) + # The center sits at the middle of the box, not on its top edge. + assert center[1] == pytest.approx((y1 + y2) / 2, abs=1e-3) + + +def _pinhole_setup(image_width: int = 160, image_height: int = 96) -> dict: + """One forward-looking pinhole camera in a lidar frame (x fwd, y left, z up).""" + cam2img = np.eye(4) + cam2img[0, 0] = cam2img[1, 1] = 100.0 + cam2img[0, 2] = image_width / 2 + cam2img[1, 2] = image_height / 2 + lidar2cam = np.array( + [ + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + return { + "img": np.zeros((1, 3, image_height, image_width), dtype=np.float32), + "gt_labels": np.array([0]), + "lidar2cam": [lidar2cam], + "camera_intrinsics": [cam2img], + } + + +def test_load_annotations_2d_projects_box_in_front_of_camera() -> None: + input_dict = _pinhole_setup() + input_dict["gt_boxes"] = np.array( + [[8.0, 0.0, -1.0, 4.0, 2.0, 1.5, 0.0, 0.0, 0.0]], dtype=np.float32 + ) + result = LoadAnnotations2DFromBoxes3D()(input_dict) + assert result["gt_bboxes_2d"][0].shape == (1, 4) + assert result["centers_2d"][0].shape == (1, 2) + assert result["gt_labels_2d"][0].tolist() == [0] + x1, y1, x2, y2 = result["gt_bboxes_2d"][0][0] + assert 0 <= x1 < x2 <= 160 + assert 0 <= y1 < y2 <= 96 + + +def test_load_annotations_2d_drops_box_behind_camera() -> None: + input_dict = _pinhole_setup() + input_dict["gt_boxes"] = np.array( + [[-10.0, 0.0, 0.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0]], dtype=np.float32 + ) + result = LoadAnnotations2DFromBoxes3D()(input_dict) + assert result["gt_bboxes_2d"][0].shape == (0, 4) + assert result["centers_2d"][0].shape == (0, 2) + assert result["gt_labels_2d"][0].shape == (0,) + + +def test_load_annotations_2d_drops_box_whose_center_leaves_the_image() -> None: + """Corners still visible but the projected center is off-image -> dropped. + + Matches the reference recipe: a clamped center could land outside its own + clipped box and would distort the center-based 2D assignment. + """ + input_dict = _pinhole_setup() + # Wide box far to the left: some corners project inside, the center at + # y = 9 m projects to x = 80 - 100 * 9/10 = -10 (outside). + input_dict["gt_boxes"] = np.array( + [[10.0, 9.0, 0.0, 2.0, 6.0, 2.0, 0.0, 0.0, 0.0]], dtype=np.float32 + ) + result = LoadAnnotations2DFromBoxes3D()(input_dict) + assert result["gt_bboxes_2d"][0].shape == (0, 4) + + +def test_load_annotations_2d_clips_partially_visible_box_to_the_image() -> None: + input_dict = _pinhole_setup() + # Center projects at x = 80 - 100 * 6.5/10 = 15 (inside); the near-left + # corners project far outside the left edge and must be clipped to 0. + input_dict["gt_boxes"] = np.array( + [[10.0, 6.5, 0.0, 2.0, 6.0, 2.0, 0.0, 0.0, 0.0]], dtype=np.float32 + ) + result = LoadAnnotations2DFromBoxes3D()(input_dict) + assert result["gt_bboxes_2d"][0].shape == (1, 4) + x1, y1, x2, y2 = result["gt_bboxes_2d"][0][0] + assert x1 == 0.0 + assert 0 < x2 <= 160 + center = result["centers_2d"][0][0] + assert 0 <= center[0] < 160 + + +def test_load_annotations_2d_assigns_boxes_per_camera() -> None: + input_dict = _pinhole_setup() + # Second camera looks backward (rotate lidar->cam by 180 deg around z). + backward = np.array( + [ + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + input_dict["lidar2cam"] = [input_dict["lidar2cam"][0], backward] + input_dict["camera_intrinsics"] = [input_dict["camera_intrinsics"][0]] * 2 + input_dict["gt_boxes"] = np.array( + [ + [10.0, 0.0, 0.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0], + [-10.0, 0.0, 0.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + input_dict["gt_labels"] = np.array([0, 1]) + result = LoadAnnotations2DFromBoxes3D()(input_dict) + assert result["gt_labels_2d"][0].tolist() == [0] + assert result["gt_labels_2d"][1].tolist() == [1] + + +def test_load_annotations_2d_handles_empty_and_invalid_gt() -> None: + input_dict = _pinhole_setup() + input_dict["gt_boxes"] = np.zeros((0, 9), dtype=np.float32) + input_dict["gt_labels"] = np.zeros((0,), dtype=np.int64) + result = LoadAnnotations2DFromBoxes3D()(input_dict) + assert result["gt_bboxes_2d"][0].shape == (0, 4) + assert result["centers_2d"][0].shape == (0, 2) + + input_dict["gt_boxes"] = np.zeros((3, 5), dtype=np.float32) + input_dict["gt_labels"] = np.zeros((3,), dtype=np.int64) + with pytest.raises(ValueError, match="gt_boxes"): + LoadAnnotations2DFromBoxes3D()(input_dict) + + +def test_multiview_loader_shuffle_order_keeps_sample_consistent(tmp_path) -> None: + import random + + import cv2 + + from autoware_ml.transforms.camera.loading import LoadMultiViewImagesFromFiles + + camera_order = [f"CAM_{i}" for i in range(5)] + images_meta = {} + for index, name in enumerate(camera_order): + path = tmp_path / f"{name}.png" + cv2.imwrite(str(path), np.full((4, 6, 3), index * 10, dtype=np.uint8)) + intrinsics = np.eye(3, dtype=np.float32) * (index + 1) + images_meta[name] = { + "img_path": str(path), + "cam2img": intrinsics, + "lidar2cam": np.eye(4, dtype=np.float32), + } + + loader = LoadMultiViewImagesFromFiles(normalize_to_unit=False, shuffle_order=True) + random.seed(3) + shuffled_seen = False + for _ in range(8): + out = loader({"images": images_meta, "camera_order": camera_order}) + names = out["camera_names"] + assert sorted(names) == sorted(camera_order) + if names != camera_order: + shuffled_seen = True + for position, name in enumerate(names): + index = camera_order.index(name) + # Image content and intrinsics must follow the shuffled order. + assert float(out["img"][position].mean()) == index * 10 + assert out["camera_intrinsics"][position][0, 0] == index + 1 + assert shuffled_seen + + fixed_loader = LoadMultiViewImagesFromFiles(normalize_to_unit=False) + out = fixed_loader({"images": images_meta, "camera_order": camera_order}) + assert out["camera_names"] == camera_order diff --git a/autoware_ml/tests/transforms/test_geometry3d.py b/autoware_ml/tests/transforms/test_geometry3d.py index f11655ce..53e08512 100644 --- a/autoware_ml/tests/transforms/test_geometry3d.py +++ b/autoware_ml/tests/transforms/test_geometry3d.py @@ -26,6 +26,7 @@ import numpy as np import pytest +from autoware_ml.transforms import geometry3d as g3d from autoware_ml.transforms.camera import geometry as cam from autoware_ml.transforms.camera_lidar import geometry as cam_lidar from autoware_ml.transforms.point_cloud import geometry as pc @@ -102,6 +103,67 @@ def test_random_flip_shared_math_matches_across_namespaces(seed: int) -> None: assert np.allclose(out_cam["points"], _sample()["points"]) +def _ego_pose() -> np.ndarray: + # Non-trivial lidar->global pose: yaw + translation. + yaw = 0.7 + pose = np.eye(4, dtype=np.float64) + pose[:2, :2] = [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]] + pose[:3, 3] = [10.0, -4.0, 1.2] + return pose + + +@pytest.mark.parametrize("transform_cls", ["rot_scale_trans", "flip"]) +def test_camera_transforms_fold_augmentation_into_ego_poses(transform_cls: str) -> None: + """After augmentation, ego_pose must map *augmented* lidar coords to the + same global point as before (StreamPETR warps temporal memory through it).""" + sample = _sample() + sample["ego_pose"] = _ego_pose() + sample["ego_pose_inv"] = np.linalg.inv(sample["ego_pose"]) + point_lidar = np.array([3.0, -2.0, 0.5, 1.0]) + point_global = sample["ego_pose"] @ point_lidar + + np.random.seed(3) + if transform_cls == "rot_scale_trans": + out = cam.GlobalRotScaleTrans( + rot_range=[-0.5, 0.5], scale_ratio_range=[0.9, 1.1], translation_std=[0.5, 0.5, 0.2] + )(sample) + augmentation = out["global_aug_matrix"] + else: + out = cam.RandomFlip3D(flip_ratio_bev_horizontal=1.0, flip_ratio_bev_vertical=1.0)(sample) + augmentation = out["bev_flip_matrix"] + + point_augmented = augmentation @ point_lidar + assert np.allclose(out["ego_pose"] @ point_augmented, point_global, atol=1e-5) + assert np.allclose(out["ego_pose_inv"] @ point_global, point_augmented, atol=1e-5) + # The pair stays mutually inverse. + assert np.allclose(out["ego_pose"] @ out["ego_pose_inv"], np.eye(4), atol=1e-5) + + +def test_camera_and_camera_lidar_fold_ego_poses_identically() -> None: + def _sample_with_pose() -> dict: + sample = _sample() + sample["ego_pose"] = _ego_pose() + sample["ego_pose_inv"] = np.linalg.inv(sample["ego_pose"]) + return sample + + kwargs = dict(rot_range=[-0.5, 0.5], scale_ratio_range=[0.9, 1.1]) + np.random.seed(11) + out_cam = cam.GlobalRotScaleTrans(**kwargs)(_sample_with_pose()) + np.random.seed(11) + out_cl = cam_lidar.GlobalRotScaleTrans(**kwargs)(_sample_with_pose()) + assert np.allclose(out_cam["ego_pose"], out_cl["ego_pose"]) + assert np.allclose(out_cam["ego_pose_inv"], out_cl["ego_pose_inv"]) + + +def test_transform_boxes_scales_velocity_with_the_world() -> None: + input_dict = { + "gt_boxes": np.array([[1.0, 2.0, 0.0, 4.0, 2.0, 1.0, 0.3, 1.5, -0.5]], dtype=np.float32) + } + identity = np.eye(3, dtype=np.float32) + g3d.transform_boxes(input_dict, identity, 0.0, 2.0, np.zeros((1, 3), dtype=np.float32)) + assert np.allclose(input_dict["gt_boxes"][0, 7:9], [3.0, -1.0]) + + def test_point_cloud_requires_a_point_representation() -> None: with pytest.raises(KeyError): pc.GlobalRotScaleTrans(rot_range=[0.1, 0.1], scale_ratio_range=[1.0, 1.0])( diff --git a/autoware_ml/transforms/camera/annotations2d.py b/autoware_ml/transforms/camera/annotations2d.py new file mode 100644 index 00000000..04982e56 --- /dev/null +++ b/autoware_ml/transforms/camera/annotations2d.py @@ -0,0 +1,191 @@ +# 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. + +"""Derive per-camera 2D annotations from 3D boxes for auxiliary supervision. + +The transform projects the (already augmented) 3D ground-truth boxes onto +every camera and emits per-camera 2D boxes, projected gravity centers, and +labels. It must therefore run after all geometric augmentations +(``ResizeCropFlipRotImage``, ``GlobalRotScaleTrans``, ``PadMultiViewImage``) +so the projection matrices and the pixels agree. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from autoware_ml.transforms.base import BaseTransform +from autoware_ml.transforms.camera.utils import as_hwc_image_list + + +def _boxes3d_corners(boxes: np.ndarray) -> np.ndarray: + """Compute the 8 corners of gravity-center 3D boxes. + + A numpy sibling of :meth:`autoware_ml.geometry.bbox_3d.LiDARBBox3D.corners` + (which is torch-based) for use inside numpy transform pipelines. + + Args: + boxes: Boxes ``(N, >=7)`` as ``[x, y, z_center, dx, dy, dz, yaw, ...]``, + where ``z_center`` is the gravity center — the convention this + repo's 3D boxes carry throughout. + + Returns: + Corners with shape ``(N, 8, 3)``. + """ + if boxes.shape[0] == 0: + return np.zeros((0, 8, 3), dtype=np.float32) + dims = boxes[:, 3:6] + signs = np.array( + [[dx, dy, dz] for dx in (-0.5, 0.5) for dy in (-0.5, 0.5) for dz in (-0.5, 0.5)], + dtype=np.float32, + ) + corners = dims[:, None, :] * signs[None, :, :] + yaw = boxes[:, 6] + cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw) + rotated_x = corners[..., 0] * cos_yaw[:, None] - corners[..., 1] * sin_yaw[:, None] + rotated_y = corners[..., 0] * sin_yaw[:, None] + corners[..., 1] * cos_yaw[:, None] + corners = np.stack([rotated_x, rotated_y, corners[..., 2]], axis=-1) + return corners + boxes[:, None, :3] + + +def _project_points( + points: np.ndarray, lidar2cam: np.ndarray, cam2img: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Project lidar-frame points into image pixels. + + Returns: + Pixel coordinates ``(N, 2)`` and an in-front-of-camera mask ``(N,)``. + """ + homogeneous = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1) + points_cam = (lidar2cam @ homogeneous.T).T + valid = points_cam[:, 2] > 0 + projected = (cam2img[:3, :3] @ points_cam[:, :3].T).T + projected = projected / np.maximum(projected[:, 2:3], 1e-6) + return projected[:, :2], valid + + +class LoadAnnotations2DFromBoxes3D(BaseTransform): + """Project 3D ground-truth boxes onto every camera as 2D annotations. + + Outputs (per camera, one entry per visible box): + ``gt_bboxes_2d``: ``(N, 4)`` clipped ``(x1, y1, x2, y2)`` pixel boxes. + ``centers_2d``: ``(N, 2)`` projected gravity centers in pixels. + ``gt_labels_2d``: ``(N,)`` class labels. + """ + + _required_keys = ["img", "gt_boxes", "gt_labels", "lidar2cam", "camera_intrinsics"] + + def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: + """Compute per-camera 2D annotations from the 3D boxes.""" + image_list, _ = as_hwc_image_list(input_dict["img"]) + image_height, image_width = int(image_list[0].shape[0]), int(image_list[0].shape[1]) + gt_boxes = np.asarray(input_dict["gt_boxes"], dtype=np.float32) + if gt_boxes.size == 0: + gt_boxes = gt_boxes.reshape(0, 9) + if gt_boxes.ndim != 2 or gt_boxes.shape[1] < 7: + raise ValueError( + f"gt_boxes must be (N, >=7) [x, y, z, dx, dy, dz, yaw, ...], " + f"got shape {gt_boxes.shape}." + ) + gt_labels = np.asarray(input_dict["gt_labels"]).reshape(-1) + corners = _boxes3d_corners(gt_boxes) + gravity_centers = gt_boxes[:, :3].copy() + + all_bboxes, all_centers, all_labels = [], [], [] + num_cams = len(input_dict["lidar2cam"]) + for camera_index in range(num_cams): + lidar2cam = np.asarray(input_dict["lidar2cam"][camera_index], dtype=np.float64) + cam2img = np.asarray(input_dict["camera_intrinsics"][camera_index], dtype=np.float64) + bboxes, centers, labels = self._project_camera( + corners, + gravity_centers, + gt_labels, + lidar2cam, + cam2img, + image_height, + image_width, + ) + all_bboxes.append(bboxes) + all_centers.append(centers) + all_labels.append(labels) + + return { + "gt_bboxes_2d": all_bboxes, + "centers_2d": all_centers, + "gt_labels_2d": all_labels, + } + + def _project_camera( + self, + corners: np.ndarray, + gravity_centers: np.ndarray, + gt_labels: np.ndarray, + lidar2cam: np.ndarray, + cam2img: np.ndarray, + image_height: int, + image_width: int, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + num_boxes = corners.shape[0] + if num_boxes == 0: + return self._empty_annotations() + + points = np.concatenate([corners.reshape(-1, 3), gravity_centers], axis=0) + pixels, in_front = _project_points(points, lidar2cam, cam2img) + corner_pixels = pixels[: num_boxes * 8].reshape(num_boxes, 8, 2) + corner_front = in_front[: num_boxes * 8].reshape(num_boxes, 8) + center_pixels = pixels[num_boxes * 8 :] + center_front = in_front[num_boxes * 8 :] + + # A box qualifies when its gravity center projects inside this camera's + # image (matching the reference recipe, which drops boxes whose center + # leaves the crop — a clamped center could land outside its own box) + # and at least one corner is in front of the camera. Corners behind + # the camera are excluded from the 2D extent: their clamped-depth + # projection is meaningless. + candidates = np.nonzero( + center_front + & corner_front.any(axis=1) + & (center_pixels[:, 0] >= 0) + & (center_pixels[:, 0] < image_width) + & (center_pixels[:, 1] >= 0) + & (center_pixels[:, 1] < image_height) + )[0] + if candidates.size == 0: + return self._empty_annotations() + + masked = np.where(corner_front[candidates, :, None], corner_pixels[candidates], np.nan) + x_min = np.clip(np.nanmin(masked[:, :, 0], axis=1), 0, image_width) + x_max = np.clip(np.nanmax(masked[:, :, 0], axis=1), 0, image_width) + y_min = np.clip(np.nanmin(masked[:, :, 1], axis=1), 0, image_height) + y_max = np.clip(np.nanmax(masked[:, :, 1], axis=1), 0, image_height) + visible = (x_min < x_max) & (y_min < y_max) + if not visible.any(): + return self._empty_annotations() + kept = candidates[visible] + + bboxes = np.stack( + [x_min[visible], y_min[visible], x_max[visible], y_max[visible]], axis=1 + ).astype(np.float32) + centers = center_pixels[kept].astype(np.float32) + return bboxes, centers, gt_labels[kept].astype(np.int64) + + @staticmethod + def _empty_annotations() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + return ( + np.zeros((0, 4), dtype=np.float32), + np.zeros((0, 2), dtype=np.float32), + np.zeros((0,), dtype=np.int64), + ) diff --git a/autoware_ml/transforms/camera/geometry.py b/autoware_ml/transforms/camera/geometry.py index 8e23c931..121637c1 100644 --- a/autoware_ml/transforms/camera/geometry.py +++ b/autoware_ml/transforms/camera/geometry.py @@ -63,7 +63,9 @@ def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: if flip_x: g3d.flip_boxes(input_dict, axis=0) flip = g3d.flip_matrix(flip_x, flip_y) - g3d.update_camera_matrices(input_dict, np.linalg.inv(flip)) + flip_inv = np.linalg.inv(flip) + g3d.update_camera_matrices(input_dict, flip_inv) + g3d.update_ego_poses(input_dict, flip, flip_inv) input_dict["bev_flip_matrix"] = flip return input_dict @@ -100,6 +102,8 @@ def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: ) g3d.transform_boxes(input_dict, rotation, rotation_angle, scale, translation) augmentation = g3d.rot_scale_trans_matrix(rotation, scale, translation) - g3d.update_camera_matrices(input_dict, np.linalg.inv(augmentation)) + augmentation_inv = np.linalg.inv(augmentation) + g3d.update_camera_matrices(input_dict, augmentation_inv) + g3d.update_ego_poses(input_dict, augmentation, augmentation_inv) input_dict["global_aug_matrix"] = augmentation return input_dict diff --git a/autoware_ml/transforms/camera/loading.py b/autoware_ml/transforms/camera/loading.py index 79a96daa..d6709927 100644 --- a/autoware_ml/transforms/camera/loading.py +++ b/autoware_ml/transforms/camera/loading.py @@ -2,6 +2,7 @@ from __future__ import annotations +import random from typing import Any import cv2 @@ -49,15 +50,27 @@ class LoadMultiViewImagesFromFiles(BaseTransform): _required_keys = ["images", "camera_order"] - def __init__(self, *, to_float32: bool = True, normalize_to_unit: bool = True) -> None: + def __init__( + self, + *, + to_float32: bool = True, + normalize_to_unit: bool = True, + shuffle_order: bool = False, + ) -> None: """Initialize the LoadMultiViewImagesFromFiles transform. Args: to_float32: Whether to cast images to ``float32``. normalize_to_unit: Whether to divide pixel values by ``255``. + shuffle_order: Shuffle the camera order per sample (train-time + regularization for camera-order-agnostic models). Every + emitted per-camera array follows the shuffled order, so the + sample stays internally consistent. Enable only in training + pipelines. """ self.to_float32 = to_float32 self.normalize_to_unit = normalize_to_unit + self.shuffle_order = shuffle_order def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: """Load images and camera matrices for all configured views. @@ -73,7 +86,10 @@ def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: lidar2cam = [] lidar2img = [] camera_names = [] - for camera_name in input_dict["camera_order"]: + camera_order = list(input_dict["camera_order"]) + if self.shuffle_order: + random.shuffle(camera_order) + for camera_name in camera_order: camera_info = input_dict["images"].get(camera_name) if camera_info is None: raise ValueError( diff --git a/autoware_ml/transforms/camera_lidar/geometry.py b/autoware_ml/transforms/camera_lidar/geometry.py index d806d21d..9158482b 100644 --- a/autoware_ml/transforms/camera_lidar/geometry.py +++ b/autoware_ml/transforms/camera_lidar/geometry.py @@ -67,7 +67,9 @@ def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: g3d.flip_normal(input_dict, axis=0) g3d.flip_boxes(input_dict, axis=0) flip = g3d.flip_matrix(flip_x, flip_y) - g3d.update_camera_matrices(input_dict, np.linalg.inv(flip)) + flip_inv = np.linalg.inv(flip) + g3d.update_camera_matrices(input_dict, flip_inv) + g3d.update_ego_poses(input_dict, flip, flip_inv) input_dict["bev_flip_matrix"] = flip return input_dict @@ -107,6 +109,8 @@ def transform(self, input_dict: dict[str, Any]) -> dict[str, Any]: g3d.transform_normal(input_dict, rotation) g3d.transform_boxes(input_dict, rotation, rotation_angle, scale, translation) augmentation = g3d.rot_scale_trans_matrix(rotation, scale, translation) - g3d.update_camera_matrices(input_dict, np.linalg.inv(augmentation)) + augmentation_inv = np.linalg.inv(augmentation) + g3d.update_camera_matrices(input_dict, augmentation_inv) + g3d.update_ego_poses(input_dict, augmentation, augmentation_inv) input_dict["global_aug_matrix"] = augmentation return input_dict diff --git a/autoware_ml/transforms/geometry3d.py b/autoware_ml/transforms/geometry3d.py index 1a80a625..9fbdfcf9 100644 --- a/autoware_ml/transforms/geometry3d.py +++ b/autoware_ml/transforms/geometry3d.py @@ -244,6 +244,33 @@ def flip_boxes(input_dict: dict[str, Any], axis: int) -> None: input_dict["gt_boxes"] = boxes +def update_ego_poses( + input_dict: dict[str, Any], + aug: npt.NDArray[np.float32], + aug_inv: npt.NDArray[np.float32], +) -> None: + """Fold a lidar-frame augmentation into ``ego_pose`` / ``ego_pose_inv``. + + Streaming temporal models (StreamPETR) warp the previous frame's memory + with ``ego_pose_inv(t) @ ego_pose(t-1)``, so after augmenting frame ``t`` + the pose pair must map between the *augmented* lidar frame and the global + frame — otherwise the memory is misaligned by the sampled augmentation: + + ``ego_pose ← ego_pose @ aug_inv``, ``ego_pose_inv ← aug @ ego_pose_inv``. + + No-op for samples without ego poses (single-frame models). + """ + if "ego_pose" not in input_dict: + return + ego_pose = np.asarray(input_dict["ego_pose"]) + input_dict["ego_pose"] = (ego_pose @ aug_inv.astype(ego_pose.dtype)).astype(ego_pose.dtype) + if "ego_pose_inv" in input_dict: + ego_pose_inv = np.asarray(input_dict["ego_pose_inv"]) + input_dict["ego_pose_inv"] = (aug.astype(ego_pose_inv.dtype) @ ego_pose_inv).astype( + ego_pose_inv.dtype + ) + + def update_camera_matrices(input_dict: dict[str, Any], aug_inv: npt.NDArray[np.float32]) -> None: """Keep camera projection consistent after a lidar-space transform.