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 <yihsiang.fang@tier4.jp>

Author: vividf

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

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])(