mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-01 05:59:44 +00:00
navigation: port voxel map + occupancy/A* mapping backbone
Copied from the dyna360 research stack into lerobot.navigation: - voxel_map.py: sparse-hash VoxelMap (5 cm default), count-weighted running-mean geometry + optional per-voxel semantic feature; carve() for DynaMem-style free-space removal (dynamic updates), query() for top-k cosine text matches, remove_voxels_in_box() for scene mutation. Added an inverse.reshape(-1) guard for numpy 2.x's np.unique. - occupancy.py: derived 3-class top-down grid (UNOBSERVED/NAVIGABLE/ OBSTACLE) projected from the voxel map, A* with no corner-cutting, obstacle inflation, and frontier extraction for exploration. Added OccupancyGrid.is_obstacle() used by SafeBaseController's occupancy gate. Pure numpy, no torch/hardware/SDK. 42 new tests (71 total in tests/navigation/). Next: LingBot-Map geometry runner + integration. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Pepijn
@@ -31,12 +31,33 @@ from .base_controller import (
|
||||
odometry_to_world_pose,
|
||||
world_velocity_to_body,
|
||||
)
|
||||
from .occupancy import (
|
||||
NAVIGABLE,
|
||||
OBSTACLE,
|
||||
UNOBSERVED,
|
||||
OccupancyGrid,
|
||||
astar,
|
||||
find_frontier_cells,
|
||||
project_voxel_map_to_grid,
|
||||
)
|
||||
from .voxel_map import CarveResult, QueryResult, VoxelMap, VoxelSnapshot
|
||||
|
||||
__all__ = [
|
||||
"NAVIGABLE",
|
||||
"OBSTACLE",
|
||||
"UNOBSERVED",
|
||||
"BaseController",
|
||||
"CarveResult",
|
||||
"OccupancyGrid",
|
||||
"QueryResult",
|
||||
"RobotBaseController",
|
||||
"SafeBaseController",
|
||||
"StubBaseController",
|
||||
"VoxelMap",
|
||||
"VoxelSnapshot",
|
||||
"astar",
|
||||
"find_frontier_cells",
|
||||
"odometry_to_world_pose",
|
||||
"project_voxel_map_to_grid",
|
||||
"world_velocity_to_body",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,371 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# 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 occupancy projection of the voxel map + A* path planning.
|
||||
|
||||
Ported from the dyna360 research stack. Derived, not maintained: every
|
||||
call to :func:`project_voxel_map_to_grid` rebuilds the 3-class grid from
|
||||
a fresh ``VoxelMap.snapshot()``, so the projection reflects whatever the
|
||||
keyframe loop most recently carved or added — no separate obstacle
|
||||
structure to keep in sync.
|
||||
|
||||
Coordinate convention: OpenCV (X right, Y *down*, Z forward), matching
|
||||
the navigation world frame. "Up" is the −Y direction. The top-down grid
|
||||
indexes the XZ plane; cell ``(iz, ix)`` covers world rectangle
|
||||
``[origin_x + ix·cell, origin_x + (ix+1)·cell]`` ×
|
||||
``[origin_z + iz·cell, origin_z + (iz+1)·cell]``.
|
||||
|
||||
Classes:
|
||||
- ``UNOBSERVED`` (0): no voxel projects here. The base must not plan
|
||||
through it (might be an unseen obstacle), but explorers treat it as
|
||||
the goal class.
|
||||
- ``NAVIGABLE`` (1): observed ground / open space.
|
||||
- ``OBSTACLE`` (2): at least one voxel in the robot-height band
|
||||
projects here.
|
||||
"""
|
||||
|
||||
# ruff: noqa: N806 — H, W, D are conventional array-dimension names (and appear verbatim in error strings)
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
import logging
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.navigation.voxel_map import VoxelMap
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
# Class constants — picked so a colormap can index directly.
|
||||
UNOBSERVED = np.int8(0)
|
||||
NAVIGABLE = np.int8(1)
|
||||
OBSTACLE = np.int8(2)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OccupancyGrid:
|
||||
"""3-class top-down grid plus its world↔cell mapping."""
|
||||
|
||||
classes: np.ndarray # (H, W) int8 — H = z-extent, W = x-extent
|
||||
cell_size: float # m per cell
|
||||
origin_x: float # world x of the LEFT edge of column 0
|
||||
origin_z: float # world z of the TOP edge of row 0
|
||||
ground_y: float # world y of the (auto-estimated or given) ground plane
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int, int]:
|
||||
return self.classes.shape # (H, W)
|
||||
|
||||
def world_to_cell(self, x: float, z: float) -> tuple[int, int]:
|
||||
"""Return ``(iz, ix)``, clipped to grid extents."""
|
||||
ix = int(np.clip(math.floor((x - self.origin_x) / self.cell_size), 0, self.shape[1] - 1))
|
||||
iz = int(np.clip(math.floor((z - self.origin_z) / self.cell_size), 0, self.shape[0] - 1))
|
||||
return iz, ix
|
||||
|
||||
def cell_to_world(self, iz: int, ix: int) -> tuple[float, float]:
|
||||
"""Return ``(x, z)`` at the *centre* of cell ``(iz, ix)``."""
|
||||
x = self.origin_x + (ix + 0.5) * self.cell_size
|
||||
z = self.origin_z + (iz + 0.5) * self.cell_size
|
||||
return x, z
|
||||
|
||||
def is_navigable(self, iz: int, ix: int) -> bool:
|
||||
H, W = self.shape
|
||||
return 0 <= iz < H and 0 <= ix < W and self.classes[iz, ix] == NAVIGABLE
|
||||
|
||||
def is_obstacle(self, iz: int, ix: int) -> bool:
|
||||
"""Whether cell ``(iz, ix)`` is a known obstacle. Out-of-bounds is
|
||||
not an obstacle (it is simply unobservable) — used by
|
||||
``SafeBaseController``'s occupancy gate."""
|
||||
H, W = self.shape
|
||||
return 0 <= iz < H and 0 <= ix < W and self.classes[iz, ix] == OBSTACLE
|
||||
|
||||
def is_in_bounds(self, iz: int, ix: int) -> bool:
|
||||
H, W = self.shape
|
||||
return 0 <= iz < H and 0 <= ix < W
|
||||
|
||||
def nearest_navigable_cell(self, iz: int, ix: int, max_radius: int = 50) -> tuple[int, int] | None:
|
||||
"""BFS outward until a navigable cell is found, or give up."""
|
||||
if self.is_navigable(iz, ix):
|
||||
return iz, ix
|
||||
for r in range(1, max_radius + 1):
|
||||
for diz in range(-r, r + 1):
|
||||
for dix in range(-r, r + 1):
|
||||
if max(abs(diz), abs(dix)) != r:
|
||||
continue # ring only, not the interior
|
||||
if self.is_navigable(iz + diz, ix + dix):
|
||||
return iz + diz, ix + dix
|
||||
return None
|
||||
|
||||
|
||||
def estimate_ground_y(xyz: np.ndarray, percentile: float = 95.0) -> float:
|
||||
"""Estimate the world-frame y of the ground plane.
|
||||
|
||||
Y is down (OpenCV), so the ground is at the LARGEST y values. Using a
|
||||
high percentile (default 95) is robust to outliers below the ground.
|
||||
"""
|
||||
if xyz.size == 0:
|
||||
return 0.0
|
||||
return float(np.percentile(xyz[:, 1], percentile))
|
||||
|
||||
|
||||
def project_voxel_map_to_grid(
|
||||
voxel_map: VoxelMap,
|
||||
*,
|
||||
cell_size: float = 0.1,
|
||||
ground_y: float | None = None,
|
||||
obstacle_y_range: tuple[float, float] = (-2.0, -0.1),
|
||||
bbox: tuple[float, float, float, float] | None = None,
|
||||
bbox_pad: float = 1.0,
|
||||
inflate_cells: int = 0,
|
||||
) -> OccupancyGrid:
|
||||
"""Snapshot the voxel map and project it into a 2D occupancy grid.
|
||||
|
||||
``obstacle_y_range`` is interpreted *relative* to ``ground_y`` with the
|
||||
Y-down convention, so the default ``(-2.0, -0.1)`` means "voxels
|
||||
between 2.0 m and 0.1 m above the ground are obstacles". Anything above
|
||||
the ceiling band or below ground level is silently ignored.
|
||||
|
||||
``inflate_cells`` dilates the OBSTACLE class by N cells of clearance
|
||||
(square morphology) — a body-radius safety margin for the base without
|
||||
resampling the voxel map.
|
||||
"""
|
||||
snap = voxel_map.snapshot()
|
||||
xyz = snap.xyz
|
||||
|
||||
if xyz.size == 0:
|
||||
# Empty map → a 1×1 grid of UNOBSERVED at world origin.
|
||||
return OccupancyGrid(
|
||||
classes=np.zeros((1, 1), dtype=np.int8),
|
||||
cell_size=float(cell_size),
|
||||
origin_x=0.0,
|
||||
origin_z=0.0,
|
||||
ground_y=ground_y if ground_y is not None else 0.0,
|
||||
)
|
||||
|
||||
if ground_y is None:
|
||||
ground_y = estimate_ground_y(xyz)
|
||||
|
||||
# Promote to float64 — VoxelMap snapshots are float32, and naive
|
||||
# ``(float32_array <= float64_scalar)`` lets numpy downcast the scalar
|
||||
# back to float32, which causes edge-case bugs (e.g. 1.0 <= 0.999999999
|
||||
# becomes True because the threshold rounds up to 1.0 in float32).
|
||||
x_arr = xyz[:, 0].astype(np.float64)
|
||||
y_arr = xyz[:, 1].astype(np.float64)
|
||||
z_arr = xyz[:, 2].astype(np.float64)
|
||||
|
||||
abs_y_top = ground_y + obstacle_y_range[0] # most-negative y (highest above ground)
|
||||
abs_y_bottom = ground_y + obstacle_y_range[1] # closer to ground
|
||||
is_obstacle = (y_arr >= abs_y_top) & (y_arr <= abs_y_bottom)
|
||||
|
||||
if bbox is None:
|
||||
x_min = float(x_arr.min()) - bbox_pad
|
||||
x_max = float(x_arr.max()) + bbox_pad
|
||||
z_min = float(z_arr.min()) - bbox_pad
|
||||
z_max = float(z_arr.max()) + bbox_pad
|
||||
else:
|
||||
x_min, z_min, x_max, z_max = bbox
|
||||
|
||||
W = max(1, int(math.ceil((x_max - x_min) / cell_size)))
|
||||
H = max(1, int(math.ceil((z_max - z_min) / cell_size)))
|
||||
classes = np.zeros((H, W), dtype=np.int8) # default UNOBSERVED
|
||||
|
||||
# eps absorbs float32→float64 representation drift so points that should
|
||||
# land exactly on a cell boundary aren't randomly bumped into the
|
||||
# previous cell. 1e-3 of a cell width is well above float32's ~1e-7
|
||||
# relative precision and well below the 0.5-cell misclassification
|
||||
# threshold.
|
||||
eps = cell_size * 1e-3
|
||||
ix = np.clip(np.floor((x_arr - x_min) / cell_size + eps).astype(np.int32), 0, W - 1)
|
||||
iz = np.clip(np.floor((z_arr - z_min) / cell_size + eps).astype(np.int32), 0, H - 1)
|
||||
|
||||
# Two-pass labelling: any voxel makes a cell observed (-> NAVIGABLE);
|
||||
# obstacle voxels then upgrade those cells to OBSTACLE.
|
||||
classes[iz, ix] = NAVIGABLE
|
||||
obs_iz = iz[is_obstacle]
|
||||
obs_ix = ix[is_obstacle]
|
||||
classes[obs_iz, obs_ix] = OBSTACLE
|
||||
|
||||
if inflate_cells > 0:
|
||||
classes = _inflate_obstacles(classes, inflate_cells)
|
||||
|
||||
return OccupancyGrid(
|
||||
classes=classes,
|
||||
cell_size=float(cell_size),
|
||||
origin_x=x_min,
|
||||
origin_z=z_min,
|
||||
ground_y=float(ground_y),
|
||||
)
|
||||
|
||||
|
||||
def _inflate_obstacles(classes: np.ndarray, radius: int) -> np.ndarray:
|
||||
"""Dilate OBSTACLE cells by `radius` cells (Chebyshev). Pure-numpy
|
||||
morphological dilation — fine for our grid sizes."""
|
||||
out = classes.copy()
|
||||
obs = classes == OBSTACLE
|
||||
H, W = classes.shape
|
||||
for diz in range(-radius, radius + 1):
|
||||
for dix in range(-radius, radius + 1):
|
||||
if diz == 0 and dix == 0:
|
||||
continue
|
||||
sl_src_iz = slice(max(0, -diz), H - max(0, diz))
|
||||
sl_src_ix = slice(max(0, -dix), W - max(0, dix))
|
||||
sl_dst_iz = slice(max(0, diz), H - max(0, -diz))
|
||||
sl_dst_ix = slice(max(0, dix), W - max(0, -dix))
|
||||
inflated = obs[sl_src_iz, sl_src_ix]
|
||||
# Only upgrade NAVIGABLE → OBSTACLE; never overwrite UNOBSERVED
|
||||
# so the frontier (NAVIGABLE↔UNOBSERVED boundary) survives.
|
||||
target = out[sl_dst_iz, sl_dst_ix]
|
||||
promote = inflated & (target == NAVIGABLE)
|
||||
target[promote] = OBSTACLE
|
||||
out[sl_dst_iz, sl_dst_ix] = target
|
||||
return out
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- A*
|
||||
|
||||
_DIAG_COST = math.sqrt(2.0)
|
||||
_NEIGHBOURS_ORTHO = ((-1, 0), (1, 0), (0, -1), (0, 1))
|
||||
_NEIGHBOURS_DIAG = ((-1, -1), (-1, 1), (1, -1), (1, 1))
|
||||
|
||||
|
||||
def astar(
|
||||
grid: OccupancyGrid,
|
||||
start_world: tuple[float, float],
|
||||
goal_world: tuple[float, float],
|
||||
*,
|
||||
allow_unobserved_goal: bool = True,
|
||||
) -> list[tuple[float, float]] | None:
|
||||
"""Plan a path from ``start_world`` to ``goal_world`` in (x, z) world m.
|
||||
|
||||
Returns a list of world ``(x, z)`` waypoints, or ``None`` if no path
|
||||
exists. Start/goal are snapped to the nearest navigable cell.
|
||||
"""
|
||||
H, W = grid.shape
|
||||
if H == 0 or W == 0:
|
||||
return None
|
||||
|
||||
s_iz, s_ix = grid.world_to_cell(*start_world)
|
||||
g_iz, g_ix = grid.world_to_cell(*goal_world)
|
||||
|
||||
if not grid.is_navigable(s_iz, s_ix):
|
||||
snapped = grid.nearest_navigable_cell(s_iz, s_ix)
|
||||
if snapped is None:
|
||||
return None
|
||||
s_iz, s_ix = snapped
|
||||
if not grid.is_navigable(g_iz, g_ix):
|
||||
if not allow_unobserved_goal:
|
||||
return None
|
||||
snapped = grid.nearest_navigable_cell(g_iz, g_ix)
|
||||
if snapped is None:
|
||||
return None
|
||||
g_iz, g_ix = snapped
|
||||
|
||||
def heuristic(iz: int, ix: int) -> float:
|
||||
d_iz = abs(iz - g_iz)
|
||||
d_ix = abs(ix - g_ix)
|
||||
return (max(d_iz, d_ix) - min(d_iz, d_ix)) + _DIAG_COST * min(d_iz, d_ix)
|
||||
|
||||
open_heap: list[tuple[float, int, tuple[int, int]]] = []
|
||||
counter = 0 # tiebreaker so heapq doesn't compare tuples on ties
|
||||
heapq.heappush(open_heap, (0.0, counter, (s_iz, s_ix)))
|
||||
came_from: dict[tuple[int, int], tuple[int, int]] = {}
|
||||
g_score: dict[tuple[int, int], float] = {(s_iz, s_ix): 0.0}
|
||||
|
||||
while open_heap:
|
||||
_, _, current = heapq.heappop(open_heap)
|
||||
if current == (g_iz, g_ix):
|
||||
return _reconstruct_path(came_from, current, grid)
|
||||
|
||||
cur_iz, cur_ix = current
|
||||
cur_g = g_score[current]
|
||||
|
||||
for diz, dix in _NEIGHBOURS_ORTHO:
|
||||
n = (cur_iz + diz, cur_ix + dix)
|
||||
if not grid.is_navigable(*n):
|
||||
continue
|
||||
tentative = cur_g + 1.0
|
||||
if tentative < g_score.get(n, float("inf")):
|
||||
came_from[n] = current
|
||||
g_score[n] = tentative
|
||||
counter += 1
|
||||
heapq.heappush(open_heap, (tentative + heuristic(*n), counter, n))
|
||||
|
||||
for diz, dix in _NEIGHBOURS_DIAG:
|
||||
n = (cur_iz + diz, cur_ix + dix)
|
||||
if not grid.is_navigable(*n):
|
||||
continue
|
||||
# Prevent corner-cutting: both perpendicular neighbours must be
|
||||
# navigable, or we'd squeeze through an obstacle's diagonal.
|
||||
if not grid.is_navigable(cur_iz + diz, cur_ix):
|
||||
continue
|
||||
if not grid.is_navigable(cur_iz, cur_ix + dix):
|
||||
continue
|
||||
tentative = cur_g + _DIAG_COST
|
||||
if tentative < g_score.get(n, float("inf")):
|
||||
came_from[n] = current
|
||||
g_score[n] = tentative
|
||||
counter += 1
|
||||
heapq.heappush(open_heap, (tentative + heuristic(*n), counter, n))
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _reconstruct_path(
|
||||
came_from: dict[tuple[int, int], tuple[int, int]],
|
||||
end: tuple[int, int],
|
||||
grid: OccupancyGrid,
|
||||
) -> list[tuple[float, float]]:
|
||||
cells = [end]
|
||||
while cells[-1] in came_from:
|
||||
cells.append(came_from[cells[-1]])
|
||||
cells.reverse()
|
||||
return [grid.cell_to_world(iz, ix) for iz, ix in cells]
|
||||
|
||||
|
||||
# ----------------------------------------------------------------- frontier
|
||||
|
||||
|
||||
def find_frontier_cells(grid: OccupancyGrid) -> np.ndarray:
|
||||
"""Return cells on the NAVIGABLE↔UNOBSERVED boundary, as ``(K, 2)`` int.
|
||||
|
||||
These are the cells exploration aims for: places we already know we
|
||||
can stand at, but with unknown adjacent territory worth visiting.
|
||||
"""
|
||||
nav = grid.classes == NAVIGABLE
|
||||
unobs = grid.classes == UNOBSERVED
|
||||
if not nav.any() or not unobs.any():
|
||||
return np.zeros((0, 2), dtype=np.int32)
|
||||
|
||||
boundary = np.zeros_like(nav)
|
||||
boundary[1:, :] |= nav[1:, :] & unobs[:-1, :]
|
||||
boundary[:-1, :] |= nav[:-1, :] & unobs[1:, :]
|
||||
boundary[:, 1:] |= nav[:, 1:] & unobs[:, :-1]
|
||||
boundary[:, :-1] |= nav[:, :-1] & unobs[:, 1:]
|
||||
iz, ix = np.where(boundary)
|
||||
return np.stack([iz, ix], axis=-1).astype(np.int32)
|
||||
|
||||
|
||||
def occupancy_to_rgb(grid: OccupancyGrid) -> np.ndarray:
|
||||
"""Render the 3-class grid as an (H, W, 3) uint8 image."""
|
||||
img = np.zeros((*grid.shape, 3), dtype=np.uint8)
|
||||
img[grid.classes == UNOBSERVED] = (40, 40, 50)
|
||||
img[grid.classes == NAVIGABLE] = (200, 200, 200)
|
||||
img[grid.classes == OBSTACLE] = (220, 60, 60)
|
||||
return img
|
||||
@@ -0,0 +1,508 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# 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.
|
||||
|
||||
"""Sparse-hash voxel memory with free-space carving + semantic features.
|
||||
|
||||
Ported from the dyna360 research stack. Per occupied voxel: voxel index,
|
||||
running-mean xyz (count-weighted), running-mean rgb (count-weighted),
|
||||
count, last_frame, last_time, and — once vision-language features have
|
||||
been fed in — a conf-weighted running-mean feature in fp16 plus the
|
||||
weight sum.
|
||||
|
||||
Storage is hybrid: a Python dict maps voxel index ``(ix, iy, iz)`` to a
|
||||
row in column-stored numpy arrays so lookup is O(1) and bulk arithmetic
|
||||
stays vectorized. ``carve`` removes voxels that fall inside a view's
|
||||
observed free space (DynaMem-style dynamic updates); ``query`` returns
|
||||
the top-k cosine matches against a text embedding.
|
||||
|
||||
Default voxel size is 5 cm. The map is geometry-only until
|
||||
``add(..., feat_map=...)`` supplies per-pixel features; occupancy /
|
||||
planning use only the geometry, so they work without any features.
|
||||
"""
|
||||
|
||||
# ruff: noqa: N806 — H, W, D are conventional array-dimension names (and appear verbatim in error strings)
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VoxelMapStats:
|
||||
"""Per-keyframe deltas, surfaced to scalar logs."""
|
||||
|
||||
n_voxels: int
|
||||
n_added: int
|
||||
n_updated: int
|
||||
n_removed: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VoxelSnapshot:
|
||||
"""Current voxel map state, materialized for visualization / export."""
|
||||
|
||||
xyz: np.ndarray # (M, 3) float32 — count-weighted mean position
|
||||
rgb: np.ndarray # (M, 3) uint8 — count-weighted mean color (RGB)
|
||||
count: np.ndarray # (M,) int64
|
||||
last_frame: np.ndarray # (M,) int64
|
||||
last_time: np.ndarray # (M,) float64
|
||||
feat: np.ndarray | None = None # (M, D) fp16 — L2-normalized per-voxel mean
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CarveResult:
|
||||
"""Output of one ``carve`` pass."""
|
||||
|
||||
n_removed: int
|
||||
removed_xyz: np.ndarray # (K, 3) float32 — centres of removed voxels, for viz
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueryResult:
|
||||
"""Top-k cosine matches against a text embedding."""
|
||||
|
||||
xyz: np.ndarray # (k, 3) float32
|
||||
score: np.ndarray # (k,) float32 — cosine similarity in [-1, 1]
|
||||
voxel_indices: np.ndarray # (k,) int64 — row indices into the map
|
||||
|
||||
|
||||
_MAX_ABS_VOXEL_INDEX = 1 << 20
|
||||
_FEAT_CHUNK_PIXELS = 16384 # bound peak per-keyframe feature contribution memory
|
||||
|
||||
|
||||
class VoxelMap:
|
||||
"""Sparse-hash voxel grid with count-weighted means and semantic features."""
|
||||
|
||||
def __init__(self, voxel_size: float = 0.05) -> None:
|
||||
if voxel_size <= 0:
|
||||
raise ValueError("voxel_size must be > 0")
|
||||
self.voxel_size = float(voxel_size)
|
||||
|
||||
self._lookup: dict[tuple[int, int, int], int] = {}
|
||||
self._idx = np.zeros((0, 3), dtype=np.int64)
|
||||
self._count = np.zeros(0, dtype=np.int64)
|
||||
self._xyz_sum = np.zeros((0, 3), dtype=np.float64)
|
||||
self._rgb_sum = np.zeros((0, 3), dtype=np.float64)
|
||||
self._last_frame = np.zeros(0, dtype=np.int64)
|
||||
self._last_time = np.zeros(0, dtype=np.float64)
|
||||
|
||||
# Lazily allocated on first add() with feat_map.
|
||||
self._feature_dim: int | None = None
|
||||
self._feat_sum: np.ndarray | None = None # (M, D) fp16
|
||||
self._feat_weight: np.ndarray | None = None # (M,) fp32
|
||||
|
||||
def __len__(self) -> int:
|
||||
return int(self._count.shape[0])
|
||||
|
||||
@property
|
||||
def feature_dim(self) -> int | None:
|
||||
return self._feature_dim
|
||||
|
||||
# ------------------------------------------------------------------ add
|
||||
|
||||
def add(
|
||||
self,
|
||||
points: np.ndarray,
|
||||
rgb: np.ndarray,
|
||||
conf: np.ndarray,
|
||||
frame: int,
|
||||
t: float,
|
||||
conf_thresh: float = 0.5,
|
||||
feat_map: np.ndarray | None = None,
|
||||
) -> VoxelMapStats:
|
||||
"""Insert / update voxels from a per-pixel observation.
|
||||
|
||||
``points``: ``(..., 3)`` world xyz, fp32.
|
||||
``rgb``: ``(..., 3)`` uint8 (RGB order).
|
||||
``conf``: ``(...,)`` in [0, 1].
|
||||
``feat_map``: optional ``(..., D)`` fp16 per-pixel feature, already
|
||||
bilinearly upsampled to the points/conf grid. First
|
||||
call with features locks the feature dimension;
|
||||
subsequent calls must match.
|
||||
"""
|
||||
pts = np.asarray(points).reshape(-1, 3)
|
||||
cols = np.asarray(rgb).reshape(-1, 3)
|
||||
cnf = np.asarray(conf).reshape(-1)
|
||||
if not (len(pts) == len(cols) == len(cnf)):
|
||||
raise ValueError(f"length mismatch: points={len(pts)}, rgb={len(cols)}, conf={len(cnf)}")
|
||||
|
||||
features: np.ndarray | None = None
|
||||
if feat_map is not None:
|
||||
features = np.asarray(feat_map).reshape(-1, feat_map.shape[-1])
|
||||
if len(features) != len(pts):
|
||||
raise ValueError(f"feat_map length {len(features)} != points length {len(pts)}")
|
||||
D = features.shape[-1]
|
||||
if self._feature_dim is None:
|
||||
self._feature_dim = int(D)
|
||||
# Pad pre-existing voxels (added before features arrived) with zeros.
|
||||
self._feat_sum = np.zeros((len(self), D), dtype=np.float16)
|
||||
self._feat_weight = np.zeros(len(self), dtype=np.float32)
|
||||
LOG.info("VoxelMap features enabled: D=%d (fp16 storage)", D)
|
||||
elif self._feature_dim != D:
|
||||
raise ValueError(f"feature dim mismatch: existing={self._feature_dim}, got={D}")
|
||||
|
||||
mask = (cnf >= conf_thresh) & np.isfinite(pts).all(axis=1)
|
||||
pts = pts[mask]
|
||||
cols = cols[mask]
|
||||
cnf_kept = cnf[mask]
|
||||
if features is not None:
|
||||
features = features[mask]
|
||||
if pts.size == 0:
|
||||
return VoxelMapStats(n_voxels=len(self), n_added=0, n_updated=0)
|
||||
|
||||
idx = np.floor(pts / self.voxel_size).astype(np.int64)
|
||||
sane = (np.abs(idx) < _MAX_ABS_VOXEL_INDEX).all(axis=1)
|
||||
if not sane.all():
|
||||
n_drop = int((~sane).sum())
|
||||
LOG.debug("dropping %d points with extreme voxel index", n_drop)
|
||||
idx = idx[sane]
|
||||
pts = pts[sane]
|
||||
cols = cols[sane]
|
||||
cnf_kept = cnf_kept[sane]
|
||||
if features is not None:
|
||||
features = features[sane]
|
||||
if idx.size == 0:
|
||||
return VoxelMapStats(n_voxels=len(self), n_added=0, n_updated=0)
|
||||
|
||||
unique_idx, inverse = np.unique(idx, axis=0, return_inverse=True)
|
||||
inverse = inverse.reshape(-1)
|
||||
n_unique = unique_idx.shape[0]
|
||||
kf_count = np.bincount(inverse, minlength=n_unique).astype(np.int64)
|
||||
kf_xyz_sum = np.zeros((n_unique, 3), dtype=np.float64)
|
||||
kf_rgb_sum = np.zeros((n_unique, 3), dtype=np.float64)
|
||||
np.add.at(kf_xyz_sum, inverse, pts.astype(np.float64))
|
||||
np.add.at(kf_rgb_sum, inverse, cols.astype(np.float64))
|
||||
|
||||
kf_feat_sum: np.ndarray | None = None
|
||||
kf_feat_weight: np.ndarray | None = None
|
||||
if features is not None:
|
||||
kf_feat_sum = np.zeros((n_unique, self._feature_dim), dtype=np.float32)
|
||||
kf_feat_weight = np.zeros(n_unique, dtype=np.float32)
|
||||
cnf_f = cnf_kept.astype(np.float32)
|
||||
# Chunked accumulation — keeps the (chunk, D) intermediate small.
|
||||
for s in range(0, features.shape[0], _FEAT_CHUNK_PIXELS):
|
||||
e = s + _FEAT_CHUNK_PIXELS
|
||||
w = cnf_f[s:e]
|
||||
contrib = w[:, None] * features[s:e].astype(np.float32)
|
||||
np.add.at(kf_feat_sum, inverse[s:e], contrib)
|
||||
np.add.at(kf_feat_weight, inverse[s:e], w)
|
||||
|
||||
existing_rows: list[int] = []
|
||||
existing_local: list[int] = []
|
||||
new_local: list[int] = []
|
||||
new_keys: list[tuple[int, int, int]] = []
|
||||
for i in range(n_unique):
|
||||
key = (int(unique_idx[i, 0]), int(unique_idx[i, 1]), int(unique_idx[i, 2]))
|
||||
row = self._lookup.get(key)
|
||||
if row is None:
|
||||
new_local.append(i)
|
||||
new_keys.append(key)
|
||||
else:
|
||||
existing_rows.append(row)
|
||||
existing_local.append(i)
|
||||
|
||||
if existing_rows:
|
||||
rows = np.asarray(existing_rows, dtype=np.int64)
|
||||
local = np.asarray(existing_local, dtype=np.int64)
|
||||
self._count[rows] += kf_count[local]
|
||||
self._xyz_sum[rows] += kf_xyz_sum[local]
|
||||
self._rgb_sum[rows] += kf_rgb_sum[local]
|
||||
self._last_frame[rows] = frame
|
||||
self._last_time[rows] = t
|
||||
if kf_feat_sum is not None:
|
||||
assert self._feat_sum is not None and self._feat_weight is not None
|
||||
# fp32 accumulator -> fp16 storage; cast on store to match storage dtype.
|
||||
self._feat_sum[rows] = (self._feat_sum[rows].astype(np.float32) + kf_feat_sum[local]).astype(
|
||||
np.float16
|
||||
)
|
||||
self._feat_weight[rows] += kf_feat_weight[local]
|
||||
|
||||
if new_local:
|
||||
base = len(self)
|
||||
local = np.asarray(new_local, dtype=np.int64)
|
||||
self._idx = np.concatenate([self._idx, unique_idx[local]], axis=0)
|
||||
self._count = np.concatenate([self._count, kf_count[local]])
|
||||
self._xyz_sum = np.concatenate([self._xyz_sum, kf_xyz_sum[local]], axis=0)
|
||||
self._rgb_sum = np.concatenate([self._rgb_sum, kf_rgb_sum[local]], axis=0)
|
||||
self._last_frame = np.concatenate(
|
||||
[self._last_frame, np.full(len(new_local), frame, dtype=np.int64)]
|
||||
)
|
||||
self._last_time = np.concatenate([self._last_time, np.full(len(new_local), t, dtype=np.float64)])
|
||||
if self._feature_dim is not None:
|
||||
assert self._feat_sum is not None and self._feat_weight is not None
|
||||
if kf_feat_sum is not None:
|
||||
new_feats = kf_feat_sum[local].astype(np.float16)
|
||||
new_weights = kf_feat_weight[local]
|
||||
else:
|
||||
# Features enabled, but this call didn't bring any — pad zeros
|
||||
# so array sizes stay aligned with _count.
|
||||
new_feats = np.zeros((len(new_local), self._feature_dim), dtype=np.float16)
|
||||
new_weights = np.zeros(len(new_local), dtype=np.float32)
|
||||
self._feat_sum = np.concatenate([self._feat_sum, new_feats], axis=0)
|
||||
self._feat_weight = np.concatenate([self._feat_weight, new_weights])
|
||||
for offset, key in enumerate(new_keys):
|
||||
self._lookup[key] = base + offset
|
||||
|
||||
return VoxelMapStats(
|
||||
n_voxels=len(self),
|
||||
n_added=len(new_local),
|
||||
n_updated=len(existing_rows),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------- hard-delete
|
||||
def remove_voxels_in_box(
|
||||
self,
|
||||
xyz_min: tuple[float, float, float],
|
||||
xyz_max: tuple[float, float, float],
|
||||
) -> int:
|
||||
"""Surgical hard-delete of every voxel whose mean position lies inside
|
||||
the axis-aligned bounding box.
|
||||
|
||||
Different from :meth:`carve` (DynaMem-style free-space removal from a
|
||||
camera frustum + depth). This one is for simulated scene mutation:
|
||||
"the couch moved away" is removing the box around the old couch then
|
||||
``add()``-ing one at the new position.
|
||||
"""
|
||||
if len(self) == 0:
|
||||
return 0
|
||||
cnt = self._count.astype(np.float64).reshape(-1, 1)
|
||||
means = self._xyz_sum / cnt
|
||||
in_box = (
|
||||
(means[:, 0] >= xyz_min[0])
|
||||
& (means[:, 0] <= xyz_max[0])
|
||||
& (means[:, 1] >= xyz_min[1])
|
||||
& (means[:, 1] <= xyz_max[1])
|
||||
& (means[:, 2] >= xyz_min[2])
|
||||
& (means[:, 2] <= xyz_max[2])
|
||||
)
|
||||
if not in_box.any():
|
||||
return 0
|
||||
keep = ~in_box
|
||||
n_removed = int(in_box.sum())
|
||||
for k in self._idx[in_box]:
|
||||
del self._lookup[(int(k[0]), int(k[1]), int(k[2]))]
|
||||
self._idx = self._idx[keep]
|
||||
self._count = self._count[keep]
|
||||
self._xyz_sum = self._xyz_sum[keep]
|
||||
self._rgb_sum = self._rgb_sum[keep]
|
||||
self._last_frame = self._last_frame[keep]
|
||||
self._last_time = self._last_time[keep]
|
||||
if self._feat_sum is not None and self._feat_weight is not None:
|
||||
self._feat_sum = self._feat_sum[keep]
|
||||
self._feat_weight = self._feat_weight[keep]
|
||||
# Row indices shifted — rebuild the lookup.
|
||||
self._lookup = {
|
||||
(int(self._idx[i, 0]), int(self._idx[i, 1]), int(self._idx[i, 2])): i
|
||||
for i in range(len(self._idx))
|
||||
}
|
||||
return n_removed
|
||||
|
||||
# ---------------------------------------------------------------- carve
|
||||
|
||||
def carve(
|
||||
self,
|
||||
local_points: np.ndarray,
|
||||
conf: np.ndarray,
|
||||
pose: np.ndarray,
|
||||
focal_px: float,
|
||||
frame: int,
|
||||
t: float,
|
||||
conf_thresh: float = 0.5,
|
||||
margin: float = 0.05,
|
||||
) -> CarveResult:
|
||||
"""Remove voxels inside this view's observed free space.
|
||||
|
||||
A voxel is carved when it projects into the image, sits in front of
|
||||
the camera, and lies closer than the observed depth (minus a margin)
|
||||
at that pixel — i.e. we can see through where it claims to be. Carve
|
||||
runs before ``add`` each keyframe so moved/removed objects disappear.
|
||||
"""
|
||||
if len(self) == 0:
|
||||
return CarveResult(0, np.zeros((0, 3), dtype=np.float32))
|
||||
|
||||
if local_points.ndim != 3 or local_points.shape[-1] != 3:
|
||||
raise ValueError(f"expected (H, W, 3), got {local_points.shape}")
|
||||
if conf.shape != local_points.shape[:2]:
|
||||
raise ValueError(f"conf shape {conf.shape} != local_points (H, W) {local_points.shape[:2]}")
|
||||
if pose.shape != (4, 4):
|
||||
raise ValueError(f"pose must be (4, 4); got {pose.shape}")
|
||||
|
||||
H, W = local_points.shape[:2]
|
||||
cx = (W - 1) / 2.0
|
||||
cy = (H - 1) / 2.0
|
||||
depth_map = local_points[..., 2]
|
||||
|
||||
cnt = self._count.astype(np.float64).reshape(-1, 1)
|
||||
xyz_world = self._xyz_sum / cnt
|
||||
|
||||
R = pose[:3, :3].astype(np.float64)
|
||||
t_vec = pose[:3, 3].astype(np.float64)
|
||||
xyz_cam = (xyz_world - t_vec[None, :]) @ R
|
||||
|
||||
d_voxel = xyz_cam[:, 2]
|
||||
front = d_voxel > 1e-3
|
||||
|
||||
d_safe = np.where(front, d_voxel, 1.0)
|
||||
u = focal_px * xyz_cam[:, 0] / d_safe + cx
|
||||
v = focal_px * xyz_cam[:, 1] / d_safe + cy
|
||||
in_bounds = (u >= 0.0) & (u < W) & (v >= 0.0) & (v < H)
|
||||
valid = front & in_bounds
|
||||
|
||||
u_i = np.clip(np.floor(u).astype(np.int64), 0, W - 1)
|
||||
v_i = np.clip(np.floor(v).astype(np.int64), 0, H - 1)
|
||||
D_at = depth_map[v_i, u_i]
|
||||
C_at = conf[v_i, u_i]
|
||||
|
||||
finite_D = np.isfinite(D_at) & (D_at > 0.0)
|
||||
free_space = valid & finite_D & (C_at >= conf_thresh) & (d_voxel < (D_at - margin))
|
||||
|
||||
n_removed = int(free_space.sum())
|
||||
if n_removed == 0:
|
||||
return CarveResult(0, np.zeros((0, 3), dtype=np.float32))
|
||||
|
||||
removed_xyz = xyz_world[free_space].astype(np.float32)
|
||||
removed_keys = self._idx[free_space]
|
||||
for k in removed_keys:
|
||||
del self._lookup[(int(k[0]), int(k[1]), int(k[2]))]
|
||||
|
||||
keep = ~free_space
|
||||
self._idx = self._idx[keep]
|
||||
self._count = self._count[keep]
|
||||
self._xyz_sum = self._xyz_sum[keep]
|
||||
self._rgb_sum = self._rgb_sum[keep]
|
||||
self._last_frame = self._last_frame[keep]
|
||||
self._last_time = self._last_time[keep]
|
||||
if self._feat_sum is not None and self._feat_weight is not None:
|
||||
self._feat_sum = self._feat_sum[keep]
|
||||
self._feat_weight = self._feat_weight[keep]
|
||||
|
||||
self._lookup = {
|
||||
(int(self._idx[i, 0]), int(self._idx[i, 1]), int(self._idx[i, 2])): i
|
||||
for i in range(len(self._idx))
|
||||
}
|
||||
LOG.debug("carve frame=%d t=%.3fs removed=%d", frame, t, n_removed)
|
||||
return CarveResult(n_removed=n_removed, removed_xyz=removed_xyz)
|
||||
|
||||
# ------------------------------------------------------------- snapshot
|
||||
|
||||
def snapshot(self, include_features: bool = False) -> VoxelSnapshot:
|
||||
"""Materialize the current map.
|
||||
|
||||
``include_features``: pay the cost of normalizing the per-voxel
|
||||
feature mean. Off by default — visualization doesn't need features.
|
||||
"""
|
||||
if len(self) == 0:
|
||||
return VoxelSnapshot(
|
||||
xyz=np.zeros((0, 3), dtype=np.float32),
|
||||
rgb=np.zeros((0, 3), dtype=np.uint8),
|
||||
count=np.zeros(0, dtype=np.int64),
|
||||
last_frame=np.zeros(0, dtype=np.int64),
|
||||
last_time=np.zeros(0, dtype=np.float64),
|
||||
feat=None,
|
||||
)
|
||||
cnt = self._count.astype(np.float64).reshape(-1, 1)
|
||||
xyz = (self._xyz_sum / cnt).astype(np.float32)
|
||||
rgb = np.clip(self._rgb_sum / cnt, 0, 255).astype(np.uint8)
|
||||
|
||||
feat = None
|
||||
if include_features and self._feat_sum is not None and self._feat_weight is not None:
|
||||
feat = self._normalized_features()
|
||||
|
||||
return VoxelSnapshot(
|
||||
xyz=xyz,
|
||||
rgb=rgb,
|
||||
count=self._count.copy(),
|
||||
last_frame=self._last_frame.copy(),
|
||||
last_time=self._last_time.copy(),
|
||||
feat=feat,
|
||||
)
|
||||
|
||||
def _normalized_features(self) -> np.ndarray:
|
||||
"""Per-voxel L2-normalized feature mean. (M, D) fp16."""
|
||||
assert self._feat_sum is not None and self._feat_weight is not None
|
||||
w = np.maximum(self._feat_weight, 1e-6).reshape(-1, 1)
|
||||
mean = self._feat_sum.astype(np.float32) / w
|
||||
norms = np.linalg.norm(mean, axis=1, keepdims=True)
|
||||
mean = mean / np.maximum(norms, 1e-6)
|
||||
return mean.astype(np.float16)
|
||||
|
||||
# ----------------------------------------------------------------- query
|
||||
|
||||
def query(self, text_embedding: np.ndarray, top_k: int = 32) -> QueryResult:
|
||||
"""Top-k cosine matches against ``text_embedding``.
|
||||
|
||||
``text_embedding``: ``(D,)`` array — does NOT need to be unit norm;
|
||||
we re-normalize.
|
||||
"""
|
||||
if self._feat_sum is None or self._feature_dim is None:
|
||||
raise RuntimeError("VoxelMap has no semantic features yet — call add(..., feat_map=...) first")
|
||||
if len(self) == 0:
|
||||
return QueryResult(
|
||||
xyz=np.zeros((0, 3), dtype=np.float32),
|
||||
score=np.zeros(0, dtype=np.float32),
|
||||
voxel_indices=np.zeros(0, dtype=np.int64),
|
||||
)
|
||||
if text_embedding.shape != (self._feature_dim,):
|
||||
raise ValueError(f"text_embedding shape {text_embedding.shape} != ({self._feature_dim},)")
|
||||
|
||||
voxel_feat = self._normalized_features().astype(np.float32)
|
||||
text_unit = text_embedding.astype(np.float32)
|
||||
text_unit = text_unit / max(float(np.linalg.norm(text_unit)), 1e-6)
|
||||
|
||||
scores = voxel_feat @ text_unit # (M,)
|
||||
k = min(int(top_k), len(scores))
|
||||
# Partition-and-sort for the top-k.
|
||||
top_idx = np.argpartition(scores, -k)[-k:]
|
||||
order = np.argsort(-scores[top_idx])
|
||||
top_idx = top_idx[order]
|
||||
|
||||
snap_xyz = (self._xyz_sum[top_idx] / self._count[top_idx].astype(np.float64).reshape(-1, 1)).astype(
|
||||
np.float32
|
||||
)
|
||||
return QueryResult(
|
||||
xyz=snap_xyz,
|
||||
score=scores[top_idx].astype(np.float32),
|
||||
voxel_indices=top_idx.astype(np.int64),
|
||||
)
|
||||
|
||||
# --------------------------------------------------------- introspection
|
||||
|
||||
def memory_bytes(self) -> dict[str, int]:
|
||||
"""Return per-array memory footprint."""
|
||||
out = {
|
||||
"idx": self._idx.nbytes,
|
||||
"count": self._count.nbytes,
|
||||
"xyz_sum": self._xyz_sum.nbytes,
|
||||
"rgb_sum": self._rgb_sum.nbytes,
|
||||
"last_frame": self._last_frame.nbytes,
|
||||
"last_time": self._last_time.nbytes,
|
||||
"lookup_dict": _approx_dict_bytes(self._lookup),
|
||||
}
|
||||
if self._feat_sum is not None:
|
||||
out["feat_sum"] = self._feat_sum.nbytes
|
||||
assert self._feat_weight is not None
|
||||
out["feat_weight"] = self._feat_weight.nbytes
|
||||
out["total"] = sum(v for k, v in out.items() if k != "total")
|
||||
return out
|
||||
|
||||
|
||||
def _approx_dict_bytes(d: dict) -> int:
|
||||
"""Rough lower-bound estimate; ~100 bytes/entry is a fine ballpark."""
|
||||
return 100 * len(d)
|
||||
Reference in New Issue
Block a user