diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index 89e37a9519d..13679a52c11 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -20,6 +20,9 @@ on: - dimos/protocol/tf/** - dimos/memory/** - dimos/mapping/occupancy/** + - dimos/mapping/pointclouds/occupancy.py + - dimos/mapping/costmapper.py + - dimos/mapping/test_generated_costmapper.py - dimos/protocol/cdr_mcap.py - docs/usage/transforms.md - dimos/hardware/sensors/camera/** @@ -227,3 +230,12 @@ jobs: run: | .native-review/bin/python -I -c 'from dimos.mapping.occupancy import operations, gradient, inflation, extrude_occupancy; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (operations, gradient, inflation, extrude_occupancy)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/mapping/occupancy/test_generated_operations.py" MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/occupancy/operations.py dimos/mapping/occupancy/gradient.py dimos/mapping/occupancy/inflation.py dimos/mapping/occupancy/extrude_occupancy.py + - name: Verify generated pointcloud costmaps and real offline Module ports + if: runner.os == 'Linux' + env: + NUMBA_NUM_THREADS: '2' + OMP_NUM_THREADS: '2' + run: | + .native-review/bin/python -m pip install 'numba==0.68.0' + .native-review/bin/python -I -c 'import dimos.mapping.costmapper as module; import dimos.mapping.pointclouds.occupancy as algorithms; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (module, algorithms)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -o 'markers=self_hosted: optional project fixture case' --strict-markers -m 'not self_hosted' -q "$PWD/dimos/mapping/test_generated_costmapper.py" + MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/pointclouds/occupancy.py dimos/mapping/costmapper.py diff --git a/dimos/mapping/costmapper.py b/dimos/mapping/costmapper.py index 31b94b164f5..62a4aa7790e 100644 --- a/dimos/mapping/costmapper.py +++ b/dimos/mapping/costmapper.py @@ -27,8 +27,9 @@ HeightCostConfig, OccupancyConfig, ) -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.msgs.occupancy import occupancy_view from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -90,19 +91,22 @@ def stop(self) -> None: def _calculate_costmap(self, msg: PointCloud2) -> OccupancyGrid: occupancy_function = OCCUPANCY_ALGOS[self.config.algo] grid = occupancy_function(msg, **asdict(self.config.config)) + if not isinstance(grid, OccupancyGrid): + raise TypeError("CostMapper algorithms must return generated OccupancyGrid values") self._apply_initial_safe_radius(grid) return grid def _apply_initial_safe_radius(self, grid: OccupancyGrid) -> None: radius_meters = self.config.initial_safe_radius_meters - if radius_meters <= 0 or grid.grid.size == 0: + if radius_meters <= 0 or grid.data.size == 0: return - resolution = grid.resolution - origin_x = grid.origin.position.x - origin_y = grid.origin.position.y + resolution = grid.info.resolution + origin_x = grid.info.origin.position.x + origin_y = grid.info.origin.position.y - rows, columns = np.ogrid[: grid.grid.shape[0], : grid.grid.shape[1]] + cells = occupancy_view(grid).copy() + rows, columns = np.ogrid[: cells.shape[0], : cells.shape[1]] cell_world_x = columns * resolution + origin_x cell_world_y = rows * resolution + origin_y distance_squared_meters = cell_world_x**2 + cell_world_y**2 @@ -111,4 +115,5 @@ def _apply_initial_safe_radius(self, grid: OccupancyGrid) -> None: # the disc. Avoids floating-point boundary flakiness from radius/resolution. effective_radius_meters = radius_meters + resolution * 0.5 safe_mask = distance_squared_meters <= effective_radius_meters**2 - grid.grid[safe_mask] = 0 + cells[safe_mask] = 0 + grid.data = cells.reshape(-1) diff --git a/dimos/mapping/pointclouds/occupancy.py b/dimos/mapping/pointclouds/occupancy.py index d723b901543..ae9700358f1 100644 --- a/dimos/mapping/pointclouds/occupancy.py +++ b/dimos/mapping/pointclouds/occupancy.py @@ -14,14 +14,21 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload import numpy as np +from numpy.typing import NDArray from scipy import ndimage -from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 as CdrCloud +from dimos.msgs.geometry_msgs.Pose import Pose as LegacyPose from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.occupancy import occupancy_from_array +from dimos.msgs.pointcloud import pointcloud_xyz if TYPE_CHECKING: from collections.abc import Callable @@ -37,6 +44,49 @@ class OccupancyConfig: frame_id: str | None = None +def _cloud_points(cloud: PointCloud2 | CdrCloud) -> NDArray[np.float64]: + if isinstance(cloud, CdrCloud): + points = pointcloud_xyz(cloud) + return points[np.isfinite(points).all(axis=1)] + points, _ = cloud.as_numpy() + return points.astype(np.float64) + + +def _legacy_stamp(cloud: PointCloud2 | CdrCloud) -> float: + if isinstance(cloud, CdrCloud): + return 0.0 # Generated output uses its original Header, never this float. + return cloud.ts if cloud.ts is not None else 0.0 + + +def _grid_message( + cloud: PointCloud2 | CdrCloud, + config: OccupancyConfig, + cells: NDArray[np.int8], + origin: tuple[float, float, float] = (0.0, 0.0, 0.0), + *, + legacy_stamp: float | None = None, +) -> OccupancyGrid | CdrGrid: + if isinstance(cloud, CdrCloud): + header = deepcopy(cloud.header) + header.frame_id = config.frame_id or header.frame_id + return occupancy_from_array( + cells, + header=header, + resolution=config.resolution, + origin=Pose( + position=Point(x=origin[0], y=origin[1], z=origin[2]), + orientation=Quaternion(w=1, x=0, y=0, z=0), + ), + ) + return OccupancyGrid( + grid=cells, + resolution=config.resolution, + origin=LegacyPose(*origin), + frame_id=config.frame_id or cloud.frame_id, + ts=legacy_stamp, + ) + + ConfigT = TypeVar("ConfigT", bound=OccupancyConfig, covariant=True) @@ -54,7 +104,7 @@ def __call__(self, cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... # Populated after function definitions below -OCCUPANCY_ALGOS: dict[str, Callable[..., OccupancyGrid]] = {} +OCCUPANCY_ALGOS: dict[str, Callable[..., OccupancyGrid | CdrGrid]] = {} @dataclass(frozen=True) @@ -67,7 +117,15 @@ class HeightCostConfig(OccupancyConfig): smoothing: float = 1.0 -def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: +@overload +def height_cost_occupancy(cloud: CdrCloud, **kwargs: Any) -> CdrGrid: ... + + +@overload +def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... + + +def height_cost_occupancy(cloud: PointCloud2 | CdrCloud, **kwargs: Any) -> OccupancyGrid | CdrGrid: """Create a costmap based on terrain slope (rate of change of height). Costs are assigned based on the gradient magnitude of the terrain height. @@ -83,17 +141,11 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with costs 0-100 based on terrain slope, -1 for unknown """ cfg = HeightCostConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + points = _cloud_points(cloud) + ts = _legacy_stamp(cloud) if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Find bounds of the point cloud in X-Y plane (use all points) min_x = np.min(points[:, 0]) @@ -113,11 +165,7 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose - origin = Pose() - origin.position.x = min_x - origin.position.y = min_y - origin.position.z = 0.0 - origin.orientation.w = 1.0 + origin = (float(min_x), float(min_y), 0.0) # Step 1: Build min and max height maps for each cell # Initialize with NaN to track which cells have observations @@ -204,13 +252,7 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: else: cost = np.full((height, width), -1, dtype=np.int8) - return OccupancyGrid( - grid=cost, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, cost, origin, legacy_stamp=ts) @dataclass(frozen=True) @@ -223,7 +265,15 @@ class GeneralOccupancyConfig(OccupancyConfig): # can remove, just needs pulling out of unitree type/map.py -def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: +@overload +def general_occupancy(cloud: CdrCloud, **kwargs: Any) -> CdrGrid: ... + + +@overload +def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... + + +def general_occupancy(cloud: PointCloud2 | CdrCloud, **kwargs: Any) -> OccupancyGrid | CdrGrid: """Create an OccupancyGrid from a PointCloud2 message. Args: @@ -235,16 +285,10 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with occupied cells where points were projected """ cfg = GeneralOccupancyConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding + points = _cloud_points(cloud) if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Filter points by height for obstacles obstacle_mask = (points[:, 2] >= cfg.min_height) & (points[:, 2] <= cfg.max_height) @@ -272,11 +316,7 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose (bottom-left corner of the grid) - origin = Pose() - origin.position.x = min_x - origin.position.y = min_y - origin.position.z = 0.0 - origin.orientation.w = 1.0 # No rotation + origin = (float(min_x), float(min_y), 0.0) # Initialize grid (all unknown) grid = np.full((height, width), -1, dtype=np.int8) @@ -328,15 +368,9 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: # Create and return OccupancyGrid # Get timestamp from cloud if available - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + ts = _legacy_stamp(cloud) - return OccupancyGrid( - grid=grid, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, grid, origin, legacy_stamp=ts) @dataclass(frozen=True) @@ -353,7 +387,15 @@ class SimpleOccupancyConfig(OccupancyConfig): smoothing: float = 1.0 -def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: +@overload +def simple_occupancy(cloud: CdrCloud, **kwargs: Any) -> CdrGrid: ... + + +@overload +def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... + + +def simple_occupancy(cloud: PointCloud2 | CdrCloud, **kwargs: Any) -> OccupancyGrid | CdrGrid: """Create a simple occupancy grid with morphological closing. Args: @@ -365,16 +407,10 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with occupied/free cells """ cfg = SimpleOccupancyConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding + points = _cloud_points(cloud) if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Find bounds of the point cloud in X-Y plane min_x = float(np.min(points[:, 0])) - 1.0 @@ -387,11 +423,7 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose (bottom-left corner of the grid) - origin = Pose() - origin.position.x = min_x - origin.position.y = min_y - origin.position.z = 0.0 - origin.orientation.w = 1.0 + origin = (float(min_x), float(min_y), 0.0) # Initialize grid (all unknown) grid = np.full((height, width), -1, dtype=np.int8) @@ -411,15 +443,9 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: cfg.max_height, ) - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + ts = _legacy_stamp(cloud) - return OccupancyGrid( - grid=grid, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, grid, origin, legacy_stamp=ts) # Populate algorithm registry diff --git a/dimos/mapping/test_generated_costmapper.py b/dimos/mapping/test_generated_costmapper.py new file mode 100644 index 00000000000..98058d4af6c --- /dev/null +++ b/dimos/mapping/test_generated_costmapper.py @@ -0,0 +1,221 @@ +# Copyright 2026 Dimensional 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. + +from contextlib import ExitStack +from copy import deepcopy +from threading import Event + +import cv2 +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.core.cdr_transport import CdrTransport +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.occupancy.visualizations import rainbow_image +from dimos.mapping.pointclouds.occupancy import ( + SimpleOccupancyConfig, + general_occupancy, + height_cost_occupancy, + simple_occupancy, +) +from dimos.mapping.pointclouds.util import read_pointcloud +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2, PointField +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.msgs.occupancy import occupancy_from_array, occupancy_view +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as LegacyCloud +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.utils.data import get_data + + +@pytest.fixture(params=[(0, 0), (-1, 999999999), (1700000000, 123456789)]) +def source_header(request): + return Header(frame_id="lidar", stamp=Time(sec=request.param[0], nanosec=request.param[1])) + + +@pytest.mark.parametrize("algorithm", [simple_occupancy, general_occupancy, height_cost_occupancy]) +def test_generated_costmaps_match_existing_math_and_keep_exact_header(source_header, algorithm): + points = np.array( + [[-1, -1, 0], [-0.5, 0, 0.4], [0, 0, 0.4], [0.5, 0.5, 0], [1, 1, 0.2]], dtype=np.float32 + ) + cloud = pointcloud_from_xyz(points, header=source_header) + legacy = LegacyCloud.from_numpy(points, frame_id="lidar", timestamp=0) + + actual = algorithm(cloud, resolution=0.25) + original = algorithm(legacy, resolution=0.25) + received = decode(encode(actual), OccupancyGrid) + + assert type(original) is LegacyGrid + assert type(actual) is OccupancyGrid + assert received.header == source_header + assert received.info.map_load_time == source_header.stamp + assert original.ts == legacy.ts + assert received.info.resolution == original.resolution + assert (received.info.origin.position.x, received.info.origin.position.y) == ( + original.origin.position.x, + original.origin.position.y, + ) + np.testing.assert_array_equal(occupancy_view(received), original.grid) + actual.header.stamp.nanosec = 0 + assert cloud.header == source_header + + +@pytest.mark.parametrize("algorithm", [simple_occupancy, general_occupancy, height_cost_occupancy]) +def test_empty_generated_cloud_retains_source_metadata_and_unknown_map(source_header, algorithm): + cloud = pointcloud_from_xyz(np.empty((0, 3), dtype=np.float32), header=source_header) + + actual = algorithm(cloud, resolution=0.25, frame_id="map") + + assert actual.header.frame_id == "map" + assert actual.header.stamp == source_header.stamp + assert actual.info.map_load_time == source_header.stamp + assert occupancy_view(actual).tolist() == [[-1]] + + +def test_generated_missing_points_are_ignored_before_rasterization(): + header = Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)) + points = np.array([[0, 0, 0.4], [1, 1, 0], [np.nan, 0, 0], [0, np.inf, 0]], dtype=np.float32) + + actual = simple_occupancy(pointcloud_from_xyz(points, header=header), resolution=0.25) + expected = simple_occupancy(pointcloud_from_xyz(points[:2], header=header), resolution=0.25) + + np.testing.assert_array_equal(actual.data, expected.data) + assert actual.header == header + + +class MemoryRPC(LCMRPC): + """Use the real RPC implementation with an in-memory transport.""" + + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +@pytest.fixture +def mapper_factory(): + with ExitStack() as cleanup: + + def create(**kwargs): + module = CostMapper(rpc_transport=MemoryRPC, **kwargs) + cleanup.callback(module.stop) + return module + + yield create, cleanup + + +def test_initial_safe_radius_uses_owned_cells_without_changing_metadata(mapper_factory): + create, _ = mapper_factory + module = create(initial_safe_radius_meters=0.25) + header = Header(frame_id="map", stamp=Time(sec=1700000000, nanosec=123456789)) + grid = occupancy_from_array(np.full((3, 3), 100, dtype=np.int8), header=header, resolution=0.25) + original = deepcopy(grid.header) + + module._apply_initial_safe_radius(grid) + + assert occupancy_view(grid).tolist() == [[0, 0, 100], [0, 0, 100], [100, 100, 100]] + assert grid.header == original + assert grid.info.map_load_time == original.stamp + + +def test_costmapper_actual_generated_ports_deliver_cdr_without_network(mapper_factory): + create, cleanup = mapper_factory + global_wire = CdrTransport("costmap_global", PointCloud2, backend="lcm", url="memq://") + merged_wire = CdrTransport("costmap_merged", PointCloud2, backend="lcm", url="memq://") + output_wire = CdrTransport("costmap_output", OccupancyGrid, backend="lcm", url="memq://") + for wire in (global_wire, merged_wire, output_wire): + cleanup.callback(wire.stop) + module = create(algo="simple", config=SimpleOccupancyConfig(resolution=0.25)) + assert module.global_map.type is PointCloud2 + assert module.global_costmap.type is OccupancyGrid + module.global_map.transport = global_wire + module.merged_map.transport = merged_wire + module.global_costmap.transport = output_wire + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(output_wire.subscribe(receive)) + header = Header(frame_id="lidar", stamp=Time(sec=1700000000, nanosec=123456789)) + cloud = pointcloud_from_xyz(np.array([[0, 0, 0.4], [1, 1, 0]], dtype=np.float32), header=header) + module.start() + + global_wire.broadcast(None, cloud) + + assert ready.wait(5) + assert len(received) == 1 + assert received[0].header == header + assert 100 in received[0].data + + ready.clear() + merged_header = Header(frame_id="merged_map", stamp=Time(sec=1700000000, nanosec=123456790)) + merged = pointcloud_from_xyz( + np.array([[0, 0, 0.4], [1, 1, 0]], dtype=np.float32), header=merged_header + ) + merged_wire.broadcast(None, merged) + assert ready.wait(5) + assert len(received) == 2 + assert received[-1].header == merged_header + + +@pytest.mark.self_hosted +@pytest.mark.parametrize( + "source_name,algorithm,expected_name,grayscale", + [ + ("apartment", simple_occupancy, "occupancy_simple.png", True), + ("big_office.ply", height_cost_occupancy, "big_office_height_cost_occupancy.png", False), + ("big_office.ply", simple_occupancy, "big_office_simple_occupancy.png", False), + ], +) +def test_generated_real_pointcloud_costmap_matches_existing_golden( + source_name, algorithm, expected_name, grayscale +): + source = get_data(source_name) + if source_name == "apartment": + source = source / "sum.ply" + points = np.asarray(read_pointcloud(source).points).astype("