From cbd4f27384e44533ab0dccfb33cf6bfe56fcb1d5 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 19 Oct 2026 22:53:08 -0700 Subject: [PATCH 1/2] refactor: encode voxel map streams with generated CDR values --- .github/workflows/message-codegen.yml | 9 + dimos/mapping/voxels/grid.py | 23 +- dimos/mapping/voxels/impl/o3d.py | 21 +- dimos/mapping/voxels/impl/packed.py | 7 +- dimos/mapping/voxels/module.py | 16 +- dimos/mapping/voxels/test_generated_stream.py | 217 ++++++++++++++++++ 6 files changed, 283 insertions(+), 10 deletions(-) create mode 100644 dimos/mapping/voxels/test_generated_stream.py diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index 2fc79174369..b3a41a22629 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -20,6 +20,7 @@ on: - dimos/protocol/tf/** - dimos/memory/** - dimos/mapping/occupancy/** + - dimos/mapping/voxels/** - dimos/mapping/pointclouds/occupancy.py - dimos/mapping/odometry_hist.py - dimos/mapping/test_generated_odometry_hist.py @@ -246,3 +247,11 @@ jobs: run: | .native-review/bin/python -I -c 'import dimos.mapping.odometry_hist as history; import dimos.visualization.rerun.cdr as viewer; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (history, viewer)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/mapping/test_generated_odometry_hist.py" "$PWD/dimos/visualization/rerun/test_cdr.py" MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/odometry_hist.py dimos/visualization/rerun/cdr.py + - name: Verify generated voxel streams and CPU stores over memory CDR + if: runner.os == 'Linux' + env: + OMP_NUM_THREADS: '2' + NUMBA_NUM_THREADS: '2' + run: | + .native-review/bin/python -I -c 'from dimos.mapping.voxels import grid, module; from dimos.mapping.voxels.impl import packed, o3d; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (grid, module, packed, o3d)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/mapping/voxels/test_generated_stream.py" "$PWD/dimos/mapping/voxels/test_stamped.py" + MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/voxels/grid.py dimos/mapping/voxels/module.py dimos/mapping/voxels/impl/packed.py dimos/mapping/voxels/impl/o3d.py diff --git a/dimos/mapping/voxels/grid.py b/dimos/mapping/voxels/grid.py index a3dc189c0d5..a9ae31c1f47 100644 --- a/dimos/mapping/voxels/grid.py +++ b/dimos/mapping/voxels/grid.py @@ -18,6 +18,9 @@ import time from typing import TYPE_CHECKING +import numpy as np +from numpy.typing import NDArray + from dimos.mapping.voxels.impl.o3d import O3dVoxels from dimos.mapping.voxels.impl.packed import PackedVoxels from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @@ -87,6 +90,24 @@ def add_frame(self, frame: PointCloud2) -> None: self.get_global_pointcloud.invalidate_cache(self) self.get_global_pointcloud2.invalidate_cache(self) + def add_points(self, points: NDArray[np.float32], *, timestamp: float | None = None) -> None: + """Accumulate finite Nx3 coordinates without requiring a message wrapper. + + The optional per-frame timestamp uses the existing store's seconds unit; + stream adapters keep their exact source Header separately. + """ + self._check_disposed() + values = np.asarray(points, dtype=np.float32) + if values.ndim != 2 or values.shape[1] != 3: + raise ValueError("Voxel points must have shape Nx3") + if not np.isfinite(values).all(): + raise ValueError("Voxel points must be finite") + if timestamp is not None: + self._latest_frame_ts = timestamp + self._impl.add_points(values, timestamp) + self.get_global_pointcloud.invalidate_cache(self) + self.get_global_pointcloud2.invalidate_cache(self) + def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: """Move every voxel to where `place` puts its stamped original position.""" self._check_disposed() @@ -113,7 +134,7 @@ def get_global_pointcloud(self) -> o3d.t.geometry.PointCloud: self._check_disposed() cpu = o3c.Device("CPU:0") out = o3d.t.geometry.PointCloud(device=cpu) - out.point["positions"] = o3c.Tensor(self._impl.points(), device=cpu) + out.point["positions"] = o3c.Tensor(self._impl.points(), dtype=o3c.float32, device=cpu) return out def size(self) -> int: diff --git a/dimos/mapping/voxels/impl/o3d.py b/dimos/mapping/voxels/impl/o3d.py index 82e80982b97..7a88ac16601 100644 --- a/dimos/mapping/voxels/impl/o3d.py +++ b/dimos/mapping/voxels/impl/o3d.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING import numpy as np +from numpy.typing import NDArray from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @@ -74,7 +75,18 @@ def add_frame(self, frame: PointCloud2) -> None: if pcd.is_empty(): return - pts = pcd.point["positions"].to(self._dev, o3c.float32) + self._insert_points(pcd.point["positions"].to(self._dev, o3c.float32)) + + def add_points(self, points: NDArray[np.float32], _timestamp: float | None = None) -> None: + import open3d.core as o3c # type: ignore[import-untyped] + + if not len(points): + return + self._insert_points(o3c.Tensor(points, dtype=o3c.float32, device=self._dev)) + + def _insert_points(self, pts: o3c.Tensor) -> None: + import open3d.core as o3c # type: ignore[import-untyped] + vox = (pts / self._voxel_size).floor().to(self._key_dtype) keys_Nx3 = vox.contiguous() @@ -108,7 +120,7 @@ def _carve_and_insert(self, new_keys: o3c.Tensor) -> None: value_element_shapes=[o3c.SizeVector([1])], device=self._dev, ) - dummy_vals = o3c.Tensor.zeros((xy_keys.shape[0], 1), o3c.uint8, self._dev) + dummy_vals = o3c.Tensor.zeros(o3c.SizeVector([xy_keys.shape[0], 1]), o3c.uint8, self._dev) xy_hashmap.insert(xy_keys, dummy_vals) active_indices = self._hashmap.active_buf_indices() @@ -142,5 +154,6 @@ def size(self) -> int: def dispose(self) -> None: """Free GPU resources.""" - self.vbg = None - self._hashmap = None + if self.vbg is not None: + self.vbg = None + del self._hashmap diff --git a/dimos/mapping/voxels/impl/packed.py b/dimos/mapping/voxels/impl/packed.py index 83f9d5f2e25..e774f592794 100644 --- a/dimos/mapping/voxels/impl/packed.py +++ b/dimos/mapping/voxels/impl/packed.py @@ -17,6 +17,7 @@ from collections.abc import Callable import numpy as np +from numpy.typing import NDArray from dimos.mapping.voxels.keys import FIELD_BITS as _BITS, FIELD_MASK as _MASK, KEY_OFFSET as _BIAS from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @@ -56,7 +57,9 @@ def _centers(self, k: np.ndarray) -> np.ndarray: return (vox + np.float32(0.5)) * np.float32(self._voxel_size) def add_frame(self, frame: PointCloud2) -> None: - pts = frame.points_f32() + self.add_points(frame.points_f32(), frame.ts) + + def add_points(self, pts: NDArray[np.float32], timestamp: float | None = None) -> None: if not len(pts): return new = np.unique(self._pack(pts)) @@ -84,7 +87,7 @@ def add_frame(self, frame: PointCloud2) -> None: pos, new = pos[fresh], new[fresh] self._keys = np.insert(keys, pos, new) if self.stamped: - self._ts = np.insert(self._ts, pos, frame.ts) + self._ts = np.insert(self._ts, pos, np.nan if timestamp is None else timestamp) self._raw = np.insert(self._raw, pos, new) def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: diff --git a/dimos/mapping/voxels/module.py b/dimos/mapping/voxels/module.py index b273889a4c0..f5f4004973d 100644 --- a/dimos/mapping/voxels/module.py +++ b/dimos/mapping/voxels/module.py @@ -14,8 +14,11 @@ from __future__ import annotations +from copy import deepcopy from typing import TYPE_CHECKING, Any +import numpy as np + from dimos.core.core import rpc from dimos.core.module import ModuleConfig from dimos.core.stream import In, Out @@ -23,7 +26,9 @@ from dimos.memory.module import StreamModule from dimos.memory.stream import Stream from dimos.memory.transform import Transformer -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.time import to_seconds if TYPE_CHECKING: from collections.abc import Iterator @@ -53,8 +58,11 @@ def _make_obs( self, grid: VoxelGrid, last_obs: Observation[PointCloud2], count: int ) -> Observation[PointCloud2]: # pose=None: the global map is in world frame, per-observation pose is meaningless + header = deepcopy(last_obs.data.header) + header.frame_id = self._grid_kwargs.get("frame_id", "world") + points = grid.get_global_pointcloud().point["positions"].numpy() return last_obs.derive( - data=grid.get_global_pointcloud2(), + data=pointcloud_from_xyz(points, header=header), pose=None, tags={**last_obs.tags, "frame_count": count}, ) @@ -68,7 +76,9 @@ def __call__( count = 0 for obs in upstream: - grid.add_frame(obs.data) + coordinates = pointcloud_xyz(obs.data) + points = coordinates[np.isfinite(coordinates).all(axis=1)].astype(np.float32) + grid.add_points(points, timestamp=to_seconds(obs.data.header.stamp)) last_obs = obs count += 1 diff --git a/dimos/mapping/voxels/test_generated_stream.py b/dimos/mapping/voxels/test_generated_stream.py new file mode 100644 index 00000000000..6019f0c3490 --- /dev/null +++ b/dimos/mapping/voxels/test_generated_stream.py @@ -0,0 +1,217 @@ +# 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 + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import open3d.core as o3c +import pytest + +from dimos.core.cdr_transport import CdrTransport +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.pointclouds.occupancy import SimpleOccupancyConfig +from dimos.mapping.voxels.grid import VoxelGrid +from dimos.mapping.voxels.impl.o3d import O3dVoxels +from dimos.mapping.voxels.impl.packed import PackedVoxels +from dimos.mapping.voxels.module import VoxelGridMapper, VoxelMapTransformer +from dimos.memory.type.observation import Observation +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as LegacyCloud +from dimos.msgs.time import time_from_nanoseconds, to_nanoseconds +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def keys(points): + return {tuple(row) for row in np.floor(np.asarray(points) / np.float32(0.05)).astype(np.int64)} + + +@pytest.mark.parametrize("carve", [False, True]) +def test_array_input_matches_legacy_voxel_ingestion(carve): + with ExitStack() as cleanup: + old = VoxelGrid(device="CPU:0", carve_columns=carve, show_startup_log=False) + new = VoxelGrid(device="CPU:0", carve_columns=carve, show_startup_log=False) + cleanup.callback(old.dispose) + cleanup.callback(new.dispose) + for index, points in enumerate( + ([[0.025, 0.025, 0.025], [0.075, 0.025, 0.025]], [[0.025, 0.025, 0.125]]) + ): + values = np.array(points, dtype=np.float32) + old.add_frame(LegacyCloud.from_numpy(values, timestamp=index)) + new.add_points(values, timestamp=index) + assert keys(old.get_global_pointcloud().point["positions"].numpy()) == keys( + new.get_global_pointcloud().point["positions"].numpy() + ) + assert len(old) == len(new) == (2 if carve else 3) + + +@pytest.mark.parametrize("carve", [False, True]) +def test_actual_open3d_cpu_array_entry_matches_packed_store(carve): + with ExitStack() as cleanup: + packed = PackedVoxels(0.05, carve) + tensor = O3dVoxels(0.05, 128, carve, o3c.Device("CPU:0")) + cleanup.callback(packed.dispose) + cleanup.callback(tensor.dispose) + for points in ([[0.025, 0.025, 0.025], [0.075, 0.025, 0.025]], [[0.025, 0.025, 0.125]]): + values = np.array(points, dtype=np.float32) + packed.add_points(values, 0) + tensor.add_points(values, 0) + assert keys(packed.points()) == keys(tensor.points()) + assert packed.size() == tensor.size() == (2 if carve else 3) + + +@pytest.mark.parametrize("points", [np.zeros((3,)), np.zeros((1, 2)), np.array([[np.nan, 0, 0]])]) +def test_array_entry_rejects_invalid_coordinates(points): + grid = VoxelGrid(device="CPU:0", show_startup_log=False) + try: + with pytest.raises(ValueError, match="Voxel points"): + grid.add_points(points) + finally: + grid.dispose() + + +@pytest.mark.parametrize("base", [0, -1, 1700000000123456789]) +@pytest.mark.parametrize("emit_every,indices", [(0, [2]), (1, [0, 1, 2]), (2, [1, 2])]) +def test_generated_stream_retains_exact_headers_batch_tail_and_observation_metadata( + base, emit_every, indices +): + frames = [] + for index in range(3): + cloud = pointcloud_from_xyz( + np.array([[0.025 + 0.05 * index, 0.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(base + index)), + ) + frames.append( + Observation(ts=42 + index, data_type=PointCloud2, _data=cloud, tags={"source": "lidar"}) + ) + original = deepcopy(frames[-1].data.header) + transform = VoxelMapTransformer( + emit_every=emit_every, + device="CPU:0", + voxel_size=0.05, + frame_id="map", + show_startup_log=False, + ) + + result = list(transform(iter(frames))) + + assert [to_nanoseconds(item.data.header.stamp) for item in result] == [ + base + index for index in indices + ] + assert [item.tags["frame_count"] for item in result] == [index + 1 for index in indices] + assert [item.ts for item in result] == [42 + index for index in indices] + assert all(item.tags["source"] == "lidar" and item.pose is None for item in result) + assert all(item.data.header.frame_id == "map" for item in result) + received = decode(encode(result[-1].data), PointCloud2) + assert len(pointcloud_xyz(received)) == 3 + assert frames[-1].data.header == original + + +def test_generated_empty_and_missing_points_keep_zero_source_stamp(): + header = Header(frame_id="world", stamp=time_from_nanoseconds(0)) + cloud = pointcloud_from_xyz(np.array([[np.nan, 0, 0]], dtype=np.float32), header=header) + transform = VoxelMapTransformer(device="CPU:0", show_startup_log=False) + + result = list(transform(iter([Observation(ts=123, _data=cloud, data_type=PointCloud2)]))) + + assert pointcloud_xyz(result[0].data).shape == (0, 3) + assert to_nanoseconds(result[0].data.header.stamp) == 0 + assert result[0].ts == 123 + + +def test_voxel_mapper_real_generated_ports_accumulate_over_cdr_memory_transport(): + with ExitStack() as cleanup: + incoming = CdrTransport("voxel_input", PointCloud2, backend="lcm", url="memq://") + outgoing = CdrTransport("voxel_map", PointCloud2, backend="lcm", url="memq://") + cleanup.callback(incoming.stop) + cleanup.callback(outgoing.stop) + module = VoxelGridMapper(rpc_transport=MemoryRPC, device="CPU:0", voxel_size=0.05) + cleanup.callback(module.stop) + assert module.lidar.type is PointCloud2 + assert module.global_map.type is PointCloud2 + module.lidar.transport = incoming + module.global_map.transport = outgoing + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(outgoing.subscribe(receive)) + module.start() + source = pointcloud_from_xyz( + np.array([[0.025, 0.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)), + ) + + incoming.broadcast(None, source) + + assert ready.wait(5) + assert len(received) == 1 + assert received[0].header == source.header + assert keys(pointcloud_xyz(received[0])) == {(0, 0, 0)} + + +def test_actual_voxel_to_costmapper_chain_preserves_source_header_over_cdr(): + with ExitStack() as cleanup: + source_wire = CdrTransport("voxel_chain_source", PointCloud2, backend="lcm", url="memq://") + map_wire = CdrTransport("voxel_chain_map", PointCloud2, backend="lcm", url="memq://") + merged_wire = CdrTransport("voxel_chain_merged", PointCloud2, backend="lcm", url="memq://") + cost_wire = CdrTransport("voxel_chain_cost", OccupancyGrid, backend="lcm", url="memq://") + for wire in (source_wire, map_wire, merged_wire, cost_wire): + cleanup.callback(wire.stop) + voxel = VoxelGridMapper(rpc_transport=MemoryRPC, device="CPU:0", voxel_size=0.05) + cost = CostMapper( + rpc_transport=MemoryRPC, algo="simple", config=SimpleOccupancyConfig(resolution=0.25) + ) + cleanup.callback(voxel.stop) + cleanup.callback(cost.stop) + voxel.lidar.transport = source_wire + voxel.global_map.transport = map_wire + cost.global_map.transport = map_wire + cost.merged_map.transport = merged_wire + cost.global_costmap.transport = cost_wire + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(cost_wire.subscribe(receive)) + cost.start() + voxel.start() + cloud = pointcloud_from_xyz( + np.array([[0.025, 0.025, 0.225], [1.025, 1.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)), + ) + + source_wire.broadcast(None, cloud) + + assert ready.wait(5) + assert len(received) == 1 + assert received[0].header == cloud.header + assert received[0].info.map_load_time == cloud.header.stamp + assert 100 in received[0].data From a027c97e524a611ecb684f4204945a3987ae6e99 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:53:58 +0000 Subject: [PATCH 2/2] [autofix.ci] apply automated fixes --- dimos/models/segmentation/edge_tam.py | 2 +- dimos/models/segmentation/test_yoloe.py | 4 ++-- dimos/models/segmentation/yoloe.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dimos/models/segmentation/edge_tam.py b/dimos/models/segmentation/edge_tam.py index e240a6529d1..b65977283c0 100644 --- a/dimos/models/segmentation/edge_tam.py +++ b/dimos/models/segmentation/edge_tam.py @@ -20,7 +20,6 @@ import tempfile from typing import TYPE_CHECKING, Any, Protocol, TypedDict -from dimos.messages.sensor_msgs.msg import Image from hydra.utils import instantiate import numpy as np from numpy.typing import NDArray @@ -28,6 +27,7 @@ from PIL import Image as PILImage import torch +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import default_torch_device from dimos.msgs.image import image_to_bgr, image_to_jpeg from dimos.perception.detection.detectors.base import Detector diff --git a/dimos/models/segmentation/test_yoloe.py b/dimos/models/segmentation/test_yoloe.py index 345fc79ac1c..5d2efaffa50 100644 --- a/dimos/models/segmentation/test_yoloe.py +++ b/dimos/models/segmentation/test_yoloe.py @@ -16,11 +16,11 @@ from typing import Any from unittest.mock import MagicMock -from dimos.messages.sensor_msgs.msg import Image -from dimos.messages.std_msgs.msg import Header import numpy as np from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor +from dimos.messages.sensor_msgs.msg import Image +from dimos.messages.std_msgs.msg import Header from dimos.models.segmentation.yoloe import YoloeBoxSegmenter from dimos.msgs.image import image_from_array, image_to_bgr from dimos.msgs.time import time_from_seconds, to_seconds diff --git a/dimos/models/segmentation/yoloe.py b/dimos/models/segmentation/yoloe.py index d19ecc15701..fd075227d27 100644 --- a/dimos/models/segmentation/yoloe.py +++ b/dimos/models/segmentation/yoloe.py @@ -19,11 +19,11 @@ from dataclasses import replace from typing import Protocol -from dimos.messages.sensor_msgs.msg import Image import numpy as np from numpy.typing import NDArray from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor +from dimos.messages.sensor_msgs.msg import Image from dimos.msgs.image import image_to_bgr from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox