From e2cbb3ebbb053c9f9077fde5215db015f095b31d Mon Sep 17 00:00:00 2001 From: bogwi Date: Thu, 20 Aug 2026 12:46:16 +0900 Subject: [PATCH 01/28] init generalizing refactor --- .../type/detection3d/pointcloud_filters.py | 32 ++ dimos/perception/memory/dandetect.py | 45 +- dimos/perception/memory/gates.py | 190 +------- dimos/perception/memory/inventory.py | 156 +++---- dimos/perception/memory/localize.py | 144 ++---- dimos/perception/memory/rig.py | 430 ++++++++++++++++++ dimos/perception/memory/support_plane.py | 44 +- dimos/perception/memory/tool_inventory.py | 67 +-- dimos/perception/memory/tool_localize.py | 92 ++-- dimos/perception/memory/types.py | 25 + 10 files changed, 722 insertions(+), 503 deletions(-) create mode 100644 dimos/perception/memory/rig.py diff --git a/dimos/perception/detection/type/detection3d/pointcloud_filters.py b/dimos/perception/detection/type/detection3d/pointcloud_filters.py index fdb2afeebb..f55d0d09f3 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud_filters.py +++ b/dimos/perception/detection/type/detection3d/pointcloud_filters.py @@ -48,6 +48,38 @@ def filter_func( return filter_func +def range_cluster(gap: float = 0.3) -> PointCloudFilter: + """Keep the camera-range cluster containing the median range. + + The projected-cloud analog of the depth-gap split in ``from_depth``: + points a mask collects across a range discontinuity are background seen + through or around the object, not the object. + """ + import numpy as np + + def filter_func( + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + ) -> PointCloud2 | None: + points, _ = pc.as_numpy() + if len(points) == 0: + return None + camera = tf.inverse().translation.to_numpy() + ranges = np.linalg.norm(points - camera, axis=1) + order = np.argsort(ranges) + ranges_sorted = ranges[order] + gaps = np.nonzero(np.diff(ranges_sorted) > gap)[0] + starts = np.concatenate(([0], gaps + 1)) + ends = np.concatenate((gaps + 1, [len(ranges_sorted)])) + median_idx = np.searchsorted(ranges_sorted, np.median(ranges_sorted)) + for start, end in zip(starts, ends, strict=False): + if start <= median_idx < end: + keep = order[start:end] + return PointCloud2.from_numpy(points[keep], frame_id=pc.frame_id, timestamp=pc.ts) + return pc + + return filter_func + + def raycast() -> PointCloudFilter: def filter_func( det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index 081cb52e61..f7c3715cd1 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -17,6 +17,10 @@ ``DanDetector`` owns the models behind :func:`embed_index`, :func:`localize`, and :func:`inventory`: enter once, query many times on warm weights, and ``stop()`` (or leave the ``with`` block) releases whatever loaded. + +Every entry point takes an optional :class:`~dimos.perception.memory.rig.Rig` +describing where poses and 3D geometry come from; without one the store's +shape decides. """ from __future__ import annotations @@ -25,12 +29,10 @@ from dimos.core.resource import Resource from dimos.memory.embed import EmbedImages -from dimos.memory.tf import StreamTF from dimos.memory.transform import throttle -from dimos.perception.memory import gates -from dimos.perception.memory.gates import OPTICAL_FRAME, TF_TOLERANCE, WORLD_FRAME from dimos.perception.memory.inventory import DEFAULT_VOCABULARY, NamingVocabulary, inventory -from dimos.perception.memory.localize import EMBED_HZ, embed_index, localize +from dimos.perception.memory.localize import embed_index, localize +from dimos.perception.memory.rig import Rig if TYPE_CHECKING: from reactivex.abc import DisposableBase @@ -79,9 +81,7 @@ def embed( before: float, *, live: Literal[False] = False, - optical_frame: str = ..., - world_frame: str = ..., - tf_tolerance: float = ..., + rig: Rig | None = ..., ) -> Stream[Any, Any]: ... @overload def embed( @@ -89,9 +89,7 @@ def embed( store: Any, *, live: Literal[True], - optical_frame: str = ..., - world_frame: str = ..., - tf_tolerance: float = ..., + rig: Rig | None = ..., ) -> Stream[Any, Any]: ... def embed( self, @@ -100,43 +98,32 @@ def embed( before: float | None = None, *, live: bool = False, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tf_tolerance: float = TF_TOLERANCE, + rig: Rig | None = None, ) -> Stream[Any, Any]: """SigLIP-embedded, world-posed frame index for :meth:`localize`. Replay mode indexes ``[after, before]`` in memory and returns when - done. ``live=True`` instead tails ``color_image`` and keeps saving - into the store's named ``color_image_embedded`` stream on a + done. ``live=True`` instead tails the rig's color stream and keeps + saving into the store's named ``color_image_embedded`` stream on a background thread; the returned stream is that named stream. """ + rig = rig or Rig.from_store(store) if not live: return embed_index( store, self.siglip, cast("float", after), cast("float", before), - optical_frame=optical_frame, - world_frame=world_frame, - tf_tolerance=tf_tolerance, + rig=rig, ) from dimos.msgs.sensor_msgs.Image import Image - tf = StreamTF.from_store(store) - if tf is None: - raise ValueError("store has no tf stream") embedded: Stream[Any, Any] = store.stream("color_image_embedded", Image) pipeline = ( - store.streams.color_image.live() - .transform(throttle(1.0 / EMBED_HZ)) - .map( - lambda obs: obs.derive( - data=obs.data, - pose=gates.camera_pose(tf, obs.ts, optical_frame, world_frame, tf_tolerance), - ) - ) + rig.color.live() + .transform(throttle(1.0 / rig.embed_hz)) + .map(lambda obs: obs.derive(data=obs.data, pose=rig.index_pose(obs))) .filter(lambda obs: obs.pose is not None) .transform(EmbedImages(self.siglip, batch_size=1)) .save(embedded) diff --git a/dimos/perception/memory/gates.py b/dimos/perception/memory/gates.py index 99ade4ccdd..145f25efdd 100644 --- a/dimos/perception/memory/gates.py +++ b/dimos/perception/memory/gates.py @@ -12,38 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-frame gates and lookups over a memory recording: poses, stillness, frames. - -Two per-frame gates live here and both are required: - -* The **camera-motion gate** differences tf. It rejects frames captured while - the wrist sweeps, because a stale or interpolated transform smears the - projection. -* The **scene-motion gate** differences images, conditioned on camera - stillness at both compared instants. It rejects frames captured while the - scene itself changes - a parked camera watching hands rearrange objects is - exactly the case tf cannot see. The reference frame is anchored at the - start of the surrounding camera-still interval, so a frame is trusted only - while the scene still matches the state it had when the camera parked. - -Every function here takes the store and/or tf plus primitives; the stillness -intervals and the grayscale memo are allocated by the caller and passed in, -one per query. +"""The scene-motion gate: image differencing conditioned on camera stillness. + +Pose-derived gates (camera pose, speed, stillness, keyframes) live on +:class:`~dimos.perception.memory.rig.Rig` - they depend on where the rig's +poses come from. What stays here is purely image-based: the scene-motion +gate differences frames of the color stream, conditioned on camera +stillness at both compared instants. It rejects frames captured while the +scene itself changes - a parked camera watching hands rearrange objects is +exactly the case poses cannot see. The reference frame is anchored at the +start of the surrounding camera-still interval, so a frame is trusted only +while the scene still matches the state it had when the camera parked. + +The stillness intervals and the grayscale memo are allocated by the caller +and passed in, one per query. """ from __future__ import annotations from bisect import bisect_right -from typing import TYPE_CHECKING, Any +from typing import Any import numpy as np -if TYPE_CHECKING: - from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.sensor_msgs.Image import Image - from dimos.protocol.tf.tf import TFLookup - OPTICAL_FRAME = "camera_color_optical_frame" WORLD_FRAME = "world" @@ -60,82 +51,6 @@ DIFF_WIDTH = 212 # px - diff resolution (1/4 of 848) -def camera_pose( - tf: TFLookup, - ts: float, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tolerance: float = TF_TOLERANCE, -) -> PoseStamped | None: - """World pose of the camera optical frame at ts - it rides the wrist.""" - transform = tf.get(optical_frame, world_frame, ts, tolerance) - return (-transform).to_pose() if transform is not None else None - - -def camera_speed( - tf: TFLookup, - ts: float, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tolerance: float = TF_TOLERANCE, - dt: float = 0.06, -) -> float | None: - """Linear speed of the camera (m/s) around ts, from tf differencing.""" - a = camera_pose(tf, ts - dt, optical_frame, world_frame, tolerance) - b = camera_pose(tf, ts + dt, optical_frame, world_frame, tolerance) - if a is None or b is None: - return None - return float((b.position - a.position).magnitude() / (2 * dt)) - - -def camera_still( - tf: TFLookup, - ts: float, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tolerance: float = TF_TOLERANCE, - speed_max: float = SPEED_MAX, - envelope: float = STILL_ENVELOPE, -) -> bool: - """Camera is still over the whole capture envelope, not just at ts.""" - for offset in (-envelope, 0.0, envelope): - speed = camera_speed(tf, ts + offset, optical_frame, world_frame, tolerance) - if speed is None or speed > speed_max: - return False - return True - - -def still_intervals( - tf: TFLookup, - t0: float, - t1: float, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tolerance: float = TF_TOLERANCE, - speed_max: float = SPEED_MAX, -) -> list[tuple[float, float]]: - """Maximal camera-still intervals inside [t0, t1], sampled at 0.25 s. - - Computed once per query by the caller; every scene-motion query resolves - its surrounding interval from this list. - """ - step = 0.25 - times = np.arange(t0, t1 + step, step) - intervals: list[tuple[float, float]] = [] - run_start: float | None = None - for t in times: - speed = camera_speed(tf, float(t), optical_frame, world_frame, tolerance) - still = speed is not None and speed <= speed_max - if still and run_start is None: - run_start = float(t) - elif not still and run_start is not None: - intervals.append((run_start, float(t) - step)) - run_start = None - if run_start is not None: - intervals.append((run_start, float(times[-1]))) - return [(a, b) for a, b in intervals if b >= a] - - def _interval_containing( ts: float, intervals: list[tuple[float, float]] ) -> tuple[float, float] | None: @@ -146,7 +61,7 @@ def _interval_containing( return (a, b) if a - 0.25 <= ts <= b + 0.25 else None -def _gray_small(store: Any, ts: float, gray: dict[float, np.ndarray | None]) -> np.ndarray | None: +def _gray_small(color: Any, ts: float, gray: dict[float, np.ndarray | None]) -> np.ndarray | None: """Downscaled grayscale of the color frame nearest ts, memoized in *gray*.""" key = round(ts, 2) if key in gray: @@ -156,7 +71,7 @@ def _gray_small(store: Any, ts: float, gray: dict[float, np.ndarray | None]) -> small_gray: np.ndarray | None = None try: - frame = store.streams.color_image.at(ts, 0.1).first().data + frame = color.at(ts, 0.1).first().data except LookupError: frame = None if frame is not None: @@ -170,10 +85,10 @@ def _gray_small(store: Any, ts: float, gray: dict[float, np.ndarray | None]) -> def _diff_fraction( - store: Any, ts_a: float, ts_b: float, gray: dict[float, np.ndarray | None] + color: Any, ts_a: float, ts_b: float, gray: dict[float, np.ndarray | None] ) -> float | None: """Fraction of pixels changed between the frames nearest the two instants.""" - a, b = _gray_small(store, ts_a, gray), _gray_small(store, ts_b, gray) + a, b = _gray_small(color, ts_a, gray), _gray_small(color, ts_b, gray) if a is None or b is None or a.shape != b.shape: return None delta = np.abs(a.astype(np.int16) - b.astype(np.int16)) @@ -181,7 +96,7 @@ def _diff_fraction( def scene_still( - store: Any, + color: Any, ts: float, intervals: list[tuple[float, float]], gray: dict[float, np.ndarray | None], @@ -206,75 +121,14 @@ def scene_still( a, b = interval anchor = min(a + 0.3, ts) - fraction = _diff_fraction(store, anchor, ts, gray) + fraction = _diff_fraction(color, anchor, ts, gray) if fraction is None or fraction > motion_threshold: return False for other in (max(a, ts - SHORT_DIFF_DT), min(b, ts + SHORT_DIFF_DT)): if abs(other - ts) < 0.05: continue - fraction = _diff_fraction(store, other, ts, gray) + fraction = _diff_fraction(color, other, ts, gray) if fraction is None or fraction > motion_threshold: return False return True - - -def depth_at(store: Any, ts: float, tolerance: float = 0.06) -> Image | None: - """Temporal join: aligned depth frame for a color timestamp.""" - try: - depth: Image = store.streams.depth_image.at(ts, tolerance).first().data - except LookupError: - return None - return depth - - -def keyframes( - store: Any, - tf: TFLookup, - t0: float, - t1: float, - stride: float, - intervals: list[tuple[float, float]], - gray: dict[float, np.ndarray | None], - motion_threshold: float = MOTION_THRESHOLD, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tolerance: float = TF_TOLERANCE, -) -> list[Observation[Image]]: - """Camera-still, scene-still color frames on a coarse grid over [t0, t1]. - - For each grid point the nearest passing frame within half a stride is - selected, so a grid point landing mid-sweep snaps to the neighboring - pause instead of being lost. - """ - selected: list[Observation[Image]] = [] - seen: set[float] = set() - offsets = [0.0] - probe = 0.35 - while probe <= stride / 2: - offsets.extend([probe, -probe]) - probe += 0.35 - - t = t0 + 0.5 - while t < t1: - for offset in offsets: - ts = t + offset - if ts < t0 or ts > t1: - continue - if not camera_still(tf, ts, optical_frame, world_frame, tolerance): - continue - if not scene_still(store, ts, intervals, gray, motion_threshold): - continue - try: - obs = store.streams.color_image.at(ts, 0.1).first() - except LookupError: - continue - if obs.ts in seen: - break - if tf.get(optical_frame, world_frame, obs.ts, tolerance) is None: - continue - seen.add(obs.ts) - selected.append(obs) - break - t += stride - return selected diff --git a/dimos/perception/memory/inventory.py b/dimos/perception/memory/inventory.py index f256701ee1..59a988bf02 100644 --- a/dimos/perception/memory/inventory.py +++ b/dimos/perception/memory/inventory.py @@ -38,15 +38,8 @@ import numpy as np -from dimos.memory.tf import StreamTF -from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC -from dimos.perception.memory import gates -from dimos.perception.memory.gates import ( - MOTION_THRESHOLD, - OPTICAL_FRAME, - TF_TOLERANCE, - WORLD_FRAME, -) +from dimos.perception.memory.gates import MOTION_THRESHOLD +from dimos.perception.memory.rig import Rig from dimos.perception.memory.support_plane import SupportPlane, fit_support_plane from dimos.perception.memory.types import ( Instance, @@ -63,24 +56,14 @@ from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector from dimos.perception.detection.type.detection2d.seg import Detection2DSeg - from dimos.protocol.tf.tf import TFLookup logger = setup_logger() -KEYFRAME_STRIDE = 2.5 # s - proposal keyframe grid MAX_PROPOSALS_PER_FRAME = 40 NAME_FRAMES_PER_INSTANCE = 5 -# An attachment must be the detector drawing a box around this member, not a -# box that merely crosses it. -NAME_ATTACH_IOU = 0.45 SUPPRESS_SCORE = 0.25 SUPPRESS_OVERLAP = 0.35 UNGROUNDED_TRACK_IOU = 0.40 -# The majority of a candidate's points must lie within the error envelope of -# the track's accumulated support. Partial and newly revealed views of one -# object satisfy this; a different object placed at a vacated rest position -# does not, which is what AABB overlap cannot express at tabletop scale. -SUPPORT_EXPLAINED = 0.5 # Groups of surface strings for one thing, canonical label first. Only groups # compete, so near-synonyms reinforce instead of splitting a box's score. A @@ -179,43 +162,36 @@ def _proposal_passes_2d(det: Detection2DSeg, image_area: float, policy: Inventor return True -# A lifted cloud plainly spanning more than one object: wider than any single -# tabletop object here, or reaching table-to-well-above-hand height. -SPLIT_EXTENT_M = 0.30 -SPLIT_HEIGHT_M = 0.10 -SPLIT_EPS_M = 0.03 - - def _split_oversized( points: np.ndarray, plane: SupportPlane | None, policy: InventoryPolicy ) -> list[np.ndarray]: """Re-segment a mask-bled cloud by 3D connectivity. - Automatic masks occasionally bleed across an object onto the table and - its neighbors; the lifted cloud then violates single-object bounds. The - repair is geometric: strip the support-surface points, then split by - spatial connectivity - distinct objects on this rig are separated by - more than the cluster gap, one object's surface is not. + Automatic masks occasionally bleed across an object onto its support + surface and its neighbors; the lifted cloud then violates single-object + bounds. The repair is geometric: strip the support-surface points, then + split by spatial connectivity - distinct objects are separated by more + than the cluster gap, one object's surface is not. """ extent = points.max(axis=0) - points.min(axis=0) - if float(extent.max()) <= SPLIT_EXTENT_M and float(extent[2]) <= SPLIT_HEIGHT_M: + if float(extent.max()) <= policy.split_extent_m and float(extent[2]) <= policy.split_height_m: return [points] if plane is None: return [points] heights = plane.height_above(points) - if float((np.abs(heights) <= 0.003).mean()) < 0.15: + if float((np.abs(heights) <= policy.min_height_above_plane_m).mean()) < 0.15: # No appreciable support-surface content: this is one oversized body, - # not a mask that bled across the table. Leave it to the extent cap. + # not a mask that bled across the surface. Leave it to the extent cap. return [points] - above = heights > 0.002 + above = heights > policy.min_height_above_plane_m * 2 / 3 body = points[above] if above.sum() >= policy.min_depth_points else points import open3d as o3d cloud = o3d.geometry.PointCloud() cloud.points = o3d.utility.Vector3dVector(body) - labels = np.asarray(cloud.cluster_dbscan(eps=SPLIT_EPS_M, min_points=20)) + labels = np.asarray(cloud.cluster_dbscan(eps=policy.split_eps_m, min_points=20)) clusters = [ body[labels == label] for label in range(labels.max() + 1) @@ -240,26 +216,20 @@ def _pixel_bbox( def _lift_frame( detections_2d: Any, - store: Any, - tf: TFLookup, - camera_info: CameraInfo, + rig: Rig, obs_ts: float, camera_position: np.ndarray, policy: InventoryPolicy, - optical_frame: str, - world_frame: str, - tf_tolerance: float, plane: SupportPlane | None = None, ) -> tuple[list[SupportObservation], list[Detection2DSeg]]: - """Depth-lift accepted proposals of one frame; returns (grounded, ungrounded).""" + """Lift accepted proposals of one frame through the rig; (grounded, ungrounded).""" from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - depth = gates.depth_at(store, obs_ts) - transform = tf.get(optical_frame, world_frame, obs_ts, tf_tolerance) - if depth is None or transform is None: + lifted = rig.lift(detections_2d) + transform = rig.world_to_optical(obs_ts) + if lifted is None or transform is None: return [], list(detections_2d) - - lifted = ImageDetections3DPC.from_depth(detections_2d, depth, camera_info, transform) + camera_info = rig.camera_info grounded: list[SupportObservation] = [] ungrounded: list[Detection2DSeg] = [] @@ -285,7 +255,7 @@ def _lift_frame( ts=obs_ts, cloud=det3d.pointcloud if whole - else PointCloud2.from_numpy(piece, frame_id="world", timestamp=obs_ts), + else PointCloud2.from_numpy(piece, frame_id=rig.world_frame, timestamp=obs_ts), centroid=piece.mean(axis=0), aabb_min=aabb_min, aabb_max=aabb_max, @@ -324,16 +294,17 @@ def _aabb_gap(a: SupportObservation, b: SupportObservation) -> float: return float(gap.max()) -def _cloud_gap(a: SupportObservation, b: SupportObservation) -> float: +def _cloud_gap(a: SupportObservation, b: SupportObservation, cut: float) -> float: """Minimum point-to-point distance between two observation clouds. The AABB gap is a poor contact test for diagonal objects - an axis-aligned box overhangs its object's true footprint and "touches" neighbors that are centimeters of clear table away. Actual cloud distance is the physical claim. The AABB test remains as a cheap - prefilter. + prefilter: beyond ``cut`` of box separation the exact distance cannot + matter to any caller. """ - if _aabb_gap(a, b) > 0.06: + if _aabb_gap(a, b) > cut: return np.inf from scipy.spatial import cKDTree @@ -385,7 +356,10 @@ def _merge_same_frame( changed = False for i in range(len(items)): for j in range(i + 1, len(items)): - if _cloud_gap(items[i], items[j]) <= policy.same_frame_merge_gap_m: + if ( + _cloud_gap(items[i], items[j], 3 * policy.same_frame_merge_gap_m) + <= policy.same_frame_merge_gap_m + ): _absorb_into(items[i], items[j]) items.pop(j) changed = True @@ -430,7 +404,7 @@ def _associate( continue t_lo, t_hi = track.aabb size_gap = np.abs((t_hi - t_lo) - (obs.aabb_max - obs.aabb_min)) - if float(size_gap.max()) > 0.25: + if float(size_gap.max()) > policy.size_gap_max_m: continue overlap = aabb_overlap( obs.aabb_min, obs.aabb_max, t_lo, t_hi, pad=policy.envelope_pad_m @@ -438,7 +412,7 @@ def _associate( if overlap < policy.overlap_accept: continue explained = _support_explained(obs_points, track.support_pts, policy.envelope_pad_m) - if explained < SUPPORT_EXPLAINED: + if explained < policy.support_explained: continue cost[i, j] = 1.0 - overlap @@ -477,7 +451,11 @@ def _tracks_are_fragments(a: _Track, b: _Track, policy: InventoryPolicy) -> bool shared = a.frame_ts & b.frame_ts for ts in shared: pairs_gap = min( - _cloud_gap(ma, mb) for ma in a.members if ma.ts == ts for mb in b.members if mb.ts == ts + _cloud_gap(ma, mb, 3 * policy.same_frame_merge_gap_m) + for ma in a.members + if ma.ts == ts + for mb in b.members + if mb.ts == ts ) if pairs_gap > 1.5 * policy.same_frame_merge_gap_m: return False @@ -514,7 +492,7 @@ def _merge_tracks(tracks: list[_Track], policy: InventoryPolicy) -> list[_Track] _support_explained(a.support_pts, b.support_pts, policy.envelope_pad_m), _support_explained(b.support_pts, a.support_pts, policy.envelope_pad_m), ) - if explained < SUPPORT_EXPLAINED: + if explained < policy.support_explained: continue for obs in b.members: a.add(obs, obs.ts) @@ -584,7 +562,7 @@ def _aggregated_label(labels: tuple[tuple[str, float], ...], policy: InventoryPo def _build_instance( - index: int, track: _Track, policy: InventoryPolicy, grounded: bool = True + index: int, track: _Track, policy: InventoryPolicy, frame_id: str, grounded: bool = True ) -> Instance: labels = tuple(sorted(track.labels.items(), key=lambda kv: -kv[1])) primary = _aggregated_label(labels, policy) @@ -602,7 +580,7 @@ def _build_instance( sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), coverage=coverage, axes_observed=axes_observed, - frame_id="world", + frame_id=frame_id, ) distinct_views = len({tuple(np.round(m.camera_position, 2)) for m in track.members}) return Instance( @@ -675,7 +653,7 @@ def _accepted_groups( def _name_and_suppress( tracks: list[_Track], tracks_2d: list[_Track2D], - store: Any, + color: Any, detector: Owlv2Detector, vocabulary: NamingVocabulary, policy: InventoryPolicy, @@ -710,7 +688,7 @@ def _name_and_suppress( ) for ts in all_ts: try: - image = store.streams.color_image.at(ts, 0.05).first().data + image = color.at(ts, 0.05).first().data except LookupError: continue @@ -724,7 +702,7 @@ def _name_and_suppress( continue bbox = (float(box[0]), float(box[1]), float(box[2]), float(box[3])) best_target: Any = None - best_iou = NAME_ATTACH_IOU + best_iou = policy.name_attach_iou for track, member in frame_members.get(ts, []): if member.bbox is None: continue @@ -774,9 +752,7 @@ def inventory( policy: InventoryPolicy | None = None, motion_threshold: float = MOTION_THRESHOLD, log_progress: bool = False, - world_frame: str = WORLD_FRAME, - optical_frame: str = OPTICAL_FRAME, - tf_tolerance: float = TF_TOLERANCE, + rig: Rig | None = None, ) -> list[Instance]: """Deduplicated object instances for the window, computed at query time. @@ -796,47 +772,28 @@ def inventory( Both models belong to the caller: nothing here is loaded or stopped, so one process can call this repeatedly on warm weights, over as many - windows as it wants. + windows as it wants. Without a ``rig`` the store's shape decides one, + and without a ``policy`` the rig supplies scale-appropriate defaults. """ - policy = policy or InventoryPolicy() - tf = StreamTF.from_store(store) - if tf is None: - raise ValueError("recording has no tf stream") - camera_info = store.streams.camera_info.first().data - lo, hi = store.streams.color_image.get_time_range() + rig = rig or Rig.from_store(store) + policy = policy or rig.default_inventory_policy() + lo, hi = rig.color.get_time_range() t0 = after if after is not None else lo t1 = before if before is not None else hi logger.info(f"inventory window: {t0 - lo:.1f}s to {t1 - lo:.1f}s ({t1 - t0:.1f}s)") - intervals = gates.still_intervals(tf, t0, t1, optical_frame, world_frame, tf_tolerance) - gray: dict[float, Any] = {} - - keyframes = gates.keyframes( - store, - tf, - t0, - t1, - KEYFRAME_STRIDE, - intervals, - gray, - motion_threshold, - optical_frame, - world_frame, - tf_tolerance, - ) - logger.info(f"gates: {len(keyframes)} keyframes pass camera-still + scene-still") + keyframes = rig.keyframes(t0, t1, policy.keyframe_stride_s, motion_threshold) + logger.info(f"gates: {len(keyframes)} keyframes pass the capture gates") if not keyframes: return [] - plane = fit_support_plane( - store, tf, camera_info, keyframes, optical_frame, world_frame, tf_tolerance - ) + plane = fit_support_plane(rig, keyframes) if plane is not None: logger.info(f"support plane: {plane.inlier_count} inliers") frames_grounded: list[tuple[float, list[SupportObservation]]] = [] frames_ungrounded: list[tuple[float, list[Detection2DSeg]]] = [] - image_area = float(camera_info.width * camera_info.height) + image_area = float(rig.camera_info.width * rig.camera_info.height) from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -854,7 +811,7 @@ def inventory( ) continue - pose = gates.camera_pose(tf, obs.ts, optical_frame, world_frame, tf_tolerance) + pose = rig.camera_pose(obs.ts) if pose is None: if log_progress: logger.info( @@ -868,15 +825,10 @@ def inventory( det.track_id = j grounded, ungrounded = _lift_frame( ImageDetections2D(obs.data, accepted), - store, - tf, - camera_info, + rig, obs.ts, camera_position, policy, - optical_frame, - world_frame, - tf_tolerance, plane, ) grounded = [o for o in grounded if _in_scope(o, plane, policy)] @@ -900,14 +852,14 @@ def inventory( tracks_2d = _track_ungrounded(frames_ungrounded) if include_ungrounded else [] logger.info(f"association: {len(tracks)} grounded instances") - _name_and_suppress(tracks, tracks_2d, store, detector, naming_vocabulary, policy) + _name_and_suppress(tracks, tracks_2d, rig.color, detector, naming_vocabulary, policy) tracks = [t for t in tracks if len(t.members) >= policy.min_member_observations] tracks.sort(key=lambda t: min(m.ts for m in t.members)) instances: list[Instance] = [] unknown = 0 for index, track in enumerate(tracks): - instance = _build_instance(index, track, policy) + instance = _build_instance(index, track, policy, rig.world_frame) if instance.primary_label is None: instance.primary_label = f"unknown-{unknown}" unknown += 1 diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index e6182fd8d6..d8cd3e321d 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -16,7 +16,8 @@ Search memory with embeddings (SigLIP, frame-level), open-vocabulary detection (OWLv2, calibrated per-box scores), -segmentation (EdgeTAM), projection to 3D through aligned depth. Two +segmentation (EdgeTAM), projection to 3D through the rig's geometry - an +aligned depth stream or a registered pointcloud stream. Two algorithm rules distinguish it from a best-crop search: * **Latest-pose semantics.** Among verified observations of the chosen @@ -36,30 +37,21 @@ import numpy as np from dimos.memory.embed import EmbedImages -from dimos.memory.tf import StreamTF from dimos.memory.transform import throttle -from dimos.perception.detection.project import sees from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D -from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC -from dimos.perception.memory import gates -from dimos.perception.memory.gates import OPTICAL_FRAME, TF_TOLERANCE, WORLD_FRAME +from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Localization, LocalizePolicy, Support from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo - from dimos.memory.stream import Stream from dimos.models.embedding.siglip import SigLIPModel from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC - from dimos.protocol.tf.tf import TFLookup logger = setup_logger() -EMBED_HZ = 1.0 -TOP_FRAMES = 12 TIME_BANDS = 6 # stratify retrieval across the window so late scans always compete BOXES_PER_FRAME = 4 CONFIRM_FLOOR = 0.22 # geometric confirmation accept for re-detections @@ -191,26 +183,19 @@ def detect(self, image: Any, query: str) -> ImageDetections2D: def _lift( detections: ImageDetections2D, - store: Any, - tf: TFLookup, - camera_info: CameraInfo, - optical_frame: str, - world_frame: str, - tf_tolerance: float, + rig: Rig, policy: LocalizePolicy, plane: Any | None = None, ) -> list[tuple[Detection3DPC, np.ndarray]]: - """Depth-lift 2D detections; returns valid (detection3d, camera_position) pairs.""" - depth = gates.depth_at(store, detections.ts) - transform = tf.get(optical_frame, world_frame, detections.ts, tf_tolerance) - if depth is None or transform is None: - return [] - pose = gates.camera_pose(tf, detections.ts, optical_frame, world_frame, tf_tolerance) + """Lift 2D detections through the rig; returns valid (detection3d, camera_position) pairs.""" + pose = rig.camera_pose(detections.ts) if pose is None: return [] camera = np.array([pose.position.x, pose.position.y, pose.position.z]) - lifted = ImageDetections3DPC.from_depth(detections, depth, camera_info, transform) + lifted = rig.lift(detections) + if lifted is None: + return [] valid: list[tuple[Detection3DPC, np.ndarray]] = [] for det3d in lifted: points = np.asarray(det3d.pointcloud.pointcloud.points) @@ -238,28 +223,19 @@ def embed_index( t0: float, t1: float, *, - optical_frame: str = OPTICAL_FRAME, - world_frame: str = WORLD_FRAME, - tf_tolerance: float = TF_TOLERANCE, + rig: Rig | None = None, ) -> Stream[Any, Any]: - """SigLIP-embedded, world-posed frame index at EMBED_HZ over the window. + """SigLIP-embedded, world-posed frame index at the rig's embed rate. Built once per window and handed to every ``localize`` call on it: the embed forwards are what a second query would otherwise repeat. """ - tf = StreamTF.from_store(store) - if tf is None: - raise ValueError("recording has no tf stream") + rig = rig or Rig.from_store(store) posed = ( - store.streams.color_image.after(t0) + rig.color.after(t0) .before(t1) - .transform(throttle(1.0 / EMBED_HZ)) - .map( - lambda obs: obs.derive( - data=obs.data, - pose=gates.camera_pose(tf, obs.ts, optical_frame, world_frame, tf_tolerance), - ) - ) + .transform(throttle(1.0 / rig.embed_hz)) + .map(lambda obs: obs.derive(data=obs.data, pose=rig.index_pose(obs))) .filter(lambda obs: obs.pose is not None) ) embedded: Stream[Any, Any] = posed.transform(EmbedImages(siglip)).materialize() @@ -267,14 +243,7 @@ def embed_index( return embedded -def _retrieve( - index: Stream[Any, Any], - tf: TFLookup, - query_embedding: Any, - optical_frame: str, - world_frame: str, - tf_tolerance: float, -) -> list[Any]: +def _retrieve(index: Stream[Any, Any], rig: Rig, query_embedding: Any, budget: int) -> list[Any]: """Top still frames by text similarity, stratified over time bands. Stratification is what keeps latest-pose semantics honest at the @@ -284,14 +253,14 @@ def _retrieve( ranked = [ obs for obs in index.search(query_embedding, k=max(index.count(), 1)) - if gates.camera_still(tf, obs.ts, optical_frame, world_frame, tf_tolerance) + if rig.camera_still(obs.ts) ] if not ranked: return [] t0, t1 = index.get_time_range() bands = max(1, min(TIME_BANDS, int((t1 - t0) / 20))) - per_band = max(1, TOP_FRAMES // bands) + per_band = max(1, budget // bands) span = (t1 - t0) / bands selected: list[Any] = [] chosen: set[float] = set() @@ -304,7 +273,7 @@ def _retrieve( chosen.add(obs.ts) selected.append(obs) for obs in ranked: # fill remaining budget by global rank - if len(selected) >= TOP_FRAMES: + if len(selected) >= budget: break if obs.ts not in chosen: chosen.add(obs.ts) @@ -320,12 +289,10 @@ def localize( siglip: SigLIPModel, detector: Owlv2Detector, segmenter: EdgeTAMImageSegmenter, + rig: Rig | None = None, require_pose: bool = True, policy: LocalizePolicy | None = None, cloud_mode: str = "latest_visible", - world_frame: str = WORLD_FRAME, - optical_frame: str = OPTICAL_FRAME, - tf_tolerance: float = TF_TOLERANCE, trace: LocalizeTrace | list[LocalizeTrace] | None = None, ) -> Localization | list[Localization | None] | None: """Latest unambiguous 3D localization of *query*, or ``None``. @@ -340,16 +307,15 @@ def localize( OWLv2 takes the whole list per frame - and returns one result per label, in input order; ``trace`` then takes a list of the same length. - The index and the three models belong to the caller: nothing here is - loaded or stopped, so one process can call this repeatedly on warm - weights, and every query on one window reuses the same embeddings. The - window is the index's - build it with :func:`embed_index`. + The index, the rig and the three models belong to the caller: nothing + here is loaded or stopped, so one process can call this repeatedly on + warm weights, and every query on one window reuses the same embeddings. + The window is the index's - build it with :func:`embed_index`. Without a + ``rig`` the store's shape decides one, and without a ``policy`` the rig + supplies scale-appropriate defaults. """ - policy = policy or LocalizePolicy() - tf = StreamTF.from_store(store) - if tf is None: - raise ValueError("recording has no tf stream") - camera_info = store.streams.camera_info.first().data + rig = rig or Rig.from_store(store) + policy = policy or rig.default_localize_policy() queries = [query] if isinstance(query, str) else query traces: list[LocalizeTrace | None] = ( @@ -358,19 +324,14 @@ def localize( cache = _DetectionCache(detector, segmenter, queries, policy.candidate_floor) results = [ _localize_one( - store, q, index=index, siglip=siglip, cache=cache, - tf=tf, - camera_info=camera_info, + rig=rig, require_pose=require_pose, policy=policy, cloud_mode=cloud_mode, - world_frame=world_frame, - optical_frame=optical_frame, - tf_tolerance=tf_tolerance, trace=t, ) for q, t in zip(queries, traces, strict=True) @@ -379,34 +340,27 @@ def localize( def _localize_one( - store: Any, query: str, *, index: Stream[Any, Any], siglip: SigLIPModel, cache: _DetectionCache, - tf: TFLookup, - camera_info: CameraInfo, + rig: Rig, require_pose: bool, policy: LocalizePolicy, cloud_mode: str, - world_frame: str, - optical_frame: str, - tf_tolerance: float, trace: LocalizeTrace | None, ) -> Localization | None: # Pass 1 - SigLIP: rank the indexed frames by the query. query_embedding = siglip.embed_text(query) - frames = _retrieve(index, tf, query_embedding, optical_frame, world_frame, tf_tolerance) + frames = _retrieve(index, rig, query_embedding, policy.retrieval_frames) logger.info(f"localize '{query}': {len(frames)} candidate frames of {index.count()} embedded") if not frames: return None from dimos.perception.memory.support_plane import fit_support_plane - plane = fit_support_plane( - store, tf, camera_info, frames, optical_frame, world_frame, tf_tolerance - ) + plane = fit_support_plane(rig, frames) # Pass 2 - OWLv2 + EdgeTAM: detect, segment, lift, verify. clusters: list[_Cluster] = [] @@ -423,17 +377,7 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: return if trace is not None and not is_verify: trace.detection_frames.append(frame_obs.derive(data=detections)) - lifted = _lift( - detections, - store, - tf, - camera_info, - optical_frame, - world_frame, - tf_tolerance, - policy, - plane, - ) + lifted = _lift(detections, rig, policy, plane) for det2d in detections: if not any(d.track_id == det2d.track_id for d, _ in lifted): best = (det2d.confidence, det2d.ts) @@ -467,27 +411,19 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: # geometrically (near + sees with occlusion), then re-detected. clusters.sort(key=lambda c: -c.max_score) for cluster in list(clusters[:4]): - predicate = sees( + predicate = rig.sees( cluster.center, - camera_info, - tf=tf, - world_frame=world_frame, - optical_frame=optical_frame, - time_tolerance=tf_tolerance, extent=np.minimum(cluster.extent, 0.4), # A large object overflows close-up frames; a third of its box in # view is still a usable re-detection pass, and those close-ups # are exactly the distinct viewpoints verification needs. min_fraction=0.35, - depth=lambda obs: gates.depth_at(store, obs.ts), - max_range=1.6, + max_range=policy.verify_radius_m, ) observing = [ obs - for obs in index.near(cluster.center, radius=1.6) - if obs.ts not in processed - and gates.camera_still(tf, obs.ts, optical_frame, world_frame, tf_tolerance) - and predicate(obs) + for obs in index.near(cluster.center, radius=policy.verify_radius_m) + if obs.ts not in processed and rig.camera_still(obs.ts) and predicate(obs) ] if len(observing) > VERIFY_FRAMES: # Even spread that always includes the endpoints: dropping the @@ -520,7 +456,7 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: ambiguity_margin=1.0, position_world_xyz=None, orientation_world_xyzw=None, - frame_id="world", + frame_id=rig.world_frame, support=None, pose_timestamp=ts, geometry_timestamp=ts, @@ -565,7 +501,7 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), coverage=_azimuth_coverage(winner.observations, winner.center), axes_observed=_axes_observed(winner.observations, winner.center), - frame_id="world", + frame_id=rig.world_frame, ) if trace is not None: @@ -583,7 +519,7 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: float(latest.centroid[2]), ), orientation_world_xyzw=orientation, - frame_id="world", + frame_id=rig.world_frame, support=support, pose_timestamp=latest.ts, geometry_timestamp=latest.ts, diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py new file mode 100644 index 0000000000..91d086d6e0 --- /dev/null +++ b/dimos/perception/memory/rig.py @@ -0,0 +1,430 @@ +# 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. + +"""The sensor rig behind a recording: intrinsics, poses, and 3D geometry. + +Two independent axes generalize the perception stack beyond one robot: + +* **Pose source.** The world-to-optical transform comes from a recorded + ``tf`` stream, or - for recordings without one - from the world pose + stamped on each observation plus a static base-to-optical mount. +* **Geometry source.** 3D geometry comes from an aligned ``depth`` stream, + unprojected per detection mask, or from a world-frame pointcloud stream + (a registered lidar), projected through the camera per detection mask. + Registered scans are sparse, so the cloud at a timestamp is the + concatenation of the scans in a short window around it - the scene is + static in world frame, which is what makes accumulation valid. + +``Rig.from_store`` recognizes both recording shapes; every field can also +be supplied directly for live stores whose streams are still filling. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from dimos.memory.tf import StreamTF +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.perception.detection.project import sees as project_sees +from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC +from dimos.perception.detection.type.detection3d.pointcloud_filters import ( + range_cluster, + statistical, +) +from dimos.perception.memory import gates +from dimos.perception.memory.gates import SPEED_MAX, STILL_ENVELOPE, TF_TOLERANCE +from dimos.perception.memory.types import InventoryPolicy, LocalizePolicy + +if TYPE_CHECKING: + from collections.abc import Callable + + from dimos_lcm.sensor_msgs import CameraInfo + + from dimos.memory.type.observation import Observation + from dimos.msgs.geometry_msgs.Pose import Pose + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.sensor_msgs.Image import Image + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D + from dimos.protocol.tf.tf import TFLookup + +# Pose-stamped rigs ride per-frame odometry, so walking does not stale the +# projection the way a sweeping wrist stales interpolated tf; the gate only +# drops speed glitches and sprints. +WALK_SPEED_MAX = 1.5 + +DEPTH_TOLERANCE = 0.06 # s - temporal join color->depth +CLOUD_ACCUM_S = 2.0 # s - scans within this of a frame form its geometry + +EMBED_HZ = 1.0 # index density for a wrist camera parked over a workspace +# A walking robot changes viewpoint every frame and its frames blur +# unevenly, so the index must sample denser for retrieval to catch the +# sharp sightings. +WALK_EMBED_HZ = 3.0 + +# Sparse projected clouds: split off background seen through the mask, then a +# loose outlier trim. The dense-cloud defaults (raycast + radius) assume a +# density registered lidar does not have. +_CLOUD_LIFT_FILTERS = [range_cluster(), statistical(nb_neighbors=12, std_ratio=2.0)] + +# Room-scale policies for mobile-robot rigs: objects are furniture-sized, +# viewpoints meters apart, odometry drifts centimeters between passes, and +# registered lidar is noisier and sparser than wrist-camera depth. +ROOM_LOCALIZE_POLICY = LocalizePolicy( + candidate_floor=0.18, + accept_score=0.32, + retrieval_frames=20, + cluster_radius_m=0.30, + min_depth_points=30, + max_object_extent_m=2.0, + min_camera_range_m=0.5, + surface_patch_max_rise_m=0.08, + surface_patch_min_drop_m=-0.06, + verify_radius_m=5.0, +) +ROOM_INVENTORY_POLICY = InventoryPolicy( + keyframe_stride_s=1.25, + min_mask_area_px=900, + min_depth_points=30, + max_object_extent_m=2.0, + min_height_above_plane_m=0.08, + band_above_plane_m=(-0.05, 1.5), + min_camera_range_m=0.5, + envelope_pad_m=0.12, + search_radius_m=0.8, + size_gap_max_m=0.8, + support_explained=0.25, + name_attach_iou=0.20, + same_frame_merge_gap_m=0.10, + split_extent_m=1.2, + split_height_m=0.9, + split_eps_m=0.10, +) + + +@dataclass +class Rig: + """Everything the stack needs to go from a 2D mask to world geometry. + + Exactly one of ``tf`` / (``base_to_optical`` + ``poses``) provides the + world-to-optical transform, and exactly one of ``depth`` / ``cloud`` + provides 3D geometry. + """ + + camera_info: CameraInfo + color: Any # color_image stream + world_frame: str + optical_frame: str + tf: TFLookup | None = None + base_to_optical: Transform | None = None + poses: Any = None # stream carrying world base poses (e.g. odom) + depth: Any = None # aligned depth stream, lifted via from_depth + cloud: Any = None # world-frame pointcloud stream, lifted via from_2d + tf_tolerance: float = TF_TOLERANCE + cloud_accum_s: float = CLOUD_ACCUM_S + speed_max: float = SPEED_MAX + scene_gate: bool = True + embed_hz: float = EMBED_HZ + _cloud_memo: tuple[float, PointCloud2] | None = field(default=None, repr=False, init=False) + + @classmethod + def from_store(cls, store: Any) -> Rig: + """Recognize the store's shape. + + A ``tf`` stream wins as pose source; without one the observations' + stamped poses are used with the Go2 front-camera mount. ``depth_image`` + wins as geometry source; a ``lidar`` stream is next, registered in + whatever world frame its scans carry; a store with neither can still + embed and retrieve, just never lift. Intrinsics and the optical frame + name come from the ``camera_info`` stream, or - for stores that carry + none - the static Go2 front-camera calibration. Stream contents are + only read where a name must be sniffed, so a live store whose streams + are still empty resolves too. + """ + streams = store.list_streams() + color = store.streams.color_image + tf = StreamTF.from_store(store) + + depth = store.streams.depth_image if "depth_image" in streams else None + cloud = None + world_frame = gates.WORLD_FRAME + if depth is None and "lidar" in streams: + cloud = store.streams.lidar + try: + world_frame = cloud.first().data.frame_id + except LookupError: + pass # live store, nothing recorded yet + + if "camera_info" in streams: + camera_info = store.streams.camera_info.first().data + optical_frame = camera_info.frame_id + else: + from dimos.robot.unitree.go2.connection import GO2Connection + + camera_info = GO2Connection.camera_info_static + # With a tf tree the optical name must match that tree; the Go2 + # calibration names apply only to the tf-less Go2 shape. + optical_frame = camera_info.frame_id if tf is None else gates.OPTICAL_FRAME + + base_to_optical = None + poses = None + if tf is None: + from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL + + base_to_optical = BASE_TO_OPTICAL + poses = store.streams.odom + + mobile = cloud is not None + return cls( + camera_info=camera_info, + color=color, + world_frame=world_frame, + optical_frame=optical_frame, + tf=tf, + base_to_optical=base_to_optical, + poses=poses, + depth=depth, + cloud=cloud, + speed_max=WALK_SPEED_MAX if mobile else SPEED_MAX, + scene_gate=not mobile, + embed_hz=WALK_EMBED_HZ if mobile else EMBED_HZ, + ) + + # pose + + def world_to_optical(self, ts: float) -> Transform | None: + if self.tf is not None: + return self.tf.get(self.optical_frame, self.world_frame, ts, self.tf_tolerance) + pose = self.pose_at(ts) + if pose is None: + return None + mount = cast("Transform", self.base_to_optical) + return -(Transform.from_pose("base_link", pose) + mount) + + def pose_at(self, ts: float) -> PoseStamped | None: + """World base pose nearest ts, from the poses stream. + + ``at().first()`` is window-earliest, not window-nearest; a walking + robot covers centimeters per pose period, so the nearest pose in the + window is what keeps the projection aligned. + """ + candidates = list(self.poses.at(ts, self.tf_tolerance)) + if not candidates: + return None + obs = min(candidates, key=lambda o: abs(o.ts - ts)) + pose: PoseStamped | None = obs.pose_stamped + return pose + + def camera_pose(self, ts: float) -> PoseStamped | None: + """World pose of the camera optical frame at ts.""" + transform = self.world_to_optical(ts) + return (-transform).to_pose() if transform is not None else None + + def index_pose(self, obs: Observation[Any]) -> Pose | PoseStamped | None: + """The pose an embedded index observation carries. + + tf rigs stamp the derived optical pose; pose-stamped rigs keep the + recorded base pose, which is what ``sees`` expects to find on an + observation when it resolves the transform through the mount. + """ + if self.tf is not None: + return self.camera_pose(obs.ts) + return obs.pose + + def camera_speed(self, ts: float, dt: float = 0.06) -> float | None: + """Linear camera speed (m/s) around ts, from pose differencing.""" + a = self.camera_pose(ts - dt) + b = self.camera_pose(ts + dt) + if a is None or b is None: + return None + return float((b.position - a.position).magnitude() / (2 * dt)) + + def camera_still(self, ts: float, envelope: float = STILL_ENVELOPE) -> bool: + """Camera below the rig's speed gate over the whole capture envelope.""" + for offset in (-envelope, 0.0, envelope): + speed = self.camera_speed(ts + offset) + if speed is None or speed > self.speed_max: + return False + return True + + def still_intervals(self, t0: float, t1: float) -> list[tuple[float, float]]: + """Maximal camera-still intervals inside [t0, t1], sampled at 0.25 s.""" + step = 0.25 + times = np.arange(t0, t1 + step, step) + intervals: list[tuple[float, float]] = [] + run_start: float | None = None + for t in times: + speed = self.camera_speed(float(t)) + still = speed is not None and speed <= self.speed_max + if still and run_start is None: + run_start = float(t) + elif not still and run_start is not None: + intervals.append((run_start, float(t) - step)) + run_start = None + if run_start is not None: + intervals.append((run_start, float(times[-1]))) + return [(a, b) for a, b in intervals if b >= a] + + # geometry + + def depth_at(self, ts: float) -> Image | None: + """Temporal join: aligned depth frame for a color timestamp.""" + try: + depth: Image = self.depth.at(ts, DEPTH_TOLERANCE).first().data + except LookupError: + return None + return depth + + def cloud_at(self, ts: float) -> PointCloud2 | None: + """World-frame geometry at ts: the scans accumulated around it.""" + if self._cloud_memo is not None and self._cloud_memo[0] == ts: + return self._cloud_memo[1] + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + scans = list(self.cloud.after(ts - self.cloud_accum_s).before(ts + self.cloud_accum_s)) + if not scans: + return None + merged: PointCloud2 + if len(scans) == 1: + merged = scans[0].data + else: + # A grid-quantized source (the Go2 occupancy stream) repeats the + # same cell in every snapshot it persists through; accumulation + # must not count one voxel once per snapshot. + points = np.unique(np.vstack([scan.data.as_numpy()[0] for scan in scans]), axis=0) + merged = PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=ts) + self._cloud_memo = (ts, merged) + return merged + + def lift(self, detections: ImageDetections2D) -> ImageDetections3DPC | None: + """2D detections to world-frame 3D clouds, or None without geometry/pose.""" + transform = self.world_to_optical(detections.ts) + if transform is None: + return None + if self.depth is not None: + depth = self.depth_at(detections.ts) + if depth is None: + return None + return ImageDetections3DPC.from_depth(detections, depth, self.camera_info, transform) + cloud = self.cloud_at(detections.ts) + if cloud is None: + return None + return ImageDetections3DPC.from_2d( + detections, cloud, self.camera_info, transform, _CLOUD_LIFT_FILTERS + ) + + def backdrop(self, ts: float) -> PointCloud2 | None: + """World-frame scene cloud around ts, for plane fits and rendering.""" + if self.depth is None: + return self.cloud_at(ts) + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + depth = self.depth_at(ts) + transform = self.world_to_optical(ts) + if depth is None or transform is None: + return None + try: + color = self.color.at(ts, 0.1).first().data + except LookupError: + return None + return PointCloud2.from_rgbd( + color, depth, self.camera_info, depth_scale=0.001, depth_trunc=1.5 + ).transform(-transform) + + # predicates + + def sees( + self, + point: Any, + *, + extent: Any | None = None, + min_fraction: float = 1.0, + max_range: float | None = None, + ) -> Callable[[Observation[Any]], bool]: + """Predicate: does an observation's camera see the world point. + + Occlusion checking through measured depth only exists on depth rigs; + projected-cloud rigs rely on the geometric visibility test alone. + """ + return project_sees( + point, + self.camera_info, + tf=self.tf, + base_to_optical=self.base_to_optical, + world_frame=self.world_frame, + optical_frame=self.optical_frame, + time_tolerance=self.tf_tolerance, + extent=extent, + min_fraction=min_fraction, + max_range=max_range, + depth=(lambda obs: self.depth_at(obs.ts)) if self.depth is not None else None, + ) + + def keyframes( + self, + t0: float, + t1: float, + stride: float, + motion_threshold: float = gates.MOTION_THRESHOLD, + ) -> list[Observation[Image]]: + """Camera-still (and, on scene-gated rigs, scene-still) frames on a grid. + + For each grid point the nearest passing frame within half a stride is + selected, so a grid point landing mid-sweep snaps to the neighboring + pause instead of being lost. + """ + intervals = self.still_intervals(t0, t1) if self.scene_gate else [] + gray: dict[float, np.ndarray | None] = {} + selected: list[Observation[Image]] = [] + seen: set[float] = set() + offsets = [0.0] + probe = 0.35 + while probe <= stride / 2: + offsets.extend([probe, -probe]) + probe += 0.35 + + t = t0 + 0.5 + while t < t1: + for offset in offsets: + ts = t + offset + if ts < t0 or ts > t1: + continue + if not self.camera_still(ts): + continue + if self.scene_gate and not gates.scene_still( + self.color, ts, intervals, gray, motion_threshold + ): + continue + try: + obs = self.color.at(ts, 0.1).first() + except LookupError: + continue + if obs.ts in seen: + break + if self.world_to_optical(obs.ts) is None: + continue + seen.add(obs.ts) + selected.append(obs) + break + t += stride + return selected + + def default_localize_policy(self) -> LocalizePolicy: + return LocalizePolicy() if self.depth is not None else ROOM_LOCALIZE_POLICY + + def default_inventory_policy(self) -> InventoryPolicy: + return InventoryPolicy() if self.depth is not None else ROOM_INVENTORY_POLICY diff --git a/dimos/perception/memory/support_plane.py b/dimos/perception/memory/support_plane.py index f3382f74e9..52a4b0f921 100644 --- a/dimos/perception/memory/support_plane.py +++ b/dimos/perception/memory/support_plane.py @@ -15,33 +15,31 @@ """Support-surface fit and the scope predicate derived from it. The plane is RANSAC-fit from the window's own frames - nothing scene-specific -is passed in and no caller supplies coordinates. The plane's inlier footprint -is the workspace; the scope predicate accepts a support when its cloud sits in -a band above the plane and its footprint intersects the plane footprint. +is passed in and no caller supplies coordinates. On a wrist-camera rig over a +workspace the dominant horizontal plane is the tabletop; on a mobile rig it +is the floor. The plane's inlier footprint is the workspace; the scope +predicate accepts a support when its cloud sits in a band above the plane and +its footprint intersects the plane footprint. """ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import numpy as np -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.perception.memory import gates from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo - from dimos.memory.type.observation import Observation from dimos.msgs.sensor_msgs.Image import Image - from dimos.protocol.tf.tf import TFLookup + from dimos.perception.memory.rig import Rig logger = setup_logger() -BACKDROP_DEPTH_TRUNC = 1.5 # m - the workspace a wrist camera actually covers -PLANE_DISTANCE = 0.01 # m - RANSAC inlier distance +PLANE_DISTANCE = 0.01 # m - RANSAC inlier distance for depth-camera clouds +PLANE_DISTANCE_CLOUD = 0.03 # m - registered lidar is noisier MIN_HORIZONTAL_DOT = 0.90 # |normal . z| for a plane to count as horizontal PLANE_SEED = 0 # RANSAC is seeded so one window always fits the same plane FOOTPRINT_DILATE_M = 0.03 @@ -78,21 +76,12 @@ def footprint_contains(self, points_xy: np.ndarray) -> np.ndarray: return MplPath(dilated).contains_points(points_xy) -def fit_support_plane( - store: Any, - tf: TFLookup, - camera_info: CameraInfo, - keyframes: list[Observation[Image]], - optical_frame: str, - world_frame: str, - tf_tolerance: float, -) -> SupportPlane | None: +def fit_support_plane(rig: Rig, keyframes: list[Observation[Image]]) -> SupportPlane | None: """Fit the dominant near-horizontal plane from a handful of window keyframes. Non-horizontal dominant planes (a wall, a screen) are peeled off and the fit repeats on the remainder. Among horizontal candidates the one with - the most inliers wins - for a wrist camera over a workspace that is the - support surface itself. + the most inliers wins. """ if not keyframes: return None @@ -100,13 +89,9 @@ def fit_support_plane( picks = keyframes[:: max(1, len(keyframes) // 5)][:5] clouds = [] for obs in picks: - depth = gates.depth_at(store, obs.ts) - transform = tf.get(optical_frame, world_frame, obs.ts, tf_tolerance) - if depth is None or transform is None: + cloud = rig.backdrop(obs.ts) + if cloud is None: continue - cloud = PointCloud2.from_rgbd( - obs.data, depth, camera_info, depth_scale=0.001, depth_trunc=BACKDROP_DEPTH_TRUNC - ).transform(-transform) clouds.append(cloud.voxel_downsample(0.01)) if not clouds: return None @@ -123,6 +108,7 @@ def fit_support_plane( # Unseeded, the plane fit lands on a different set of inliers each run o3d.utility.random.seed(PLANE_SEED) + distance = PLANE_DISTANCE if rig.depth is not None else PLANE_DISTANCE_CLOUD remaining = o3d.geometry.PointCloud() remaining.points = o3d.utility.Vector3dVector(points) best: tuple[np.ndarray, np.ndarray] | None = None # (coefficients, inlier points) @@ -130,7 +116,7 @@ def fit_support_plane( if len(remaining.points) < 500: break model, inlier_idx = remaining.segment_plane( - distance_threshold=PLANE_DISTANCE, ransac_n=3, num_iterations=1000, probability=1.0 + distance_threshold=distance, ransac_n=3, num_iterations=1000, probability=1.0 ) inliers = np.asarray(remaining.points)[inlier_idx] normal = np.array(model[:3]) diff --git a/dimos/perception/memory/tool_inventory.py b/dimos/perception/memory/tool_inventory.py index dc8c20afaf..5c0ababea2 100644 --- a/dimos/perception/memory/tool_inventory.py +++ b/dimos/perception/memory/tool_inventory.py @@ -52,13 +52,12 @@ import argparse from pathlib import Path import sys -from typing import Any, cast +from typing import cast from dimos.memory.store.sqlite import SqliteStore -from dimos.memory.tf import StreamTF from dimos.memory.transform import throttle -from dimos.perception.memory import gates from dimos.perception.memory.inventory import DEFAULT_VOCABULARY, NamingVocabulary +from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Instance, SupportObservation from dimos.utils.data import get_data @@ -85,7 +84,7 @@ def instance_label(instance: Instance) -> str: return f"{instance.instance_id} {instance.primary_label}" -def render(out: str, store: Any, instances: list[Instance], t0: float, t1: float) -> None: +def render(out: str, rig: Rig, instances: list[Instance], t0: float, t1: float) -> None: """Write the .rrd - rerun stays an inline import. Entity contract: ``map`` backdrop, ``camera/image`` the live feed carrying @@ -96,13 +95,8 @@ def render(out: str, store: Any, instances: list[Instance], t0: float, t1: float import rerun.blueprint as rrb from dimos.memory.vis.color import Color - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.visualization.rerun.init import rerun_init - tf = StreamTF.from_store(store) - assert tf is not None - camera_info = store.streams.camera_info.first().data - rerun_init("memory-inventory") rr.save(out) rr.send_blueprint( @@ -115,7 +109,7 @@ def render(out: str, store: Any, instances: list[Instance], t0: float, t1: float ) ) - POINT_SIZE = 0.005 + point_size = 0.005 if rig.depth is not None else 0.015 def at(ts: float) -> None: rr.set_time("ts", timestamp=ts) @@ -136,24 +130,37 @@ def at(ts: float) -> None: for member in instance.members: frames.setdefault(member.ts, []).append((i, member)) - # scene backdrop from the last keyframe that carried an instance - backdrop_ts = max(frames, default=None) - if backdrop_ts is not None: - color = store.streams.color_image.at(backdrop_ts, 0.1).first().data - depth = gates.depth_at(store, backdrop_ts) - transform = tf.get(gates.OPTICAL_FRAME, gates.WORLD_FRAME, backdrop_ts, gates.TF_TOLERANCE) - assert depth is not None and transform is not None - backdrop = PointCloud2.from_rgbd(color, depth, camera_info, depth_scale=0.001).transform( - -transform - ) - rr.log("map", backdrop.voxel_downsample(0.01).to_rerun(voxel_size=POINT_SIZE), static=True) + # scene backdrop: for depth rigs the last instance keyframe's RGBD cloud, + # for pointcloud rigs the window's scans merged into one map + if rig.depth is not None: + backdrop_ts = max(frames, default=None) + if backdrop_ts is not None: + backdrop = rig.backdrop(backdrop_ts) + if backdrop is not None: + rr.log( + "map", + backdrop.voxel_downsample(0.01).to_rerun(voxel_size=point_size), + static=True, + ) + else: + import numpy as np + + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + scans = [ + obs.data.as_numpy()[0] + for obs in rig.cloud.after(t0).before(t1).transform(throttle(2.0)) + ] + if scans: + merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) + rr.log("map", merged.voxel_downsample(0.05).to_rerun(voxel_size=0.01), static=True) # live camera feed + frustum; the empty box clears the overlay off non-keyframes - rr.log("camera", camera_info.to_rerun(), static=True) + rr.log("camera", rig.camera_info.to_rerun(), static=True) feed_throttle = 0.1 if (t1 - t0) <= 160 else 0.4 - feed = store.streams.color_image.after(t0).before(t1).transform(throttle(feed_throttle)) + feed = rig.color.after(t0).before(t1).transform(throttle(feed_throttle)) for obs in feed: - pose = gates.camera_pose(tf, obs.ts) + pose = rig.camera_pose(obs.ts) if pose is None: continue at(obs.ts) @@ -164,9 +171,9 @@ def at(ts: float) -> None: # keyframes, logged after the feed so their boxes win the shared timestamps for ts, entries in sorted(frames.items()): at(ts) - keyframe_pose = gates.camera_pose(tf, ts) + keyframe_pose = rig.camera_pose(ts) assert keyframe_pose is not None - rr.log("camera/image", store.streams.color_image.at(ts, 0.05).first().data.to_rerun()) + rr.log("camera/image", rig.color.at(ts, 0.05).first().data.to_rerun()) rr.log("camera", keyframe_pose.to_rerun()) rr.log( "camera/image/instances", @@ -180,7 +187,7 @@ def at(ts: float) -> None: ), ) for i, member in entries: - rr.log(paths[i], member.cloud.to_rerun(voxel_size=POINT_SIZE, colors=colors[i])) + rr.log(paths[i], member.cloud.to_rerun(voxel_size=point_size, colors=colors[i])) # the reported instance: one labeled box, static so it holds over the whole timeline for i, instance in enumerate(grounded): @@ -250,7 +257,8 @@ def main() -> int: "xarm6_worldbelief_20260729_203624_161992.db" ) store = SqliteStore(path=dataset) - lo, hi = store.streams.color_image.get_time_range() + rig = Rig.from_store(store) + lo, hi = rig.color.get_time_range() after = lo + args.start before = lo + args.start + args.duration if args.duration is not None else None @@ -271,6 +279,7 @@ def main() -> int: before=before, include_ungrounded=args.include_ungrounded, log_progress=args.log_progress, + rig=rig, ) print(f"instances: {len(instances)}") @@ -297,7 +306,7 @@ def main() -> int: ) if out is not None: - render(out, store, instances, after, before if before is not None else hi) + render(out, rig, instances, after, before if before is not None else hi) print(f"saved {out}") return 0 diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index b1ec9e9190..8224e31101 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -12,13 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Query memory for objects, localize them in 3D via depth, render in rerun. +"""Query memory for objects, localize them in 3D, render in rerun. Run: uv run python -m dimos.perception.memory.tool_localize [query ...] [out.rrd] - [--from ] [--duration ] [--multi] + [--dataset ] [--from ] [--duration ] [--multi] -Queries share one model load and one .rrd; with --multi they go to -localize() as one list, sharing a single detection pass per frame. +The recording's shape decides the rig: an xArm-style store lifts through +aligned depth and tf, a mobile-robot store (Go2/G1 replay) lifts through +registered lidar and stamped poses. Queries share one model load and one +.rrd; with --multi they go to localize() as one list, sharing a single +detection pass per frame. Exit code 0 with a printed position per verified hit; exit code 1 when no query is verified, with "no verified detection of ..." per miss - the honest answer that the object @@ -29,13 +32,11 @@ import argparse from pathlib import Path import sys -from typing import Any, cast from dimos.memory.store.sqlite import SqliteStore -from dimos.memory.tf import StreamTF from dimos.memory.transform import throttle -from dimos.perception.memory import gates from dimos.perception.memory.localize import LocalizeTrace +from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Localization from dimos.utils.data import get_data @@ -43,7 +44,7 @@ def render( - out: str, store: Any, traces: list[tuple[str, LocalizeTrace]], t0: float, t1: float + out: str, rig: Rig, traces: list[tuple[str, LocalizeTrace]], t0: float, t1: float ) -> None: """Write the .rrd - rerun stays an inline import. @@ -55,12 +56,8 @@ def render( import rerun as rr import rerun.blueprint as rrb - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.visualization.rerun.init import rerun_init - tf = cast("StreamTF", StreamTF.from_store(store)) - camera_info = store.streams.camera_info.first().data - rerun_init("memory-localize") rr.save(out) rr.send_blueprint( @@ -74,40 +71,48 @@ def render( ) GREEN, RED, BLUE = [46, 204, 113], [231, 76, 60], [52, 120, 246] - POINT_SIZE = 0.005 + point_size = 0.005 if rig.depth is not None else 0.015 def at(ts: float) -> None: rr.set_time("ts", timestamp=ts) - # scene backdrop from an answer frame's depth (or the first detection) - backdrop_ts = next((t.backdrop_ts for _, t in traces if t.backdrop_ts is not None), None) - if backdrop_ts is None: - backdrop_ts = next((t.matched[0][0] for _, t in traces if t.matched), None) - if backdrop_ts is not None: - try: - color = store.streams.color_image.at(backdrop_ts, 0.1).first().data - depth = gates.depth_at(store, backdrop_ts) - transform = tf.get( - gates.OPTICAL_FRAME, gates.WORLD_FRAME, backdrop_ts, gates.TF_TOLERANCE - ) - if depth is not None and transform is not None: - backdrop = PointCloud2.from_rgbd( - color, depth, camera_info, depth_scale=0.001 - ).transform(-transform) + # scene backdrop: for depth rigs the answer frame's RGBD cloud, for + # pointcloud rigs the window's scans merged into one map + if rig.depth is not None: + backdrop_ts = next((t.backdrop_ts for _, t in traces if t.backdrop_ts is not None), None) + if backdrop_ts is None: + backdrop_ts = next((t.matched[0][0] for _, t in traces if t.matched), None) + if backdrop_ts is not None: + backdrop = rig.backdrop(backdrop_ts) + if backdrop is not None: rr.log( "map", - backdrop.voxel_downsample(0.01).to_rerun(voxel_size=POINT_SIZE), + backdrop.voxel_downsample(0.01).to_rerun(voxel_size=point_size), static=True, ) - except LookupError: - pass + else: + import numpy as np + + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + scans = [ + obs.data.as_numpy()[0] + for obs in rig.cloud.after(t0).before(t1).transform(throttle(2.0)) + ] + if scans: + merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) + rr.log( + "map", + merged.voxel_downsample(0.05).to_rerun(voxel_size=0.01), + static=True, + ) - # live camera feed + frustum tracking the wrist along the timeline - rr.log("camera", camera_info.to_rerun(), static=True) + # live camera feed + frustum tracking the capture pose along the timeline + rr.log("camera", rig.camera_info.to_rerun(), static=True) feed_throttle = 0.1 if (t1 - t0) <= 160 else 0.4 - feed = store.streams.color_image.after(t0).before(t1).transform(throttle(feed_throttle)) + feed = rig.color.after(t0).before(t1).transform(throttle(feed_throttle)) for obs in feed: - pose = gates.camera_pose(tf, obs.ts) + pose = rig.camera_pose(obs.ts) if pose is None: continue at(obs.ts) @@ -119,7 +124,7 @@ def at(ts: float) -> None: # marked frames: into the live feed, plus a frozen frustum at the capture pose for i, obs in enumerate(trace.detection_frames): - pose = gates.camera_pose(tf, obs.ts) + pose = rig.camera_pose(obs.ts) if pose is None: continue at(obs.ts) @@ -127,7 +132,7 @@ def at(ts: float) -> None: rr.log("camera/image", annotated.to_rerun()) frame = f"{root}/frames/{i}" rr.log(frame, pose.to_rerun()) - rr.log(frame, camera_info.to_rerun()) + rr.log(frame, rig.camera_info.to_rerun()) rr.log(f"{frame}/image", annotated.to_rerun()) # 3d detections: green = matched candidates, red = cross-view re-detections @@ -139,7 +144,7 @@ def at(ts: float) -> None: at(ts) rr.log( f"{root}/{tag}/{i}_{det.name.replace(' ', '_')}", - det.pointcloud.to_rerun(voxel_size=POINT_SIZE, colors=rgb), + det.pointcloud.to_rerun(voxel_size=point_size, colors=rgb), ) # the answer: always blue, whatever the query @@ -147,7 +152,7 @@ def at(ts: float) -> None: at(trace.answer.ts) rr.log( f"{root}/answer", - trace.answer.pointcloud.to_rerun(voxel_size=POINT_SIZE, colors=BLUE), + trace.answer.pointcloud.to_rerun(voxel_size=point_size, colors=BLUE), ) @@ -213,7 +218,8 @@ def main() -> int: "xarm6_worldbelief_20260729_203624_161992.db" ) store = SqliteStore(path=dataset) - lo, hi = store.streams.color_image.get_time_range() + rig = Rig.from_store(store) + lo, hi = rig.color.get_time_range() after = lo + args.start before = lo + args.start + args.duration if args.duration is not None else hi @@ -222,13 +228,14 @@ def main() -> int: traces: list[tuple[str, LocalizeTrace]] = [] hits = 0 with DanDetector() as dan: - index = dan.embed(store, after, before) + index = dan.embed(store, after, before, rig=rig) if args.multi: qtraces = [LocalizeTrace() for _ in queries] results = dan.localize( store, queries, index=index, + rig=rig, require_pose=not args.allow_no_pose, trace=qtraces, ) @@ -243,6 +250,7 @@ def main() -> int: store, query, index=index, + rig=rig, require_pose=not args.allow_no_pose, trace=trace, ) @@ -254,7 +262,7 @@ def main() -> int: return 1 if out is not None: - render(out, store, traces, after, before) + render(out, rig, traces, after, before) print(f"saved {out}") return 0 diff --git a/dimos/perception/memory/types.py b/dimos/perception/memory/types.py index 328590d71d..552a4199fe 100644 --- a/dimos/perception/memory/types.py +++ b/dimos/perception/memory/types.py @@ -115,6 +115,10 @@ class LocalizePolicy: accept_score: float = 0.40 refusal_margin: float = 0.15 min_views: int = 2 # a support seen from one pose only is unconfirmed + # Frames retrieved per query for the detection pass. Wrist-camera frames + # each cover most of the workspace; room-scale sweeps dilute the + # embedding signal across viewpoints and need a larger budget. + retrieval_frames: int = 12 cluster_radius_m: float = 0.08 # observations within this are the same support min_depth_points: int = 60 @@ -124,6 +128,8 @@ class LocalizePolicy: # object: every real object rises above the plane, a surface patch does not. surface_patch_max_rise_m: float = 0.003 surface_patch_min_drop_m: float = -0.02 + # Cross-view verification looks for re-detections this far from a support. + verify_radius_m: float = 1.6 @dataclass(frozen=True) @@ -137,6 +143,7 @@ class InventoryPolicy: argument, not as policy. """ + keyframe_stride_s: float = 2.5 # proposal keyframe grid min_mask_area_px: int = 400 max_mask_area_fraction: float = 0.25 min_depth_points: int = 60 @@ -148,6 +155,20 @@ class InventoryPolicy: envelope_pad_m: float = 0.015 search_radius_m: float = 0.15 overlap_accept: float = 0.20 + # Same-object views may differ in bounding size by partiality alone; a + # gap beyond this is two different bodies. + size_gap_max_m: float = 0.25 + # The majority of a candidate's points must lie within the error envelope + # of the track's accumulated support. Partial and newly revealed views of + # one object satisfy this; a different object placed at a vacated rest + # position does not, which is what AABB overlap cannot express. + support_explained: float = 0.5 + # A lifted cloud plainly spanning more than one object: wider than any + # single object at this rig's scale, or taller than one body. Repaired by + # stripping support-surface points and splitting by 3D connectivity. + split_extent_m: float = 0.30 + split_height_m: float = 0.10 + split_eps_m: float = 0.03 # Same-frame observations whose clouds touch within this gap are one # body - rigid objects cannot interpenetrate, and distinct objects on a # workspace sit apart by more than sensor noise. This is what fuses @@ -163,6 +184,10 @@ class InventoryPolicy: # frame and again over a track. Otherwise the instance stays unknown-N. name_accept_score: float = 0.18 name_refusal_margin: float = 0.06 + # An attachment must be the detector drawing a box around this member. + # Whole-object masks want a strict overlap; fragment masks of a large + # object overlap their object's box only partially. + name_attach_iou: float = 0.45 include_object_parts: bool = False include_surfaces: bool = False From 32e97632914f67bfe64bb894fb7a5fb2ee5124a9 Mon Sep 17 00:00:00 2001 From: bogwi Date: Fri, 21 Aug 2026 18:04:16 +0900 Subject: [PATCH 02/28] improve perception stack; test across different sourced --- dimos/memory/stream.py | 11 +- dimos/memory/utils/sqlite.py | 3 + dimos/perception/detection/detectors/owlv2.py | 79 ++-- dimos/perception/detection/identity.py | 101 +++++ .../type/detection3d/imageDetections3DPC.py | 16 +- .../detection/type/detection3d/pointcloud.py | 140 ++++-- dimos/perception/memory/dandetect.py | 4 +- dimos/perception/memory/localize.py | 256 ++++++----- dimos/perception/memory/rig.py | 428 +++++++++++++++--- dimos/perception/memory/tool_inventory.py | 22 +- dimos/perception/memory/tool_localize.py | 22 +- 11 files changed, 826 insertions(+), 256 deletions(-) create mode 100644 dimos/perception/detection/identity.py diff --git a/dimos/memory/stream.py b/dimos/memory/stream.py index c5b4178244..da3aa9dc3d 100644 --- a/dimos/memory/stream.py +++ b/dimos/memory/stream.py @@ -568,8 +568,15 @@ def drain(self) -> int: return n def drain_thread(self) -> DisposableBase: - """Drain this stream on the dimos thread pool; returns a disposable.""" - return self.subscribe(lambda _: None) + """Drain this stream on the dimos thread pool; returns a disposable. + + A drain has no consumer to surface an iteration error, so it logs - + a dead pipeline must not die silently. + """ + return self.subscribe( + lambda _: None, + on_error=lambda e: logger.error("drain_thread() pipeline died: %s", e, exc_info=e), + ) def observable(self) -> reactivex.Observable[O]: """Convert this stream to an RxPY Observable. diff --git a/dimos/memory/utils/sqlite.py b/dimos/memory/utils/sqlite.py index 02a48f22b7..e46187c5d4 100644 --- a/dimos/memory/utils/sqlite.py +++ b/dimos/memory/utils/sqlite.py @@ -27,6 +27,9 @@ def open_sqlite_connection(path: str | Path) -> sqlite3.Connection: conn = sqlite3.connect(path, check_same_thread=False) conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") + # WAL has one writer at a time; concurrent writer threads (a recorder + # feed plus a live embed save) can starve past the 5s default and die. + conn.execute("PRAGMA busy_timeout=30000") conn.enable_load_extension(True) sqlite_vec.load(conn) conn.enable_load_extension(False) diff --git a/dimos/perception/detection/detectors/owlv2.py b/dimos/perception/detection/detectors/owlv2.py index d88ee75b8d..3139424069 100644 --- a/dimos/perception/detection/detectors/owlv2.py +++ b/dimos/perception/detection/detectors/owlv2.py @@ -30,6 +30,8 @@ class Owlv2Config(HuggingFaceModelConfig): model_name: str = "google/owlv2-base-patch16-ensemble" + # float16 runs the forward under autocast at roughly half the latency; + # scores jitter by a few thousandths, so threshold-edge boxes may flip. dtype: torch.dtype = torch.float32 @@ -62,6 +64,13 @@ def _processor(self): # type: ignore[no-untyped-def] return Owlv2Processor.from_pretrained(self.config.model_name) + def _autocast(self) -> torch.autocast: + return torch.autocast( + device_type="cuda", + dtype=self.config.dtype, + enabled=self.config.dtype is not torch.float32 and "cuda" in str(self.config.device), + ) + def query_detections( self, image: Image, @@ -74,38 +83,54 @@ def query_detections( ``confidence`` is the calibrated per-box score. ``class_id`` indexes into ``queries``. """ - pil = PILImage.fromarray(image.to_rgb().data) - with torch.inference_mode(): - inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to( - self.config.device - ) + return self.query_detections_batch([image], queries, threshold)[0] + + def query_detections_batch( + self, + images: list[Image], + queries: list[str], + threshold: float = 0.1, + ) -> list[ImageDetections2D]: + """``query_detections`` over several images in one forward pass. + + Per-call preprocessing, text encoding and kernel launches amortize + across the batch, which is what makes many-frame sweeps affordable; + results are per-image, in input order. + """ + pils = [PILImage.fromarray(image.to_rgb().data) for image in images] + with torch.inference_mode(), self._autocast(): + inputs = self._processor( + text=[queries] * len(pils), images=pils, return_tensors="pt" + ).to(self.config.device) outputs = self._model(**inputs) results = self._processor.post_process_grounded_object_detection( outputs=outputs, - target_sizes=torch.tensor([(pil.height, pil.width)]), + target_sizes=torch.tensor([(pil.height, pil.width) for pil in pils]), threshold=threshold, - )[0] - - detections: list[Detection2DBBox] = [] - w, h = float(pil.width), float(pil.height) - for box, score, label in zip( - results["boxes"], results["scores"], results["labels"], strict=False - ): - x1, y1, x2, y2 = (float(v) for v in box) - bbox = (max(0.0, x1), max(0.0, y1), min(w, x2), min(h, y2)) - det = Detection2DBBox( - bbox=bbox, - track_id=-1, - class_id=int(label), - confidence=float(score), - name=queries[int(label)], - ts=image.ts, - image=image, ) - if det.is_valid(): - detections.append(det) - return ImageDetections2D(image=image, detections=detections) + batch: list[ImageDetections2D] = [] + for image, pil, result in zip(images, pils, results, strict=True): + detections: list[Detection2DBBox] = [] + w, h = float(pil.width), float(pil.height) + for box, score, label in zip( + result["boxes"], result["scores"], result["labels"], strict=False + ): + x1, y1, x2, y2 = (float(v) for v in box) + bbox = (max(0.0, x1), max(0.0, y1), min(w, x2), min(h, y2)) + det = Detection2DBBox( + bbox=bbox, + track_id=-1, + class_id=int(label), + confidence=float(score), + name=queries[int(label)], + ts=image.ts, + image=image, + ) + if det.is_valid(): + detections.append(det) + batch.append(ImageDetections2D(image=image, detections=detections)) + return batch def query_score_rows( self, @@ -123,7 +148,7 @@ def query_score_rows( their score rows. """ pil = PILImage.fromarray(image.to_rgb().data) - with torch.inference_mode(): + with torch.inference_mode(), self._autocast(): inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to( self.config.device ) diff --git a/dimos/perception/detection/identity.py b/dimos/perception/detection/identity.py new file mode 100644 index 0000000000..b36652e080 --- /dev/null +++ b/dimos/perception/detection/identity.py @@ -0,0 +1,101 @@ +# 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. + +"""Object identity over detection streams: many sightings in, one per object out. + +``Identity`` is the aggregation stage of a search pipeline:: + + detections_2d.transform(ProjectTo3D(cloud, ...)).transform(Identity()) + +As a stream transformer it is a batch search processor: it consumes the full +upstream of 3D detections, groups the sightings of one physical object, and +then emits one merged :class:`Detection3DPC` per object - the union cloud of +every viewpoint that saw it. What counts as "the same object" is a pluggable +``is_same(a, b)`` strategy; v0 is spatial only ("is it roughly at the same +spot"), so there is no permanence: an object that moved registers as a new +object at its new rest position. + +The grouping core (``add`` plus ``groups``) is usable directly for callers +that need the members of each identity rather than the merged stream output; +``localize`` forms its support candidates with it. +""" + +from __future__ import annotations + +import operator +from typing import TYPE_CHECKING, Any + +from dimos.memory.transform import Transformer +from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC +from dimos.perception.detection.type.imageDetections import ImageDetections + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from dimos.memory.type.observation import Observation + + +def spatial(radius: float = 0.1) -> Callable[[Detection3DPC, Detection3DPC], bool]: + """Same object when the cloud centers sit within *radius* meters.""" + + def is_same(a: Detection3DPC, b: Detection3DPC) -> bool: + return float((a.center - b.center).magnitude()) <= radius + + return is_same + + +class Identity(Transformer[Any, Detection3DPC]): + """One detection3D per object, aggregated from every sighting. + + Each incoming detection is matched against the running merged + representative of every known object with ``is_same``; a match joins + that object and folds into its representative with ``merge``, otherwise + it founds a new object. Upstream observations may carry a single + :class:`Detection3DPC` or a per-frame :class:`ImageDetections` batch. + """ + + def __init__( + self, + is_same: Callable[[Detection3DPC, Detection3DPC], bool] | None = None, + merge: Callable[[Detection3DPC, Detection3DPC], Detection3DPC] | None = None, + ) -> None: + self.is_same = is_same or spatial() + self.merge = merge or operator.add + self.groups: list[list[Detection3DPC]] = [] + self.merged: list[Detection3DPC] = [] + + def add(self, detection: Detection3DPC) -> int: + """Assign one sighting to its object; returns the object's index.""" + for index, representative in enumerate(self.merged): + if self.is_same(representative, detection): + self.groups[index].append(detection) + self.merged[index] = self.merge(representative, detection) + return index + self.groups.append([detection]) + self.merged.append(detection) + return len(self.groups) - 1 + + def __call__( + self, upstream: Iterator[Observation[Any]] + ) -> Iterator[Observation[Detection3DPC]]: + template: Observation[Any] | None = None + for obs in upstream: + template = obs + data = obs.data + for detection in data if isinstance(data, ImageDetections) else [data]: + self.add(detection) + if template is None: + return + for merged in self.merged: + yield template.derive(data=merged, ts=merged.ts, pose=merged.pose) diff --git a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py index e63d415b8a..50e0735b9a 100644 --- a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py +++ b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py @@ -41,16 +41,26 @@ def from_2d( world_to_optical_transform: Transform, filters: list[PointCloudFilter] | None = None, ) -> ImageDetections3DPC: - """Project every 2D detection into 3D, dropping any that yield no valid points.""" + """Project every 2D detection into 3D, dropping any that yield no valid points. + + The cloud is projected through the camera once; each detection then + selects its points from that shared projection. + """ + world_points, points_2d = Detection3DPC.project_cloud( + world_pointcloud, camera_info, world_to_optical_transform + ) detections_3d = [ d3d for det in detections_2d if ( - d3d := Detection3DPC.from_2d( + d3d := Detection3DPC.from_projection( det, - world_pointcloud, + world_points, + points_2d, camera_info, world_to_optical_transform, + world_pointcloud.frame_id, + world_pointcloud.ts, filters, ) ) diff --git a/dimos/perception/detection/type/detection3d/pointcloud.py b/dimos/perception/detection/type/detection3d/pointcloud.py index 691c7ae32a..786f2152ae 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -47,6 +47,28 @@ class Detection3DPC(Detection3D): def center(self) -> Vector3: return Vector3(*self.pointcloud.center) + def __add__(self, other: Detection3DPC) -> Detection3DPC: + """Union of two sightings of one object. + + The cloud is the union of both; identity metadata follows the + higher-confidence sighting and time context follows the later one, + matching latest-pose semantics. + """ + later = self if self.ts >= other.ts else other + stronger = self if self.confidence >= other.confidence else other + return Detection3DPC( + image=later.image, + bbox=later.bbox, + track_id=self.track_id, + class_id=stronger.class_id, + confidence=stronger.confidence, + name=stronger.name, + ts=later.ts, + pointcloud=self.pointcloud + other.pointcloud, + transform=later.transform, + frame_id=self.frame_id, + ) + @functools.cached_property def pose(self) -> PoseStamped: """Convert detection to a PoseStamped using pointcloud center. @@ -198,52 +220,22 @@ def from_depth( frame_id=detection_pc.frame_id, ) - @classmethod - def from_2d( # type: ignore[override] - cls, - det: Detection2DBBox, + @staticmethod + def project_cloud( world_pointcloud: PointCloud2, camera_info: CameraInfo, world_to_optical_transform: Transform, - # filters are to be adjusted based on the sensor noise characteristics if feeding - # sensor data directly - filters: list[PointCloudFilter] | None = None, - ) -> Detection3DPC | None: - """Create a Detection3D from a 2D detection by projecting world pointcloud. + ) -> tuple[np.ndarray, np.ndarray]: + """Project a world cloud through the camera once. - This method handles: - 1. Projecting world pointcloud to camera frame - 2. Filtering points within the 2D detection bounding box - 3. Cleaning up the pointcloud (height filter, outlier removal) - 4. Hidden point removal from camera perspective - - Args: - det: The 2D detection - world_pointcloud: Full pointcloud in world frame - camera_info: Camera calibration info - world_to_camerlka_transform: Transform from world to camera frame - filters: List of functions to apply to the pointcloud for filtering - Returns: - Detection3D with filtered pointcloud, or None if no valid points + Returns the world points that land inside the image and their pixel + coordinates - the detection-independent half of ``from_2d``, shared + by every detection of one frame via ``from_projection``. """ - # Set default filters if none provided - if filters is None: - filters = [ - # height_filter(0.1), - raycast(), - radius_outlier(), - statistical(), - ] - - # Extract camera parameters fx, fy = camera_info.K[0], camera_info.K[4] cx, cy = camera_info.K[2], camera_info.K[5] - image_width = camera_info.width - image_height = camera_info.height - camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) - # Convert pointcloud to numpy array world_points, _ = world_pointcloud.as_numpy() # Project points to camera frame @@ -257,7 +249,7 @@ def from_2d( # type: ignore[override] world_points = world_points[valid_mask] if len(world_points) == 0: - return None + return world_points, np.empty((0, 2)) # Project to 2D points_2d_homogeneous = (camera_matrix @ points_camera[:, :3].T).T @@ -266,12 +258,38 @@ def from_2d( # type: ignore[override] # Filter points within image bounds in_image_mask = ( (points_2d[:, 0] >= 0) - & (points_2d[:, 0] < image_width) + & (points_2d[:, 0] < camera_info.width) & (points_2d[:, 1] >= 0) - & (points_2d[:, 1] < image_height) + & (points_2d[:, 1] < camera_info.height) ) - points_2d = points_2d[in_image_mask] - world_points = world_points[in_image_mask] + return world_points[in_image_mask], points_2d[in_image_mask] + + @classmethod + def from_projection( + cls, + det: Detection2DBBox, + world_points: np.ndarray, + points_2d: np.ndarray, + camera_info: CameraInfo, + world_to_optical_transform: Transform, + frame_id: str, + timestamp: float, + filters: list[PointCloudFilter] | None = None, + ) -> Detection3DPC | None: + """Create a Detection3D by selecting from a shared frame projection. + + ``world_points`` and ``points_2d`` come from ``project_cloud`` for + this detection's frame; only the mask selection and the per-detection + filters run here. + """ + # Set default filters if none provided + if filters is None: + filters = [ + # height_filter(0.1), + raycast(), + radius_outlier(), + statistical(), + ] if len(world_points) == 0: return None @@ -296,14 +314,13 @@ def from_2d( # type: ignore[override] detection_points = world_points[in_det_mask] if detection_points.shape[0] == 0: - # print(f"No points found in detection bbox after projection. {det.name}") return None # Create initial pointcloud for this detection initial_pc = PointCloud2.from_numpy( detection_points, - frame_id=world_pointcloud.frame_id, - timestamp=world_pointcloud.ts, + frame_id=frame_id, + timestamp=timestamp, ) # Apply filters - each filter gets all arguments @@ -329,5 +346,36 @@ def from_2d( # type: ignore[override] ts=det.ts, pointcloud=detection_pc, transform=world_to_optical_transform, - frame_id=world_pointcloud.frame_id, + frame_id=frame_id, + ) + + @classmethod + def from_2d( # type: ignore[override] + cls, + det: Detection2DBBox, + world_pointcloud: PointCloud2, + camera_info: CameraInfo, + world_to_optical_transform: Transform, + # filters are to be adjusted based on the sensor noise characteristics if feeding + # sensor data directly + filters: list[PointCloudFilter] | None = None, + ) -> Detection3DPC | None: + """Create a Detection3D from a 2D detection by projecting world pointcloud. + + One-detection convenience over ``project_cloud`` + ``from_projection``; + callers lifting several detections of one frame should project once + and call ``from_projection`` per detection instead. + """ + world_points, points_2d = cls.project_cloud( + world_pointcloud, camera_info, world_to_optical_transform + ) + return cls.from_projection( + det, + world_points, + points_2d, + camera_info, + world_to_optical_transform, + world_pointcloud.frame_id, + world_pointcloud.ts, + filters, ) diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index f7c3715cd1..c94578522a 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -57,12 +57,14 @@ class DanDetector(Resource): segmenter: EdgeTAMImageSegmenter def start(self) -> None: + import torch + from dimos.models.embedding.siglip import SigLIPModel from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector self.siglip = SigLIPModel() - self.detector = Owlv2Detector() + self.detector = Owlv2Detector(dtype=torch.float16) self.segmenter = EdgeTAMImageSegmenter() self._live: list[DisposableBase] = [] diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index d8cd3e321d..7ad30fd861 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -38,6 +38,7 @@ from dimos.memory.embed import EmbedImages from dimos.memory.transform import throttle +from dimos.perception.detection.identity import Identity, spatial from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Localization, LocalizePolicy, Support @@ -54,51 +55,54 @@ TIME_BANDS = 6 # stratify retrieval across the window so late scans always compete BOXES_PER_FRAME = 4 +# Frames per OWLv2 forward. The forward's softmax transient is fp32 even +# under autocast (~580MB per image at 960px), so the profitable batch is +# whatever VRAM holds beyond the resident models - 1 on an 8GB card. +DETECT_BATCH = 1 CONFIRM_FLOOR = 0.22 # geometric confirmation accept for re-detections VERIFY_FRAMES = 20 -@dataclass -class _ClusterObservation: - ts: float - score: float - centroid: np.ndarray - cloud: Any - camera_position: np.ndarray - detection: Detection3DPC +# A support candidate is an identity group: the member sightings of one +# object. Everything a group reports is a plain function over its members. -@dataclass -class _Cluster: - center: np.ndarray - observations: list[_ClusterObservation] = field(default_factory=list) +def _centroid(det: Detection3DPC) -> np.ndarray: + centroid: np.ndarray = np.asarray(det.pointcloud.pointcloud.points).mean(axis=0) + return centroid + + +def _camera_position(det: Detection3DPC) -> np.ndarray: + position = (-det.transform).translation + return np.array([position.x, position.y, position.z]) + + +def _group_center(members: list[Detection3DPC]) -> np.ndarray: + center: np.ndarray = np.mean(np.stack([_centroid(d) for d in members]), axis=0) + return center + + +def _max_score(members: list[Detection3DPC]) -> float: + return max(d.confidence for d in members) + - def add(self, obs: _ClusterObservation) -> None: - self.observations.append(obs) - self.center = np.mean(np.stack([o.centroid for o in self.observations]), axis=0) +def _latest(members: list[Detection3DPC]) -> Detection3DPC: + return max(members, key=lambda d: d.ts) - @property - def max_score(self) -> float: - return max(o.score for o in self.observations) - @property - def latest(self) -> _ClusterObservation: - return max(self.observations, key=lambda o: o.ts) +def _interval(members: list[Detection3DPC]) -> tuple[float, float]: + times = [d.ts for d in members] + return min(times), max(times) - @property - def interval(self) -> tuple[float, float]: - times = [o.ts for o in self.observations] - return min(times), max(times) - @property - def n_views(self) -> int: - return len({tuple(np.round(o.camera_position, 2)) for o in self.observations}) +def _n_views(members: list[Detection3DPC]) -> int: + return len({tuple(np.round(_camera_position(d), 2)) for d in members}) - @property - def extent(self) -> np.ndarray: - points = np.concatenate([np.asarray(o.cloud.pointcloud.points) for o in self.observations]) - extent: np.ndarray = points.max(axis=0) - points.min(axis=0) - return extent + +def _group_extent(members: list[Detection3DPC]) -> np.ndarray: + points = np.concatenate([np.asarray(d.pointcloud.pointcloud.points) for d in members]) + extent: np.ndarray = points.max(axis=0) - points.min(axis=0) + return extent @dataclass @@ -119,10 +123,10 @@ def _quaternion_from_matrix(rotation: np.ndarray) -> tuple[float, float, float, return (float(x), float(y), float(z), float(w)) -def _azimuth_coverage(observations: list[_ClusterObservation], center: np.ndarray) -> float: +def _azimuth_coverage(members: list[Detection3DPC], center: np.ndarray) -> float: directions = [] - for obs in observations: - v = obs.camera_position - center + for det in members: + v = _camera_position(det) - center norm = np.linalg.norm(v) if norm > 1e-6: directions.append(v / norm) @@ -134,12 +138,10 @@ def _azimuth_coverage(observations: list[_ClusterObservation], center: np.ndarra return len(bins) / 8.0 -def _axes_observed( - observations: list[_ClusterObservation], center: np.ndarray -) -> tuple[bool, bool, bool]: +def _axes_observed(members: list[Detection3DPC], center: np.ndarray) -> tuple[bool, bool, bool]: directions = [] - for obs in observations: - v = obs.camera_position - center + for det in members: + v = _camera_position(det) - center norm = np.linalg.norm(v) if norm > 1e-6: directions.append(v / norm) @@ -150,54 +152,82 @@ def _axes_observed( class _DetectionCache: - """One OWLv2 + EdgeTAM pass per unique frame, shared across queries and clusters.""" + """One OWLv2 + EdgeTAM + lift pass per unique frame, shared across queries and clusters.""" - def __init__(self, detector: Any, segmenter: Any, queries: list[str], floor: float) -> None: + def __init__( + self, detector: Any, segmenter: Any, rig: Rig, queries: list[str], floor: float + ) -> None: self.detector = detector self.segmenter = segmenter + self.rig = rig self.queries = queries self.floor = floor self._cache: dict[float, ImageDetections2D] = {} + self._lifted: dict[float, list[Detection3DPC]] = {} + + def _ingest(self, image: Any, detections: ImageDetections2D) -> None: + ranked = sorted(detections.detections, key=lambda d: -d.confidence) + kept = [ + det + for label in self.queries + for det in [d for d in ranked if d.name == label][:BOXES_PER_FRAME] + ] + for i, det in enumerate(kept): + det.track_id = i + cached = ImageDetections2D(image, kept) + if len(cached): + cached = self.segmenter.segment(cached) + self._cache[image.ts] = cached + + def prefetch(self, images: list[Any]) -> None: + """Detect and segment every uncached frame, batching the OWLv2 forwards.""" + todo: dict[float, Any] = {} + for image in images: + if image.ts not in self._cache and image.ts not in todo: + todo[image.ts] = image + pending = list(todo.values()) + for start in range(0, len(pending), DETECT_BATCH): + chunk = pending[start : start + DETECT_BATCH] + for image, detections in zip( + chunk, + self.detector.query_detections_batch(chunk, self.queries, threshold=self.floor), + strict=True, + ): + self._ingest(image, detections) def detect(self, image: Any, query: str) -> ImageDetections2D: - key = image.ts - cached = self._cache.get(key) - if cached is None: - detections: ImageDetections2D = self.detector.query_detections( - image, self.queries, threshold=self.floor + if image.ts not in self._cache: + self._ingest( + image, self.detector.query_detections(image, self.queries, threshold=self.floor) ) - ranked = sorted(detections.detections, key=lambda d: -d.confidence) - cached = ImageDetections2D( - image, - [ - det - for label in self.queries - for det in [d for d in ranked if d.name == label][:BOXES_PER_FRAME] - ], - ) - if len(cached): - cached = self.segmenter.segment(cached) - self._cache[key] = cached - return cached.filter(lambda d: d.name == query) + return self._cache[image.ts].filter(lambda d: d.name == query) + + def lift(self, ts: float) -> list[Detection3DPC]: + """All of a frame's detections lifted once, whatever their label.""" + if ts not in self._lifted: + lifted = self.rig.lift(self._cache[ts]) + self._lifted[ts] = list(lifted) if lifted is not None else [] + return self._lifted[ts] def _lift( detections: ImageDetections2D, + cache: _DetectionCache, rig: Rig, policy: LocalizePolicy, plane: Any | None = None, ) -> list[tuple[Detection3DPC, np.ndarray]]: - """Lift 2D detections through the rig; returns valid (detection3d, camera_position) pairs.""" + """Gate a query's share of the frame's shared lift; (detection3d, camera_position) pairs.""" pose = rig.camera_pose(detections.ts) if pose is None: return [] camera = np.array([pose.position.x, pose.position.y, pose.position.z]) - lifted = rig.lift(detections) - if lifted is None: - return [] + track_ids = {det.track_id for det in detections} valid: list[tuple[Detection3DPC, np.ndarray]] = [] - for det3d in lifted: + for det3d in cache.lift(detections.ts): + if det3d.track_id not in track_ids: + continue points = np.asarray(det3d.pointcloud.pointcloud.points) if len(points) < policy.min_depth_points: continue @@ -321,7 +351,7 @@ def localize( traces: list[LocalizeTrace | None] = ( list(trace) if isinstance(trace, list) else [trace] * len(queries) ) - cache = _DetectionCache(detector, segmenter, queries, policy.candidate_floor) + cache = _DetectionCache(detector, segmenter, rig, queries, policy.candidate_floor) results = [ _localize_one( q, @@ -361,9 +391,12 @@ def _localize_one( from dimos.perception.memory.support_plane import fit_support_plane plane = fit_support_plane(rig, frames) + cache.prefetch([obs.data for obs in frames]) - # Pass 2 - OWLv2 + EdgeTAM: detect, segment, lift, verify. - clusters: list[_Cluster] = [] + # Pass 2 - OWLv2 + EdgeTAM: detect, segment, lift, verify. Support + # candidates are identity groups: sightings joined by the spatial + # is_same strategy at the policy's cluster radius. + identity = Identity(is_same=spatial(policy.cluster_radius_m)) ungrounded_best: tuple[float, float] | None = None # (score, ts) processed: set[float] = set() @@ -377,43 +410,29 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: return if trace is not None and not is_verify: trace.detection_frames.append(frame_obs.derive(data=detections)) - lifted = _lift(detections, rig, policy, plane) + lifted = _lift(detections, cache, rig, policy, plane) for det2d in detections: if not any(d.track_id == det2d.track_id for d, _ in lifted): best = (det2d.confidence, det2d.ts) if ungrounded_best is None or best[0] > ungrounded_best[0]: ungrounded_best = best - for det3d, camera in lifted: - observation = _ClusterObservation( - ts=det3d.ts, - score=det3d.confidence, - centroid=np.asarray(det3d.pointcloud.pointcloud.points).mean(axis=0), - cloud=det3d.pointcloud, - camera_position=camera, - detection=det3d, - ) + for det3d, _camera in lifted: if trace is not None: (trace.verified if is_verify else trace.matched).append((det3d.ts, det3d)) - for cluster in clusters: - distance = float(np.linalg.norm(observation.centroid - cluster.center)) - if distance <= policy.cluster_radius_m: - cluster.add(observation) - break - else: - clusters.append(_Cluster(center=observation.centroid, observations=[observation])) + identity.add(det3d) for frame_obs in frames: _absorb(frame_obs, is_verify=False) - logger.info(f"detection: {len(clusters)} support candidates") + logger.info(f"detection: {len(identity.groups)} support candidates") # Cross-view verification: a support seen from one pose only is # unconfirmed. Frames whose camera could observe the support are found # geometrically (near + sees with occlusion), then re-detected. - clusters.sort(key=lambda c: -c.max_score) - for cluster in list(clusters[:4]): + for members in sorted(identity.groups, key=_max_score, reverse=True)[:4]: + center = _group_center(members) predicate = rig.sees( - cluster.center, - extent=np.minimum(cluster.extent, 0.4), + center, + extent=np.minimum(_group_extent(members), 0.4), # A large object overflows close-up frames; a third of its box in # view is still a usable re-detection pass, and those close-ups # are exactly the distinct viewpoints verification needs. @@ -422,7 +441,7 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: ) observing = [ obs - for obs in index.near(cluster.center, radius=policy.verify_radius_m) + for obs in index.near(center, radius=policy.verify_radius_m) if obs.ts not in processed and rig.camera_still(obs.ts) and predicate(obs) ] if len(observing) > VERIFY_FRAMES: @@ -430,16 +449,19 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: # latest seeing frames would bias the latest-pose answer early. picks = np.unique(np.linspace(0, len(observing) - 1, VERIFY_FRAMES).astype(int)) observing = [observing[i] for i in picks] + cache.prefetch([obs.data for obs in observing]) for frame_obs in observing: _absorb(frame_obs, is_verify=True) verified = [ - c for c in clusters if c.max_score >= policy.accept_score and c.n_views >= policy.min_views + members + for members in identity.groups + if _max_score(members) >= policy.accept_score and _n_views(members) >= policy.min_views ] logger.info( "verification: " + ", ".join( - f"score={c.max_score:.2f} views={c.n_views} obs={len(c.observations)}" for c in clusters + f"score={_max_score(g):.2f} views={_n_views(g)} obs={len(g)}" for g in identity.groups ) ) @@ -469,54 +491,56 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: ) return None - winner = max(verified, key=lambda c: c.latest.ts) - w_lo, w_hi = winner.interval + winner = max(verified, key=lambda g: _latest(g).ts) + w_lo, w_hi = _interval(winner) rival_scores = [ - c.max_score - for c in verified - if c is not winner - and not (c.interval[1] < w_lo or c.interval[0] > w_hi) # coexisting in time + _max_score(g) + for g in verified + if g is not winner + and not (_interval(g)[1] < w_lo or _interval(g)[0] > w_hi) # coexisting in time ] - margin = winner.max_score - max(rival_scores) if rival_scores else 1.0 + margin = _max_score(winner) - max(rival_scores) if rival_scores else 1.0 reason = "ambiguous_between_coexisting_candidates" if margin < policy.refusal_margin else None - latest = winner.latest - points = np.asarray(latest.cloud.pointcloud.points) + latest = _latest(winner) + points = np.asarray(latest.pointcloud.pointcloud.points) aabb_min, aabb_max = points.min(axis=0), points.max(axis=0) try: - orientation = _quaternion_from_matrix(np.asarray(latest.cloud.oriented_bounding_box.R)) + orientation = _quaternion_from_matrix(np.asarray(latest.pointcloud.oriented_bounding_box.R)) except Exception: orientation = (0.0, 0.0, 0.0, 1.0) center = (aabb_min + aabb_max) / 2 extent = np.maximum(aabb_max - aabb_min, 0.005) sigma = ( - np.stack([o.centroid for o in winner.observations]).std(axis=0) - if len(winner.observations) > 1 + np.stack([_centroid(d) for d in winner]).std(axis=0) + if len(winner) > 1 else np.full(3, 0.01) ) + winner_center = _group_center(winner) support = Support( center_xyz=(float(center[0]), float(center[1]), float(center[2])), extent_xyz_m=(float(extent[0]), float(extent[1]), float(extent[2])), orientation_xyzw=(0.0, 0.0, 0.0, 1.0), sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), - coverage=_azimuth_coverage(winner.observations, winner.center), - axes_observed=_axes_observed(winner.observations, winner.center), + coverage=_azimuth_coverage(winner, winner_center), + axes_observed=_axes_observed(winner, winner_center), frame_id=rig.world_frame, ) if trace is not None: - trace.answer = latest.detection + trace.answer = latest trace.backdrop_ts = latest.ts + latest_centroid = _centroid(latest) return Localization( instance_id="query-0", - semantic_score=winner.max_score, - identity_score=min(1.0, winner.n_views / 4.0), + semantic_score=_max_score(winner), + identity_score=min(1.0, _n_views(winner) / 4.0), ambiguity_margin=margin, position_world_xyz=( - float(latest.centroid[0]), - float(latest.centroid[1]), - float(latest.centroid[2]), + float(latest_centroid[0]), + float(latest_centroid[1]), + float(latest_centroid[2]), ), orientation_world_xyzw=orientation, frame_id=rig.world_frame, @@ -524,9 +548,9 @@ def _absorb(frame_obs: Any, is_verify: bool) -> None: pose_timestamp=latest.ts, geometry_timestamp=latest.ts, last_seen_timestamp=latest.ts, - point_cloud=latest.cloud, + point_cloud=latest.pointcloud, cloud_mode=cloud_mode, coverage=support.coverage, - n_views=winner.n_views, + n_views=_n_views(winner), reason=reason, ) diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 91d086d6e0..a3554d8935 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -32,12 +32,16 @@ from __future__ import annotations +from collections import OrderedDict from dataclasses import dataclass, field +import json +from pathlib import Path from typing import TYPE_CHECKING, Any, cast import numpy as np from dimos.memory.tf import StreamTF +from dimos.memory.transform import Transformer from dimos.msgs.geometry_msgs.Transform import Transform from dimos.perception.detection.project import sees as project_sees from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC @@ -48,9 +52,10 @@ from dimos.perception.memory import gates from dimos.perception.memory.gates import SPEED_MAX, STILL_ENVELOPE, TF_TOLERANCE from dimos.perception.memory.types import InventoryPolicy, LocalizePolicy +from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Iterator from dimos_lcm.sensor_msgs import CameraInfo @@ -62,13 +67,21 @@ from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.protocol.tf.tf import TFLookup +logger = setup_logger() + # Pose-stamped rigs ride per-frame odometry, so walking does not stale the # projection the way a sweeping wrist stales interpolated tf; the gate only # drops speed glitches and sprints. WALK_SPEED_MAX = 1.5 DEPTH_TOLERANCE = 0.06 # s - temporal join color->depth -CLOUD_ACCUM_S = 2.0 # s - scans within this of a frame form its geometry +# Scans within this of a frame form its geometry. Wide enough that a +# spinning lidar's near-floor blind ring is filled by scans taken from +# earlier and later poses - the scene is static in world frame. +CLOUD_ACCUM_S = 4.0 + +_SCAN_CACHE_MAX = 256 # registered scans held per rig; a window needs a few dozen +_CELL_OFFSET = 1 << 20 # shifts lattice cell indices positive for 21-bit key packing EMBED_HZ = 1.0 # index density for a wrist camera parked over a workspace # A walking robot changes viewpoint every frame and its frames blur @@ -96,6 +109,108 @@ surface_patch_min_drop_m=-0.06, verify_radius_m=5.0, ) + + +MOBILE_SPAN_M = 3.0 # camera translation beyond this means a mobile base + + +def _tf_root(store: Any, tf_name: str) -> str | None: + """The tf tree's root frame: a parent that is never a child.""" + parents: set[str] = set() + children: set[str] = set() + for obs in store.stream(tf_name).limit(500): + for transform in obs.data.transforms: + parents.add(transform.frame_id) + children.add(transform.child_frame_id) + roots = parents - children + return roots.pop() if len(roots) == 1 else None + + +def _stream_rate(stream: Any) -> float: + count: int = stream.count() + if count < 2: + return float(count) + t0, t1 = stream.get_time_range() + return count / max(float(t1 - t0), 1e-6) + + +def _camera_span(rig: Rig) -> float: + """Diagonal of the camera positions' bounding box over the recording.""" + if rig.tf is None and rig.poses is None: + return 0.0 # no pose source at all: an embed-only store being seeded + try: + t0, t1 = rig.color.get_time_range() + except LookupError: + return 0.0 # live store, nothing recorded yet + positions = [] + for k in range(12): + pose = rig.camera_pose(t0 + (t1 - t0) * k / 11) + if pose is not None: + positions.append([pose.position.x, pose.position.y, pose.position.z]) + if len(positions) < 2: + return 0.0 + spread = np.array(positions) + return float(np.linalg.norm(spread.max(axis=0) - spread.min(axis=0))) + + +def _lattice_quantum(points: np.ndarray) -> float | None: + """The grid pitch when coordinates lie on a lattice; None for continuous scans. + + Grid-quantized sources (an occupancy map streamed as clouds) repeat the + same cell across snapshots and dedup by cell key; continuous scans never + collide and skip dedup entirely. + """ + sample = points[:2048] + x = np.unique(sample[:, 0]) + if len(x) < 8: + return None + diffs = np.diff(x) + diffs = diffs[diffs > 1e-9] + if len(diffs) == 0: + return None + quantum = float(diffs.min()) + if quantum < 1e-4: + return None + scaled = sample / quantum + if float(np.abs(scaled - np.round(scaled)).max()) > 0.01: + return None + return quantum + + +class RegisterScans(Transformer["PointCloud2", "PointCloud2"]): + """Map sensor-frame scans into a world frame through tf, one transform per scan. + + A plain stream transformer, so the registered cloud is a derived stream: + ``scans.transform(RegisterScans(tf, world))`` yields world-frame clouds, + ``.save(...)`` persists them, and a live pipeline can tail-register the + same way the embed pipeline does. Scans already in the world frame pass + through untouched; scans with no transform at their time are dropped. + """ + + def __init__(self, tf: TFLookup, world_frame: str, tolerance: float = TF_TOLERANCE) -> None: + self.tf = tf + self.world_frame = world_frame + self.tolerance = tolerance + + def __call__( + self, upstream: Iterator[Observation[PointCloud2]] + ) -> Iterator[Observation[PointCloud2]]: + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + for obs in upstream: + if obs.data.frame_id == self.world_frame: + yield obs + continue + transform = self.tf.get(self.world_frame, obs.data.frame_id, obs.ts, self.tolerance) + if transform is None: + continue + matrix = transform.to_matrix() + points = obs.data.as_numpy()[0] @ matrix[:3, :3].T + matrix[:3, 3] + yield obs.derive( + data=PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=obs.ts) + ) + + ROOM_INVENTORY_POLICY = InventoryPolicy( keyframe_stride_s=1.25, min_mask_area_px=900, @@ -133,64 +248,211 @@ class Rig: base_to_optical: Transform | None = None poses: Any = None # stream carrying world base poses (e.g. odom) depth: Any = None # aligned depth stream, lifted via from_depth - cloud: Any = None # world-frame pointcloud stream, lifted via from_2d + cloud: Any = None # pointcloud stream, lifted via from_2d after registration tf_tolerance: float = TF_TOLERANCE cloud_accum_s: float = CLOUD_ACCUM_S speed_max: float = SPEED_MAX scene_gate: bool = True embed_hz: float = EMBED_HZ + mobile: bool = False # camera rides a mobile base: room-scale policies _cloud_memo: tuple[float, PointCloud2] | None = field(default=None, repr=False, init=False) + _scan_cache: OrderedDict[float, np.ndarray | None] = field( + default_factory=OrderedDict, repr=False, init=False + ) + _quantum: float | None = field(default=None, repr=False, init=False) + _quantum_known: bool = field(default=False, repr=False, init=False) @classmethod - def from_store(cls, store: Any) -> Rig: - """Recognize the store's shape. - - A ``tf`` stream wins as pose source; without one the observations' - stamped poses are used with the Go2 front-camera mount. ``depth_image`` - wins as geometry source; a ``lidar`` stream is next, registered in - whatever world frame its scans carry; a store with neither can still - embed and retrieve, just never lift. Intrinsics and the optical frame - name come from the ``camera_info`` stream, or - for stores that carry - none - the static Go2 front-camera calibration. Stream contents are - only read where a name must be sniffed, so a live store whose streams - are still empty resolves too. + def from_store( + cls, + store: Any, + manifest: dict[str, Any] | None = None, + overrides: dict[str, str] | None = None, + ) -> Rig: + """Recognize the store's shape without depending on stream names. + + Resolution order per role: ``overrides`` (explicit stream names from + a caller or CLI), then ``manifest`` (passed in, or the ``.rig.json`` + sidecar next to the recording), then discovery by stream data type + and content: TFMessage-typed stream as tf, Image streams classified + by their frames (uint16/float is depth, distinct-channel uint8 is + color; a lossy-coded gray stream is neither), a PointCloud2 stream as + geometry when there is no metric depth, and - when tf is absent - the + highest-rate pose-stamped stream as pose source. On tf rigs the world + frame is the tf tree's root; ambiguity or missing calibration raises + with the candidates rather than guessing. + + The manifest carries roles as stream names plus, for recordings whose + calibration was never recorded, an inline ``camera_info`` dict and a + ``base_to_optical`` mount. """ - streams = store.list_streams() - color = store.streams.color_image - tf = StreamTF.from_store(store) + from dimos.msgs.geometry_msgs.Quaternion import Quaternion + from dimos.msgs.geometry_msgs.Vector3 import Vector3 + from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo as CameraInfoMsg + from dimos.msgs.sensor_msgs.Image import Image as ImageMsg + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as PointCloudMsg + from dimos.msgs.tf2_msgs.TFMessage import TFMessage + + names = store.list_streams() + types = {name: store.stream(name).data_type for name in names} + + if manifest is None: + path = getattr(store.config, "path", None) + if path: + sidecar = Path(f"{path}.rig.json") + if sidecar.exists(): + manifest = json.loads(sidecar.read_text()) + logger.info(f"rig: manifest {sidecar}") + roles: dict[str, Any] = dict(manifest or {}) + roles.update(overrides or {}) + + claimed = {value for value in roles.values() if isinstance(value, str)} + color_name: str | None = roles.get("color") + depth_name: str | None = roles.get("depth") + cloud_name: str | None = roles.get("cloud") + poses_name: str | None = roles.get("poses") + + tf_names = [n for n in names if types[n] is TFMessage] + tf_name = tf_names[0] if len(tf_names) == 1 else ("tf" if "tf" in tf_names else None) + tf = StreamTF.from_store(store, tf_name) if tf_name is not None else None + + # image streams classify by content: metric depth or genuine color + image_names = [n for n in names if types[n] is ImageMsg and n not in claimed] + color_candidates: list[str] = [] + depth_candidates: list[str] = [] + empty_images: list[str] = [] + for name in image_names: + stream = store.stream(name) + if stream.count() == 0: + empty_images.append(name) + continue + frame = stream.first().data.data + if frame.dtype == np.uint16 or frame.dtype.kind == "f": + depth_candidates.append(name) + elif ( + frame.ndim == 3 + and frame.shape[2] == 3 + and not np.array_equal(frame[..., 0], frame[..., 1]) + ): + color_candidates.append(name) + else: + logger.info(f"rig: image stream {name!r} is neither color nor metric depth") + if color_name is None: + if len(color_candidates) > 1: + color_name = max(color_candidates, key=lambda n: _stream_rate(store.stream(n))) + logger.info(f"rig: several color streams, using highest-rate {color_name!r}") + elif color_candidates: + color_name = color_candidates[0] + elif len(empty_images) == 1: + color_name = empty_images[0] # a live store's not-yet-filled feed + if color_name is None: + raise ValueError(f"no color image stream among {names}; pass a manifest or --color") + claimed.add(color_name) + if depth_name is None: + if len(depth_candidates) > 1: + raise ValueError(f"several depth streams {depth_candidates}; pass --depth") + depth_name = depth_candidates[0] if depth_candidates else None + + if cloud_name is None and depth_name is None: + cloud_names = [n for n in names if types[n] is PointCloudMsg and n not in claimed] + if len(cloud_names) > 1: + raise ValueError(f"several pointcloud streams {cloud_names}; pass --cloud") + cloud_name = cloud_names[0] if cloud_names else None + + color = store.stream(color_name) + depth = store.stream(depth_name) if depth_name is not None else None + cloud = store.stream(cloud_name) if cloud_name is not None and depth is None else None + if cloud_name is not None: + claimed.add(cloud_name) - depth = store.streams.depth_image if "depth_image" in streams else None - cloud = None world_frame = gates.WORLD_FRAME - if depth is None and "lidar" in streams: - cloud = store.streams.lidar + if tf is not None: + world_frame = _tf_root(store, cast("str", tf_name)) or world_frame + elif cloud is not None: try: world_frame = cloud.first().data.frame_id except LookupError: pass # live store, nothing recorded yet - if "camera_info" in streams: - camera_info = store.streams.camera_info.first().data - optical_frame = camera_info.frame_id + # intrinsics: inline manifest dict, named stream, or discovery by + # type with the color camera's frame deciding among several + camera_info = None + camera_info_role = roles.get("camera_info") + if isinstance(camera_info_role, dict): + camera_info = CameraInfoMsg( + height=camera_info_role["height"], + width=camera_info_role["width"], + distortion_model=camera_info_role.get("distortion_model", ""), + D=camera_info_role.get("D"), + K=camera_info_role["K"], + R=camera_info_role.get("R"), + P=camera_info_role.get("P"), + frame_id=camera_info_role["frame_id"], + ) else: - from dimos.robot.unitree.go2.connection import GO2Connection - - camera_info = GO2Connection.camera_info_static - # With a tf tree the optical name must match that tree; the Go2 - # calibration names apply only to the tf-less Go2 shape. - optical_frame = camera_info.frame_id if tf is None else gates.OPTICAL_FRAME + ci_name = camera_info_role if isinstance(camera_info_role, str) else None + if ci_name is None: + candidates = [ + n for n in names if types[n] is CameraInfoMsg and store.stream(n).count() + ] + try: + color_frame = color.first().data.frame_id + except LookupError: + color_frame = None + matching = [ + n for n in candidates if store.stream(n).first().data.frame_id == color_frame + ] + if matching: + ci_name = sorted(matching)[0] + elif len(candidates) == 1: + ci_name = candidates[0] + elif len(candidates) > 1: + raise ValueError(f"several camera_info streams {candidates}; pass a manifest") + if ci_name is not None: + camera_info = store.stream(ci_name).first().data base_to_optical = None + mount = roles.get("base_to_optical") + if isinstance(mount, dict): + base_to_optical = Transform( + translation=Vector3(*mount["translation"]), + rotation=Quaternion(*mount["rotation"]), + frame_id="base_link", + child_frame_id=camera_info.frame_id if camera_info else "camera_optical", + ) + poses = None if tf is None: - from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL - - base_to_optical = BASE_TO_OPTICAL - poses = store.streams.odom - - mobile = cloud is not None - return cls( - camera_info=camera_info, + if poses_name is None: + posed = [ + n + for n in names + if n not in claimed + and n != color_name + and store.stream(n).count() + and store.stream(n).first().pose_tuple is not None + ] + if posed: + poses_name = max(posed, key=lambda n: _stream_rate(store.stream(n))) + poses = store.stream(poses_name) if poses_name is not None else None + + if depth is not None or cloud is not None: + if camera_info is None: + raise ValueError( + "store has 3D geometry but no camera calibration; add a CameraInfo " + "stream role or an inline camera_info to the .rig.json manifest" + ) + if tf is None and (poses is None or base_to_optical is None): + raise ValueError( + "store has no tf; a pose-stamped rig needs a poses stream and a " + "base_to_optical mount in the .rig.json manifest" + ) + + # embed-only stores (no geometry) may carry no calibration at all; + # every geometry API raises on use, embedding never touches it + optical_frame = camera_info.frame_id if camera_info is not None else gates.OPTICAL_FRAME + rig = cls( + camera_info=cast("CameraInfo", camera_info), color=color, world_frame=world_frame, optical_frame=optical_frame, @@ -199,12 +461,21 @@ def from_store(cls, store: Any) -> Rig: poses=poses, depth=depth, cloud=cloud, - speed_max=WALK_SPEED_MAX if mobile else SPEED_MAX, - scene_gate=not mobile, - embed_hz=WALK_EMBED_HZ if mobile else EMBED_HZ, ) - # pose + span = _camera_span(rig) + rig.mobile = span > MOBILE_SPAN_M + if rig.mobile: + rig.speed_max = WALK_SPEED_MAX + rig.scene_gate = False + rig.embed_hz = WALK_EMBED_HZ + logger.info( + f"rig: color={color_name!r} depth={depth_name!r} cloud={cloud_name!r} " + f"tf={tf_name!r} world={world_frame!r} span={span:.1f}m mobile={rig.mobile}" + ) + return rig + + # pose def world_to_optical(self, ts: float) -> Transform | None: if self.tf is not None: @@ -279,7 +550,7 @@ def still_intervals(self, t0: float, t1: float) -> list[tuple[float, float]]: intervals.append((run_start, float(times[-1]))) return [(a, b) for a, b in intervals if b >= a] - # geometry + # geometry def depth_at(self, ts: float) -> Image | None: """Temporal join: aligned depth frame for a color timestamp.""" @@ -289,24 +560,67 @@ def depth_at(self, ts: float) -> Image | None: return None return depth + def registered_scan(self, scan: Observation[PointCloud2]) -> np.ndarray | None: + """A scan's points in the world frame, registered via tf when needed. + + Decode and registration run once per scan per run: accumulation + windows of neighboring frames overlap almost entirely, and every + query of a multi-label run shares every window. + """ + key = scan.ts + if key in self._scan_cache: + self._scan_cache.move_to_end(key) + return self._scan_cache[key] + points: np.ndarray | None = scan.data.as_numpy()[0] + frame = scan.data.frame_id + if frame != self.world_frame: + transform = ( + self.tf.get(self.world_frame, frame, scan.ts, self.tf_tolerance) + if self.tf is not None + else None + ) + if transform is None: + points = None + else: + matrix = transform.to_matrix() + points = points @ matrix[:3, :3].T + matrix[:3, 3] + self._scan_cache[key] = points + if len(self._scan_cache) > _SCAN_CACHE_MAX: + self._scan_cache.popitem(last=False) + return points + + def _cloud_quantum(self, points: np.ndarray) -> float | None: + if not self._quantum_known: + self._quantum = _lattice_quantum(points) + self._quantum_known = True + return self._quantum + def cloud_at(self, ts: float) -> PointCloud2 | None: """World-frame geometry at ts: the scans accumulated around it.""" if self._cloud_memo is not None and self._cloud_memo[0] == ts: return self._cloud_memo[1] from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - scans = list(self.cloud.after(ts - self.cloud_accum_s).before(ts + self.cloud_accum_s)) - if not scans: + scans = self.cloud.after(ts - self.cloud_accum_s).before(ts + self.cloud_accum_s) + parts = [p for p in (self.registered_scan(scan) for scan in scans) if p is not None] + if not parts: return None - merged: PointCloud2 - if len(scans) == 1: - merged = scans[0].data + if len(parts) == 1: + points = parts[0] else: - # A grid-quantized source (the Go2 occupancy stream) repeats the - # same cell in every snapshot it persists through; accumulation - # must not count one voxel once per snapshot. - points = np.unique(np.vstack([scan.data.as_numpy()[0] for scan in scans]), axis=0) - merged = PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=ts) + stacked = np.vstack(parts) + quantum = self._cloud_quantum(parts[0]) + if quantum is None: + points = stacked + else: + # A grid-quantized source repeats the same cell in every + # snapshot it persists through; accumulation must not count + # one voxel once per snapshot. Cell keys dedup in one pass. + cells = np.round(stacked / quantum).astype(np.int64) + _CELL_OFFSET + keys = (cells[:, 0] << 42) | (cells[:, 1] << 21) | cells[:, 2] + _, index = np.unique(keys, return_index=True) + points = stacked[index] + merged = PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=ts) self._cloud_memo = (ts, merged) return merged @@ -345,7 +659,7 @@ def backdrop(self, ts: float) -> PointCloud2 | None: color, depth, self.camera_info, depth_scale=0.001, depth_trunc=1.5 ).transform(-transform) - # predicates + # predicates def sees( self, @@ -424,7 +738,7 @@ def keyframes( return selected def default_localize_policy(self) -> LocalizePolicy: - return LocalizePolicy() if self.depth is not None else ROOM_LOCALIZE_POLICY + return ROOM_LOCALIZE_POLICY if self.mobile else LocalizePolicy() def default_inventory_policy(self) -> InventoryPolicy: - return InventoryPolicy() if self.depth is not None else ROOM_INVENTORY_POLICY + return ROOM_INVENTORY_POLICY if self.mobile else InventoryPolicy() diff --git a/dimos/perception/memory/tool_inventory.py b/dimos/perception/memory/tool_inventory.py index 5c0ababea2..850d3150e1 100644 --- a/dimos/perception/memory/tool_inventory.py +++ b/dimos/perception/memory/tool_inventory.py @@ -50,6 +50,7 @@ """ import argparse +import json from pathlib import Path import sys from typing import cast @@ -148,8 +149,9 @@ def at(ts: float) -> None: from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 scans = [ - obs.data.as_numpy()[0] + points for obs in rig.cloud.after(t0).before(t1).transform(throttle(2.0)) + if (points := rig.registered_scan(obs)) is not None ] if scans: merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) @@ -216,6 +218,11 @@ def main() -> int: "out", nargs="?", default=None, help="rerun recording to write; omitted writes none" ) parser.add_argument("--dataset", type=Path, help="memory recording database") + parser.add_argument("--manifest", type=Path, help="rig manifest json (default: .rig.json)") + parser.add_argument("--color", help="stream name override for the color role") + parser.add_argument("--depth", help="stream name override for the depth role") + parser.add_argument("--cloud", help="stream name override for the pointcloud role") + parser.add_argument("--odom", help="stream name override for the poses role") parser.add_argument( "--from", dest="start", type=float, default=0.0, help="start offset into the recording (s)" ) @@ -256,8 +263,19 @@ def main() -> int: "xarm6_worldbelief_realsense_d435i_stationery_calibrated/" "xarm6_worldbelief_20260729_203624_161992.db" ) + manifest = json.loads(args.manifest.read_text()) if args.manifest else None + overrides = { + role: name + for role, name in [ + ("color", args.color), + ("depth", args.depth), + ("cloud", args.cloud), + ("poses", args.odom), + ] + if name + } store = SqliteStore(path=dataset) - rig = Rig.from_store(store) + rig = Rig.from_store(store, manifest=manifest, overrides=overrides) lo, hi = rig.color.get_time_range() after = lo + args.start before = lo + args.start + args.duration if args.duration is not None else None diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index 8224e31101..ad94ec84cb 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -30,6 +30,7 @@ """ import argparse +import json from pathlib import Path import sys @@ -96,8 +97,9 @@ def at(ts: float) -> None: from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 scans = [ - obs.data.as_numpy()[0] + points for obs in rig.cloud.after(t0).before(t1).transform(throttle(2.0)) + if (points := rig.registered_scan(obs)) is not None ] if scans: merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) @@ -192,6 +194,11 @@ def main() -> int: help="one or more object queries, optionally followed by an out.rrd to write", ) parser.add_argument("--dataset", type=Path, help="memory recording database") + parser.add_argument("--manifest", type=Path, help="rig manifest json (default: .rig.json)") + parser.add_argument("--color", help="stream name override for the color role") + parser.add_argument("--depth", help="stream name override for the depth role") + parser.add_argument("--cloud", help="stream name override for the pointcloud role") + parser.add_argument("--odom", help="stream name override for the poses role") parser.add_argument( "--from", dest="start", type=float, default=0.0, help="start offset into the recording (s)" ) @@ -217,8 +224,19 @@ def main() -> int: "xarm6_worldbelief_realsense_d435i_stationery_calibrated/" "xarm6_worldbelief_20260729_203624_161992.db" ) + manifest = json.loads(args.manifest.read_text()) if args.manifest else None + overrides = { + role: name + for role, name in [ + ("color", args.color), + ("depth", args.depth), + ("cloud", args.cloud), + ("poses", args.odom), + ] + if name + } store = SqliteStore(path=dataset) - rig = Rig.from_store(store) + rig = Rig.from_store(store, manifest=manifest, overrides=overrides) lo, hi = rig.color.get_time_range() after = lo + args.start before = lo + args.start + args.duration if args.duration is not None else hi From 35c65c8c476204a0820b16e67f699dad83839580 Mon Sep 17 00:00:00 2001 From: bogwi Date: Sat, 22 Aug 2026 11:57:31 +0900 Subject: [PATCH 03/28] refactor localize to memory transforms --- dimos/perception/memory/dandetect.py | 5 +- dimos/perception/memory/localize.py | 330 ++++++++--------------- dimos/perception/memory/rig.py | 1 - dimos/perception/memory/tool_localize.py | 6 +- dimos/perception/memory/types.py | 6 +- 5 files changed, 126 insertions(+), 222 deletions(-) diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index c94578522a..81885e728c 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -29,7 +29,7 @@ from dimos.core.resource import Resource from dimos.memory.embed import EmbedImages -from dimos.memory.transform import throttle +from dimos.memory.transform import QualityWindow from dimos.perception.memory.inventory import DEFAULT_VOCABULARY, NamingVocabulary, inventory from dimos.perception.memory.localize import embed_index, localize from dimos.perception.memory.rig import Rig @@ -124,7 +124,8 @@ def embed( embedded: Stream[Any, Any] = store.stream("color_image_embedded", Image) pipeline = ( rig.color.live() - .transform(throttle(1.0 / rig.embed_hz)) + .filter(lambda obs: obs.data.brightness > 0.1) + .transform(QualityWindow(lambda img: img.sharpness, window=1.0 / rig.embed_hz)) .map(lambda obs: obs.derive(data=obs.data, pose=rig.index_pose(obs))) .filter(lambda obs: obs.pose is not None) .transform(EmbedImages(self.siglip, batch_size=1)) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 7ad30fd861..d01e7ea223 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -37,8 +37,9 @@ import numpy as np from dimos.memory.embed import EmbedImages -from dimos.memory.transform import throttle +from dimos.memory.transform import QualityWindow, peaks from dimos.perception.detection.identity import Identity, spatial +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Localization, LocalizePolicy, Support @@ -53,20 +54,15 @@ logger = setup_logger() -TIME_BANDS = 6 # stratify retrieval across the window so late scans always compete -BOXES_PER_FRAME = 4 -# Frames per OWLv2 forward. The forward's softmax transient is fp32 even -# under autocast (~580MB per image at 960px), so the profitable batch is -# whatever VRAM holds beyond the resident models - 1 on an 8GB card. -DETECT_BATCH = 1 -CONFIRM_FLOOR = 0.22 # geometric confirmation accept for re-detections -VERIFY_FRAMES = 20 - # A support candidate is an identity group: the member sightings of one # object. Everything a group reports is a plain function over its members. +def _similarity(obs: Any) -> float: + return float(obs.similarity) + + def _centroid(det: Detection3DPC) -> np.ndarray: centroid: np.ndarray = np.asarray(det.pointcloud.pointcloud.points).mean(axis=0) return centroid @@ -99,12 +95,6 @@ def _n_views(members: list[Detection3DPC]) -> int: return len({tuple(np.round(_camera_position(d), 2)) for d in members}) -def _group_extent(members: list[Detection3DPC]) -> np.ndarray: - points = np.concatenate([np.asarray(d.pointcloud.pointcloud.points) for d in members]) - extent: np.ndarray = points.max(axis=0) - points.min(axis=0) - return extent - - @dataclass class LocalizeTrace: """Intermediate artifacts collected for rendering; filled when passed in.""" @@ -151,83 +141,23 @@ def _axes_observed(members: list[Detection3DPC], center: np.ndarray) -> tuple[bo return tuple(bool((np.abs(dirs[:, i]) > 0.3).any()) for i in range(3)) # type: ignore[return-value] -class _DetectionCache: - """One OWLv2 + EdgeTAM + lift pass per unique frame, shared across queries and clusters.""" - - def __init__( - self, detector: Any, segmenter: Any, rig: Rig, queries: list[str], floor: float - ) -> None: - self.detector = detector - self.segmenter = segmenter - self.rig = rig - self.queries = queries - self.floor = floor - self._cache: dict[float, ImageDetections2D] = {} - self._lifted: dict[float, list[Detection3DPC]] = {} - - def _ingest(self, image: Any, detections: ImageDetections2D) -> None: - ranked = sorted(detections.detections, key=lambda d: -d.confidence) - kept = [ - det - for label in self.queries - for det in [d for d in ranked if d.name == label][:BOXES_PER_FRAME] - ] - for i, det in enumerate(kept): - det.track_id = i - cached = ImageDetections2D(image, kept) - if len(cached): - cached = self.segmenter.segment(cached) - self._cache[image.ts] = cached - - def prefetch(self, images: list[Any]) -> None: - """Detect and segment every uncached frame, batching the OWLv2 forwards.""" - todo: dict[float, Any] = {} - for image in images: - if image.ts not in self._cache and image.ts not in todo: - todo[image.ts] = image - pending = list(todo.values()) - for start in range(0, len(pending), DETECT_BATCH): - chunk = pending[start : start + DETECT_BATCH] - for image, detections in zip( - chunk, - self.detector.query_detections_batch(chunk, self.queries, threshold=self.floor), - strict=True, - ): - self._ingest(image, detections) - - def detect(self, image: Any, query: str) -> ImageDetections2D: - if image.ts not in self._cache: - self._ingest( - image, self.detector.query_detections(image, self.queries, threshold=self.floor) - ) - return self._cache[image.ts].filter(lambda d: d.name == query) - - def lift(self, ts: float) -> list[Detection3DPC]: - """All of a frame's detections lifted once, whatever their label.""" - if ts not in self._lifted: - lifted = self.rig.lift(self._cache[ts]) - self._lifted[ts] = list(lifted) if lifted is not None else [] - return self._lifted[ts] - - def _lift( detections: ImageDetections2D, - cache: _DetectionCache, rig: Rig, policy: LocalizePolicy, plane: Any | None = None, -) -> list[tuple[Detection3DPC, np.ndarray]]: - """Gate a query's share of the frame's shared lift; (detection3d, camera_position) pairs.""" +) -> list[Detection3DPC]: + """Gate a frame's lifted detections.""" pose = rig.camera_pose(detections.ts) if pose is None: return [] camera = np.array([pose.position.x, pose.position.y, pose.position.z]) + lifted = rig.lift(detections) + if lifted is None: + return [] - track_ids = {det.track_id for det in detections} - valid: list[tuple[Detection3DPC, np.ndarray]] = [] - for det3d in cache.lift(detections.ts): - if det3d.track_id not in track_ids: - continue + valid: list[Detection3DPC] = [] + for det3d in lifted: points = np.asarray(det3d.pointcloud.pointcloud.points) if len(points) < policy.min_depth_points: continue @@ -243,7 +173,7 @@ def _lift( high = float(np.quantile(heights, 0.95)) if low > policy.surface_patch_min_drop_m and high < policy.surface_patch_max_rise_m: continue - valid.append((det3d, camera)) + valid.append(det3d) return valid @@ -264,7 +194,8 @@ def embed_index( posed = ( rig.color.after(t0) .before(t1) - .transform(throttle(1.0 / rig.embed_hz)) + .filter(lambda obs: obs.data.brightness > 0.1) + .transform(QualityWindow(lambda img: img.sharpness, window=1.0 / rig.embed_hz)) .map(lambda obs: obs.derive(data=obs.data, pose=rig.index_pose(obs))) .filter(lambda obs: obs.pose is not None) ) @@ -273,44 +204,6 @@ def embed_index( return embedded -def _retrieve(index: Stream[Any, Any], rig: Rig, query_embedding: Any, budget: int) -> list[Any]: - """Top still frames by text similarity, stratified over time bands. - - Stratification is what keeps latest-pose semantics honest at the - candidate stage: the top frames of the whole window may all be early, - and a support that only exists late must still get a detection pass. - """ - ranked = [ - obs - for obs in index.search(query_embedding, k=max(index.count(), 1)) - if rig.camera_still(obs.ts) - ] - if not ranked: - return [] - - t0, t1 = index.get_time_range() - bands = max(1, min(TIME_BANDS, int((t1 - t0) / 20))) - per_band = max(1, budget // bands) - span = (t1 - t0) / bands - selected: list[Any] = [] - chosen: set[float] = set() - for band in range(bands): - band_lo = t0 + band * span - band_hi = band_lo + span - in_band = [obs for obs in ranked if band_lo <= obs.ts < band_hi] - for obs in in_band[:per_band]: - if obs.ts not in chosen: - chosen.add(obs.ts) - selected.append(obs) - for obs in ranked: # fill remaining budget by global rank - if len(selected) >= budget: - break - if obs.ts not in chosen: - chosen.add(obs.ts) - selected.append(obs) - return selected - - def localize( store: Any, query: str | list[str], @@ -333,9 +226,10 @@ def localize( coexisting candidates is returned with ``ambiguity_margin`` below ``refusal_margin`` - a flagged hit, never a silent guess. - A list *query* runs every label through one shared detection cache - - OWLv2 takes the whole list per frame - and returns one result per label, - in input order; ``trace`` then takes a list of the same length. + A list *query* shares one detection pass: every label's semantic peaks + select frames, and each unique frame is scored against every label in a + single OWLv2 forward, segmented and lifted once. One result per label, + in input order. ``trace`` then takes a list of the same length. The index, the rig and the three models belong to the caller: nothing here is loaded or stopped, so one process can call this repeatedly on @@ -351,124 +245,138 @@ def localize( traces: list[LocalizeTrace | None] = ( list(trace) if isinstance(trace, list) else [trace] * len(queries) ) - cache = _DetectionCache(detector, segmenter, rig, queries, policy.candidate_floor) + + peaks_per_label: list[Stream[Any, Any]] = [] + for q in queries: + query_embedding = siglip.embed_text(q) + label_peaks: Stream[Any, Any] = ( + index.search(query_embedding) + .order_by("ts") + .transform(peaks(key=_similarity, distance=1.0)) + .materialize() + ) + logger.info( + f"localize {q!r}: {label_peaks.count()} semantic peaks of {index.count()} embedded" + ) + peaks_per_label.append(label_peaks) + + frames: dict[float, Any] = {} + expanded: set[float] = set() + for label_peaks in peaks_per_label: + for peak in label_peaks: + frames.setdefault(peak.ts, peak) + if peak.ts in expanded: + continue + expanded.add(peak.ts) + nearby: Stream[Any, Any] = index.near( + peak.pose_stamped, radius=policy.verify_radius_m + ).transform(QualityWindow(lambda img: img.sharpness, window=0.5)) + for obs in nearby: + frames.setdefault(obs.ts, obs) + ordered = sorted(frames.values(), key=lambda obs: obs.ts) + logger.info(f"detection: {len(ordered)} candidate frames for {len(queries)} labels") + + if not ordered: + results: list[Localization | None] = [None] * len(queries) + return None if isinstance(query, str) else results + + from dimos.perception.memory.support_plane import fit_support_plane + + plane = fit_support_plane(rig, ordered) + identities = [Identity(is_same=spatial(policy.cluster_radius_m)) for _ in queries] + ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) + + for obs in ordered: + boxes, rows = detector.query_score_rows(obs.data, queries, threshold=policy.candidate_floor) + candidates: list[Detection2DBBox] = [] + for box, row in zip(boxes, rows, strict=True): + bbox = (float(box[0]), float(box[1]), float(box[2]), float(box[3])) + for j, score in enumerate(row): + if score < policy.candidate_floor: + continue + det = Detection2DBBox( + bbox=bbox, + track_id=len(candidates), + class_id=j, + confidence=float(score), + name=queries[j], + ts=obs.data.ts, + image=obs.data, + ) + if det.is_valid() and det.bbox_2d_volume() > 3000: + candidates.append(det) + if not candidates: + continue + + frame = segmenter.segment(ImageDetections2D(image=obs.data, detections=candidates)) + lifted = _lift(frame, rig, policy, plane) + grounded = {det3d.track_id for det3d in lifted} + for det2d in frame: + j = det2d.class_id + best = ungrounded[j] + if det2d.track_id not in grounded and (best is None or det2d.confidence > best[0]): + ungrounded[j] = (det2d.confidence, det2d.ts) + for det3d in lifted: + label_trace = traces[det3d.class_id] + if label_trace is not None: + label_trace.matched.append((det3d.ts, det3d)) + identities[det3d.class_id].add(det3d) + for j, label_trace in enumerate(traces): + if label_trace is None: + continue + label_dets = [det for det in frame if det.class_id == j] + if label_dets: + label_trace.detection_frames.append( + obs.derive(data=ImageDetections2D(image=obs.data, detections=label_dets)) + ) + results = [ - _localize_one( + _finalize( q, - index=index, - siglip=siglip, - cache=cache, + identity=identities[j], + ungrounded_best=ungrounded[j], rig=rig, require_pose=require_pose, policy=policy, cloud_mode=cloud_mode, - trace=t, + trace=traces[j], ) - for q, t in zip(queries, traces, strict=True) + for j, q in enumerate(queries) ] return results[0] if isinstance(query, str) else results -def _localize_one( +def _finalize( query: str, *, - index: Stream[Any, Any], - siglip: SigLIPModel, - cache: _DetectionCache, + identity: Identity, + ungrounded_best: tuple[float, float] | None, rig: Rig, require_pose: bool, policy: LocalizePolicy, cloud_mode: str, trace: LocalizeTrace | None, ) -> Localization | None: - # Pass 1 - SigLIP: rank the indexed frames by the query. - query_embedding = siglip.embed_text(query) - frames = _retrieve(index, rig, query_embedding, policy.retrieval_frames) - logger.info(f"localize '{query}': {len(frames)} candidate frames of {index.count()} embedded") - if not frames: - return None - - from dimos.perception.memory.support_plane import fit_support_plane - - plane = fit_support_plane(rig, frames) - cache.prefetch([obs.data for obs in frames]) - - # Pass 2 - OWLv2 + EdgeTAM: detect, segment, lift, verify. Support - # candidates are identity groups: sightings joined by the spatial - # is_same strategy at the policy's cluster radius. - identity = Identity(is_same=spatial(policy.cluster_radius_m)) - ungrounded_best: tuple[float, float] | None = None # (score, ts) - processed: set[float] = set() - - def _absorb(frame_obs: Any, is_verify: bool) -> None: - nonlocal ungrounded_best - if frame_obs.ts in processed: - return - processed.add(frame_obs.ts) - detections = cache.detect(frame_obs.data, query) - if not len(detections): - return - if trace is not None and not is_verify: - trace.detection_frames.append(frame_obs.derive(data=detections)) - lifted = _lift(detections, cache, rig, policy, plane) - for det2d in detections: - if not any(d.track_id == det2d.track_id for d, _ in lifted): - best = (det2d.confidence, det2d.ts) - if ungrounded_best is None or best[0] > ungrounded_best[0]: - ungrounded_best = best - for det3d, _camera in lifted: - if trace is not None: - (trace.verified if is_verify else trace.matched).append((det3d.ts, det3d)) - identity.add(det3d) - - for frame_obs in frames: - _absorb(frame_obs, is_verify=False) - logger.info(f"detection: {len(identity.groups)} support candidates") - - # Cross-view verification: a support seen from one pose only is - # unconfirmed. Frames whose camera could observe the support are found - # geometrically (near + sees with occlusion), then re-detected. - for members in sorted(identity.groups, key=_max_score, reverse=True)[:4]: - center = _group_center(members) - predicate = rig.sees( - center, - extent=np.minimum(_group_extent(members), 0.4), - # A large object overflows close-up frames; a third of its box in - # view is still a usable re-detection pass, and those close-ups - # are exactly the distinct viewpoints verification needs. - min_fraction=0.35, - max_range=policy.verify_radius_m, - ) - observing = [ - obs - for obs in index.near(center, radius=policy.verify_radius_m) - if obs.ts not in processed and rig.camera_still(obs.ts) and predicate(obs) - ] - if len(observing) > VERIFY_FRAMES: - # Even spread that always includes the endpoints: dropping the - # latest seeing frames would bias the latest-pose answer early. - picks = np.unique(np.linspace(0, len(observing) - 1, VERIFY_FRAMES).astype(int)) - observing = [observing[i] for i in picks] - cache.prefetch([obs.data for obs in observing]) - for frame_obs in observing: - _absorb(frame_obs, is_verify=True) - verified = [ members for members in identity.groups if _max_score(members) >= policy.accept_score and _n_views(members) >= policy.min_views ] logger.info( - "verification: " + f"verification {query!r}: " + ", ".join( f"score={_max_score(g):.2f} views={_n_views(g)} obs={len(g)}" for g in identity.groups ) ) + if trace is not None: + for members in verified: + for det3d in members: + trace.verified.append((det3d.ts, det3d)) if not verified: if ungrounded_best is not None and ungrounded_best[0] >= policy.accept_score: if require_pose: - logger.info("best candidate has no valid depth and require_pose is set") + logger.info(f"{query!r}: best candidate has no valid depth and require_pose is set") return None score, ts = ungrounded_best return Localization( diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index a3554d8935..6af0a8092b 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -100,7 +100,6 @@ ROOM_LOCALIZE_POLICY = LocalizePolicy( candidate_floor=0.18, accept_score=0.32, - retrieval_frames=20, cluster_radius_m=0.30, min_depth_points=30, max_object_extent_m=2.0, diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index ad94ec84cb..fbab1c5e85 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -20,8 +20,8 @@ The recording's shape decides the rig: an xArm-style store lifts through aligned depth and tf, a mobile-robot store (Go2/G1 replay) lifts through registered lidar and stamped poses. Queries share one model load and one -.rrd; with --multi they go to localize() as one list, sharing a single -detection pass per frame. +.rrd; with --multi they go to localize() as one list and share one +detection pass over the union of the labels' candidate frames. Exit code 0 with a printed position per verified hit; exit code 1 when no query is verified, with "no verified detection of ..." per miss - the honest answer that the object @@ -211,7 +211,7 @@ def main() -> int: parser.add_argument( "--multi", action="store_true", - help="pass all queries to localize() as one list (one shared detection pass per frame)", + help="pass all queries to localize() as one list sharing one detection pass", ) args = parser.parse_args() diff --git a/dimos/perception/memory/types.py b/dimos/perception/memory/types.py index 552a4199fe..44cc3f8c2f 100644 --- a/dimos/perception/memory/types.py +++ b/dimos/perception/memory/types.py @@ -115,10 +115,6 @@ class LocalizePolicy: accept_score: float = 0.40 refusal_margin: float = 0.15 min_views: int = 2 # a support seen from one pose only is unconfirmed - # Frames retrieved per query for the detection pass. Wrist-camera frames - # each cover most of the workspace; room-scale sweeps dilute the - # embedding signal across viewpoints and need a larger budget. - retrieval_frames: int = 12 cluster_radius_m: float = 0.08 # observations within this are the same support min_depth_points: int = 60 @@ -128,7 +124,7 @@ class LocalizePolicy: # object: every real object rises above the plane, a surface patch does not. surface_patch_max_rise_m: float = 0.003 surface_patch_min_drop_m: float = -0.02 - # Cross-view verification looks for re-detections this far from a support. + # Images gathered around each semantic peak for the detection pass. verify_radius_m: float = 1.6 From 37d22aefc977d71f9779288b7dc50e2c758c4e52 Mon Sep 17 00:00:00 2001 From: bogwi Date: Sat, 22 Aug 2026 14:10:26 +0900 Subject: [PATCH 04/28] impl: (a) All instances per query. (b) Global point cloud per instance --- dimos/perception/memory/dandetect.py | 2 +- dimos/perception/memory/localize.py | 224 ++++++++++++----------- dimos/perception/memory/tool_localize.py | 68 +++---- dimos/perception/memory/types.py | 7 +- 4 files changed, 160 insertions(+), 141 deletions(-) diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index 81885e728c..0553d7533d 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -141,7 +141,7 @@ def localize( *, index: Stream[Any, Any], **kwargs: Any, - ) -> Localization | list[Localization | None] | None: + ) -> list[Localization] | list[list[Localization]]: """:func:`localize` on this resource's models.""" return localize( store, diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index d01e7ea223..e0287a5462 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -20,12 +20,14 @@ aligned depth stream or a registered pointcloud stream. Two algorithm rules distinguish it from a best-crop search: -* **Latest-pose semantics.** Among verified observations of the chosen - support, the greatest timestamp wins. The answer is "where is it now", - never "where did it match best". +* **Latest-pose semantics.** Every verified instance is returned, + latest-seen first, and each instance's position follows its latest + sighting: "where is it now", never "where did it match best". The + instance's cloud is the union of every viewpoint that saw it. * **Calibrated refusal.** Every stage carries a score and the answer can be - ``None``: no accept-level detection, no multi-view confirmation, or an - ambiguity between coexisting candidates below the refusal margin. + empty: no accept-level detection, no multi-view confirmation. Coexisting + same-label instances below the refusal margin are flagged, never merged + or silently dropped. """ from __future__ import annotations @@ -102,7 +104,7 @@ class LocalizeTrace: detection_frames: list[Any] = field(default_factory=list) # Observation[ImageDetections2D] matched: list[tuple[float, Detection3DPC]] = field(default_factory=list) verified: list[tuple[float, Detection3DPC]] = field(default_factory=list) - answer: Detection3DPC | None = None + answers: list[Detection3DPC] = field(default_factory=list) # merged union per instance backdrop_ts: float | None = None @@ -215,21 +217,22 @@ def localize( rig: Rig | None = None, require_pose: bool = True, policy: LocalizePolicy | None = None, - cloud_mode: str = "latest_visible", trace: LocalizeTrace | list[LocalizeTrace] | None = None, -) -> Localization | list[Localization | None] | None: - """Latest unambiguous 3D localization of *query*, or ``None``. +) -> list[Localization] | list[list[Localization]]: + """Every verified 3D instance of *query*, latest-seen first. - ``None`` is a first-class answer: nothing reached the accept score, no - support was confirmed from a second viewpoint, or the best candidate had - no valid depth and ``require_pose`` holds. An ambiguity between - coexisting candidates is returned with ``ambiguity_margin`` below - ``refusal_margin`` - a flagged hit, never a silent guess. + Each instance's ``point_cloud`` is the union of every viewpoint that saw + it, and its position follows the latest sighting. An empty list is a + first-class answer: nothing reached the accept score, no support was + confirmed from a second viewpoint, or the best candidate had no valid + depth and ``require_pose`` holds. Coexisting instances of one label are + all returned; each carries ``ambiguity_margin`` against its rivals and + is flagged below ``refusal_margin`` - never a silent guess. A list *query* shares one detection pass: every label's semantic peaks select frames, and each unique frame is scored against every label in a - single OWLv2 forward, segmented and lifted once. One result per label, - in input order. ``trace`` then takes a list of the same length. + single OWLv2 forward, segmented and lifted once. One instance list per + label, in input order. ``trace`` then takes a list of the same length. The index, the rig and the three models belong to the caller: nothing here is loaded or stopped, so one process can call this repeatedly on @@ -277,8 +280,8 @@ def localize( logger.info(f"detection: {len(ordered)} candidate frames for {len(queries)} labels") if not ordered: - results: list[Localization | None] = [None] * len(queries) - return None if isinstance(query, str) else results + empty: list[list[Localization]] = [[] for _ in queries] + return [] if isinstance(query, str) else empty from dimos.perception.memory.support_plane import fit_support_plane @@ -338,7 +341,6 @@ def localize( rig=rig, require_pose=require_pose, policy=policy, - cloud_mode=cloud_mode, trace=traces[j], ) for j, q in enumerate(queries) @@ -354,12 +356,11 @@ def _finalize( rig: Rig, require_pose: bool, policy: LocalizePolicy, - cloud_mode: str, trace: LocalizeTrace | None, -) -> Localization | None: +) -> list[Localization]: verified = [ - members - for members in identity.groups + (merged, members) + for merged, members in zip(identity.merged, identity.groups, strict=True) if _max_score(members) >= policy.accept_score and _n_views(members) >= policy.min_views ] logger.info( @@ -369,7 +370,7 @@ def _finalize( ) ) if trace is not None: - for members in verified: + for _merged, members in verified: for det3d in members: trace.verified.append((det3d.ts, det3d)) @@ -377,88 +378,99 @@ def _finalize( if ungrounded_best is not None and ungrounded_best[0] >= policy.accept_score: if require_pose: logger.info(f"{query!r}: best candidate has no valid depth and require_pose is set") - return None + return [] score, ts = ungrounded_best - return Localization( - instance_id="query-0", - semantic_score=score, - identity_score=0.0, - ambiguity_margin=1.0, - position_world_xyz=None, - orientation_world_xyzw=None, - frame_id=rig.world_frame, - support=None, - pose_timestamp=ts, - geometry_timestamp=ts, - last_seen_timestamp=ts, - point_cloud=None, - cloud_mode=cloud_mode, - coverage=0.0, - n_views=1, - reason="no_valid_depth", + return [ + Localization( + instance_id="query-0", + semantic_score=score, + identity_score=0.0, + ambiguity_margin=1.0, + position_world_xyz=None, + orientation_world_xyzw=None, + frame_id=rig.world_frame, + support=None, + pose_timestamp=ts, + geometry_timestamp=ts, + last_seen_timestamp=ts, + point_cloud=None, + coverage=0.0, + n_views=1, + reason="no_valid_depth", + ) + ] + return [] + + verified.sort(key=lambda pair: _latest(pair[1]).ts, reverse=True) + instances: list[Localization] = [] + for k, (merged, members) in enumerate(verified): + m_lo, m_hi = _interval(members) + rival_scores = [ + _max_score(others) + for _m, others in verified + if others is not members + and not (_interval(others)[1] < m_lo or _interval(others)[0] > m_hi) # coexisting + ] + margin = _max_score(members) - max(rival_scores) if rival_scores else 1.0 + reason = ( + "ambiguous_between_coexisting_candidates" if margin < policy.refusal_margin else None + ) + + latest = _latest(members) + union = merged.pointcloud + points = np.asarray(union.pointcloud.points) + aabb_min, aabb_max = points.min(axis=0), points.max(axis=0) + try: + orientation = _quaternion_from_matrix( + np.asarray(latest.pointcloud.oriented_bounding_box.R) ) - return None - - winner = max(verified, key=lambda g: _latest(g).ts) - w_lo, w_hi = _interval(winner) - rival_scores = [ - _max_score(g) - for g in verified - if g is not winner - and not (_interval(g)[1] < w_lo or _interval(g)[0] > w_hi) # coexisting in time - ] - margin = _max_score(winner) - max(rival_scores) if rival_scores else 1.0 - reason = "ambiguous_between_coexisting_candidates" if margin < policy.refusal_margin else None - - latest = _latest(winner) - points = np.asarray(latest.pointcloud.pointcloud.points) - aabb_min, aabb_max = points.min(axis=0), points.max(axis=0) - try: - orientation = _quaternion_from_matrix(np.asarray(latest.pointcloud.oriented_bounding_box.R)) - except Exception: - orientation = (0.0, 0.0, 0.0, 1.0) - center = (aabb_min + aabb_max) / 2 - extent = np.maximum(aabb_max - aabb_min, 0.005) - sigma = ( - np.stack([_centroid(d) for d in winner]).std(axis=0) - if len(winner) > 1 - else np.full(3, 0.01) - ) - winner_center = _group_center(winner) - support = Support( - center_xyz=(float(center[0]), float(center[1]), float(center[2])), - extent_xyz_m=(float(extent[0]), float(extent[1]), float(extent[2])), - orientation_xyzw=(0.0, 0.0, 0.0, 1.0), - sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), - coverage=_azimuth_coverage(winner, winner_center), - axes_observed=_axes_observed(winner, winner_center), - frame_id=rig.world_frame, - ) + except Exception: + orientation = (0.0, 0.0, 0.0, 1.0) + center = (aabb_min + aabb_max) / 2 + extent = np.maximum(aabb_max - aabb_min, 0.005) + sigma = ( + np.stack([_centroid(d) for d in members]).std(axis=0) + if len(members) > 1 + else np.full(3, 0.01) + ) + group_center = _group_center(members) + support = Support( + center_xyz=(float(center[0]), float(center[1]), float(center[2])), + extent_xyz_m=(float(extent[0]), float(extent[1]), float(extent[2])), + orientation_xyzw=(0.0, 0.0, 0.0, 1.0), + sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), + coverage=_azimuth_coverage(members, group_center), + axes_observed=_axes_observed(members, group_center), + frame_id=rig.world_frame, + ) - if trace is not None: - trace.answer = latest - trace.backdrop_ts = latest.ts - - latest_centroid = _centroid(latest) - return Localization( - instance_id="query-0", - semantic_score=_max_score(winner), - identity_score=min(1.0, _n_views(winner) / 4.0), - ambiguity_margin=margin, - position_world_xyz=( - float(latest_centroid[0]), - float(latest_centroid[1]), - float(latest_centroid[2]), - ), - orientation_world_xyzw=orientation, - frame_id=rig.world_frame, - support=support, - pose_timestamp=latest.ts, - geometry_timestamp=latest.ts, - last_seen_timestamp=latest.ts, - point_cloud=latest.pointcloud, - cloud_mode=cloud_mode, - coverage=support.coverage, - n_views=_n_views(winner), - reason=reason, - ) + if trace is not None: + trace.answers.append(merged) + if k == 0: + trace.backdrop_ts = latest.ts + + latest_centroid = _centroid(latest) + instances.append( + Localization( + instance_id=f"query-{k}", + semantic_score=_max_score(members), + identity_score=min(1.0, _n_views(members) / 4.0), + ambiguity_margin=margin, + position_world_xyz=( + float(latest_centroid[0]), + float(latest_centroid[1]), + float(latest_centroid[2]), + ), + orientation_world_xyzw=orientation, + frame_id=rig.world_frame, + support=support, + pose_timestamp=latest.ts, + geometry_timestamp=latest.ts, + last_seen_timestamp=latest.ts, + point_cloud=union, + coverage=support.coverage, + n_views=_n_views(members), + reason=reason, + ) + ) + return instances diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index fbab1c5e85..c5ff11e3c5 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -22,11 +22,13 @@ registered lidar and stamped poses. Queries share one model load and one .rrd; with --multi they go to localize() as one list and share one detection pass over the union of the labels' candidate frames. -Exit code 0 with a printed -position per verified hit; exit code 1 when no query is verified, with -"no verified detection of ..." per miss - the honest answer that the object -is not there. An ambiguous hit (identical twins in view) is printed with its -ambiguity margin flagged. +Every query prints one +line per verified instance - all the coke bottles, not one winner - each +with the union cloud of every viewpoint that saw it. Exit code 0 when any +query verified; exit code 1 when none did, with "no verified detection +of ..." per miss - the honest answer that the object is not there. An +ambiguous instance (identical twins in view) is printed with its ambiguity +margin flagged. """ import argparse @@ -51,8 +53,8 @@ def render( Entity contract (the acceptance color cheat sheet): ``map`` backdrop, then one subtree per query - ``detections//matched/*`` green, - ``detections//verified/*`` red, ``detections//answer`` - always blue. + ``detections//verified/*`` red, ``detections//answer/`` + always blue, one per verified instance. """ import rerun as rr import rerun.blueprint as rrb @@ -149,39 +151,41 @@ def at(ts: float) -> None: det.pointcloud.to_rerun(voxel_size=point_size, colors=rgb), ) - # the answer: always blue, whatever the query - if trace.answer is not None: - at(trace.answer.ts) + # the answers: one blue union cloud per verified instance + for i, answer in enumerate(trace.answers): + at(answer.ts) rr.log( - f"{root}/answer", - trace.answer.pointcloud.to_rerun(voxel_size=point_size, colors=BLUE), + f"{root}/answer/{i}", + answer.pointcloud.to_rerun(voxel_size=point_size, colors=BLUE), ) -def report(query: str, hit: Localization | None, lo: float) -> bool: - """Print one query's outcome; True when it counts as a hit.""" - if hit is None: +def report(query: str, hits: list[Localization], lo: float) -> bool: + """Print one query's instances; True when any is verified.""" + if not hits: print(f"no verified detection of {query!r}") return False - offset = hit.pose_timestamp - lo - if hit.position_world_xyz is None: - print( - f"hit {query!r} without pose: reason={hit.reason} " - f"score={hit.semantic_score:.2f} ts_offset={offset:.1f}s" - ) - return True - x, y, z = hit.position_world_xyz - cloud_points = len(hit.point_cloud) if hit.point_cloud is not None else 0 - print( - f"hit {query!r}: position=({x:.3f}, {y:.3f}, {z:.3f}) frame={hit.frame_id} " - f"ts_offset={offset:.1f}s points={cloud_points} views={hit.n_views} " - f"score={hit.semantic_score:.2f} margin={hit.ambiguity_margin:.2f}" - ) - if hit.ambiguity_margin < REFUSAL_MARGIN: + for hit in hits: + offset = hit.pose_timestamp - lo + if hit.position_world_xyz is None: + print( + f"hit {query!r} without pose: reason={hit.reason} " + f"score={hit.semantic_score:.2f} ts_offset={offset:.1f}s" + ) + continue + x, y, z = hit.position_world_xyz + cloud_points = len(hit.point_cloud) if hit.point_cloud is not None else 0 print( - f"ambiguity: margin {hit.ambiguity_margin:.2f} below refusal threshold " - f"{REFUSAL_MARGIN:.2f} - multiple coexisting matches, this pick is flagged" + f"hit {query!r} [{hit.instance_id}]: position=({x:.3f}, {y:.3f}, {z:.3f}) " + f"frame={hit.frame_id} ts_offset={offset:.1f}s points={cloud_points} " + f"views={hit.n_views} score={hit.semantic_score:.2f} " + f"margin={hit.ambiguity_margin:.2f}" ) + if hit.ambiguity_margin < REFUSAL_MARGIN: + print( + f"ambiguity: margin {hit.ambiguity_margin:.2f} below refusal threshold " + f"{REFUSAL_MARGIN:.2f} - coexisting matches, this instance is flagged" + ) return True diff --git a/dimos/perception/memory/types.py b/dimos/perception/memory/types.py index 44cc3f8c2f..fbbeaadc89 100644 --- a/dimos/perception/memory/types.py +++ b/dimos/perception/memory/types.py @@ -78,7 +78,11 @@ class Instance: @dataclass class Localization: - """Latest unambiguous localization of a queried object.""" + """One verified instance of a queried object. + + ``point_cloud`` is the union of every viewpoint that saw the instance; + position and timestamps follow the latest sighting. + """ instance_id: str semantic_score: float @@ -95,7 +99,6 @@ class Localization: last_seen_timestamp: float point_cloud: PointCloud2 | None - cloud_mode: str coverage: float n_views: int reason: str | None = None From eef787dfb1fd277ac6af0a423fcb98378a5836ce Mon Sep 17 00:00:00 2001 From: bogwi Date: Mon, 24 Aug 2026 09:40:08 +0900 Subject: [PATCH 05/28] filter before decode in _vector_search --- dimos/memory/backend.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/dimos/memory/backend.py b/dimos/memory/backend.py index 81509ede79..1194ba1ca0 100644 --- a/dimos/memory/backend.py +++ b/dimos/memory/backend.py @@ -207,13 +207,16 @@ def _vector_search(self, query: StreamQuery) -> Iterator[Observation[T]]: ranked: list[Observation[T]] = [] for obs_id, sim in hits: match = obs_by_id.get(obs_id) - if match is not None: - ranked.append( - match.derive(data=match.data, embedding=query.search_vec, similarity=sim) - ) + if match is None: + continue + if not all(f.matches(match) for f in query.filters): + continue + ranked.append( + match.derive(data=match.data, embedding=query.search_vec, similarity=sim) + ) - # Apply remaining query ops (skip vector search) - rest = replace(query, search_vec=None, search_k=None) + # Apply remaining query ops (filters already applied; skip vector search) + rest = replace(query, search_vec=None, search_k=None, filters=()) yield from rest.apply(iter(ranked)) def _iterate_live( From fe42124635e9337416f6f02794df2eb905f5293f Mon Sep 17 00:00:00 2001 From: bogwi Date: Mon, 24 Aug 2026 10:55:43 +0900 Subject: [PATCH 06/28] support plane cached per pose cell --- dimos/perception/memory/localize.py | 10 +++++++++- dimos/perception/memory/rig.py | 4 ++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index e0287a5462..4653b89490 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -285,7 +285,15 @@ def localize( from dimos.perception.memory.support_plane import fit_support_plane - plane = fit_support_plane(rig, ordered) + anchors = [peak.pose_tuple for label_peaks in peaks_per_label for peak in label_peaks] + mx = sum(t[0] for t in anchors) / len(anchors) + my = sum(t[1] for t in anchors) / len(anchors) + cell = (round(mx / 2.0), round(my / 2.0)) + plane = rig._plane_cache.get(cell) + if plane is None: + plane = fit_support_plane(rig, ordered) + if plane is not None: + rig._plane_cache[cell] = plane identities = [Identity(is_same=spatial(policy.cluster_radius_m)) for _ in queries] ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 6af0a8092b..f70d24a174 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -65,6 +65,7 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D + from dimos.perception.memory.support_plane import SupportPlane from dimos.protocol.tf.tf import TFLookup logger = setup_logger() @@ -258,6 +259,9 @@ class Rig: _scan_cache: OrderedDict[float, np.ndarray | None] = field( default_factory=OrderedDict, repr=False, init=False ) + _plane_cache: dict[tuple[int, int], SupportPlane] = field( + default_factory=dict, repr=False, init=False + ) _quantum: float | None = field(default=None, repr=False, init=False) _quantum_known: bool = field(default=False, repr=False, init=False) From 840463a2d626302ac4165a56ce333e9325174f50 Mon Sep 17 00:00:00 2001 From: bogwi Date: Mon, 24 Aug 2026 14:23:21 +0900 Subject: [PATCH 07/28] OWLv2 score-row cache keyed by (frame ts, label), batched misses --- dimos/perception/detection/detectors/owlv2.py | 113 ++++++++++++++++++ dimos/perception/memory/localize.py | 40 +++++-- 2 files changed, 143 insertions(+), 10 deletions(-) diff --git a/dimos/perception/detection/detectors/owlv2.py b/dimos/perception/detection/detectors/owlv2.py index 3139424069..7e7b306502 100644 --- a/dimos/perception/detection/detectors/owlv2.py +++ b/dimos/perception/detection/detectors/owlv2.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections import OrderedDict from functools import cached_property import numpy as np @@ -27,6 +28,9 @@ from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D +# ~3.8 MB of GPU memory per cached frame +_FEATURE_CACHE_MAX = 128 + class Owlv2Config(HuggingFaceModelConfig): model_name: str = "google/owlv2-base-patch16-ensemble" @@ -47,6 +51,22 @@ class Owlv2Detector(HuggingFaceModel): config: Owlv2Config + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + # image-tower forwards run so far; instrumentation reads deltas + self.forwards = 0 + # (frame ts, label, floor) to that label's floor-filtered + # (boxes_kx4, scores_k); filled and read by localize + self.score_cache: OrderedDict[tuple[float, str, float], tuple[np.ndarray, np.ndarray]] = ( + OrderedDict() + ) + # frame ts to text-independent tower outputs; any query set scores + # against a cached frame without rerunning the tower + self._features: OrderedDict[ + float, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = OrderedDict() + self._text_embeds: dict[str, torch.Tensor] = {} + @cached_property def _model(self): # type: ignore[no-untyped-def] from transformers import Owlv2ForObjectDetection @@ -97,6 +117,7 @@ def query_detections_batch( across the batch, which is what makes many-frame sweeps affordable; results are per-image, in input order. """ + self.forwards += len(images) pils = [PILImage.fromarray(image.to_rgb().data) for image in images] with torch.inference_mode(), self._autocast(): inputs = self._processor( @@ -147,6 +168,7 @@ def query_score_rows( queries and refuse. Returns pixel ``(x1, y1, x2, y2)`` boxes and their score rows. """ + self.forwards += 1 pil = PILImage.fromarray(image.to_rgb().data) with torch.inference_mode(), self._autocast(): inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to( @@ -168,7 +190,98 @@ def query_score_rows( boxes[:, 1::2] = boxes[:, 1::2].clip(0.0, float(pil.height)) return boxes, kept + def _frame_features( + self, image: Image + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Text-independent tower outputs for one frame, cached by frame ts. + + Returns unit-normalized per-box class embeddings, the class head's + logit shift and scale, and clipped pixel ``(x1, y1, x2, y2)`` boxes. + A miss runs the image tower once; a hit runs nothing. + """ + key = image.ts + if key in self._features: + self._features.move_to_end(key) + return self._features[key] + + from transformers.image_transforms import center_to_corners_format + + self.forwards += 1 + pil = PILImage.fromarray(image.to_rgb().data) + model = self._model + head = model.class_head + with torch.inference_mode(), self._autocast(): + pixel_values = self._processor(images=pil, return_tensors="pt").pixel_values.to( + self.config.device + ) + feature_map = model.image_embedder(pixel_values)[0] + batch, height, width, dim = feature_map.shape + image_feats = feature_map.reshape(batch, height * width, dim) + pred_boxes = model.box_predictor(image_feats, feature_map) + class_embeds = head.dense0(image_feats) + class_embeds = class_embeds / ( + torch.linalg.norm(class_embeds, dim=-1, keepdim=True) + 1e-6 + ) + logit_shift = head.logit_shift(image_feats) + logit_scale = head.elu(head.logit_scale(image_feats)) + 1 + + # the same conversion post_process_grounded_object_detection runs: + # corners in model dtype, scaled to the padded square in float32 + boxes = center_to_corners_format(pred_boxes)[0].float() * float( + max(pil.width, pil.height) + ) + boxes[:, 0::2] = boxes[:, 0::2].clip(0.0, float(pil.width)) + boxes[:, 1::2] = boxes[:, 1::2].clip(0.0, float(pil.height)) + + entry = (class_embeds[0], logit_shift[0], logit_scale[0], boxes) + self._features[key] = entry + if len(self._features) > _FEATURE_CACHE_MAX: + self._features.popitem(last=False) + return entry + + def _query_embeds(self, queries: list[str]) -> torch.Tensor: + """Unit-normalized text embeddings, one row per query, cached per string.""" + missing = [q for q in queries if q not in self._text_embeds] + if missing: + with torch.inference_mode(), self._autocast(): + inputs = self._processor(text=[missing], return_tensors="pt").to(self.config.device) + embeds = self._model.owlv2.get_text_features(**inputs) + embeds = embeds / torch.linalg.norm(embeds, ord=2, dim=-1, keepdim=True) + for query, embed in zip(missing, embeds, strict=True): + self._text_embeds[query] = embed + return torch.stack([self._text_embeds[q] for q in queries]) + + def query_score_rows_batch( + self, + images: list[Image], + queries: list[str], + threshold: float = 0.1, + ) -> list[tuple[np.ndarray, np.ndarray]]: + """``query_score_rows`` over several images: one ``(boxes, rows)`` each. + + The image tower runs only for frames missing from the feature cache; + scoring any query set against a cached frame is a matmul against its + stored class embeddings, so repeated windows and new labels on seen + frames cost no model forwards. + """ + with torch.inference_mode(): + query_embeds = self._query_embeds(queries) + query_embeds = query_embeds / ( + torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6 + ) + out: list[tuple[np.ndarray, np.ndarray]] = [] + for image in images: + class_embeds, logit_shift, logit_scale, boxes_px = self._frame_features(image) + logits = (class_embeds @ query_embeds.T + logit_shift) * logit_scale + scores = torch.sigmoid(logits.to(torch.float32)) + keep = scores.max(dim=-1).values > threshold + out.append((boxes_px[keep].cpu().numpy(), scores[keep].cpu().numpy())) + return out + def stop(self) -> None: + self.score_cache.clear() + self._features.clear() + self._text_embeds.clear() if "_processor" in self.__dict__: del self.__dict__["_processor"] super().stop() diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 4653b89490..8f70ec109b 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -56,6 +56,8 @@ logger = setup_logger() +_SCORE_CACHE_MAX = 8192 + # A support candidate is an identity group: the member sightings of one # object. Everything a group reports is a plain function over its members. @@ -297,29 +299,47 @@ def localize( identities = [Identity(is_same=spatial(policy.cluster_radius_m)) for _ in queries] ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) + floor = policy.candidate_floor + cache = detector.score_cache + misses = [obs for obs in ordered if any((obs.ts, q, floor) not in cache for q in queries)] + if misses: + scored = detector.query_score_rows_batch( + [obs.data for obs in misses], queries, threshold=floor + ) + for obs, (boxes, rows) in zip(misses, scored, strict=True): + for j, q in enumerate(queries): + keep = rows[:, j] >= floor + cache[(obs.ts, q, floor)] = (boxes[keep], rows[keep, j]) + if len(cache) > _SCORE_CACHE_MAX: + cache.popitem(last=False) + for obs in ordered: - boxes, rows = detector.query_score_rows(obs.data, queries, threshold=policy.candidate_floor) + rows_per_label: list[tuple[Any, Any]] = [] + for q in queries: + key = (obs.ts, q, floor) + cache.move_to_end(key) + rows_per_label.append(cache[key]) + if not any(len(scores) for _boxes, scores in rows_per_label): + continue + img = obs.data candidates: list[Detection2DBBox] = [] - for box, row in zip(boxes, rows, strict=True): - bbox = (float(box[0]), float(box[1]), float(box[2]), float(box[3])) - for j, score in enumerate(row): - if score < policy.candidate_floor: - continue + for j, (boxes, scores) in enumerate(rows_per_label): + for box, score in zip(boxes, scores, strict=True): det = Detection2DBBox( - bbox=bbox, + bbox=(float(box[0]), float(box[1]), float(box[2]), float(box[3])), track_id=len(candidates), class_id=j, confidence=float(score), name=queries[j], - ts=obs.data.ts, - image=obs.data, + ts=img.ts, + image=img, ) if det.is_valid() and det.bbox_2d_volume() > 3000: candidates.append(det) if not candidates: continue - frame = segmenter.segment(ImageDetections2D(image=obs.data, detections=candidates)) + frame = segmenter.segment(ImageDetections2D(image=img, detections=candidates)) lifted = _lift(frame, rig, policy, plane) grounded = {det3d.track_id for det3d in lifted} for det2d in frame: From 131e03dfba8b0c5eb72fb1d32ab67d34328fc09d Mon Sep 17 00:00:00 2001 From: bogwi Date: Mon, 24 Aug 2026 17:34:25 +0900 Subject: [PATCH 08/28] add persistent identity groups, `identity_store.py` --- dimos/perception/memory/identity_store.py | 55 ++++++++++++++++ dimos/perception/memory/localize.py | 79 ++++++++++++++++------- 2 files changed, 110 insertions(+), 24 deletions(-) create mode 100644 dimos/perception/memory/identity_store.py diff --git a/dimos/perception/memory/identity_store.py b/dimos/perception/memory/identity_store.py new file mode 100644 index 0000000000..a0c9cddd4a --- /dev/null +++ b/dimos/perception/memory/identity_store.py @@ -0,0 +1,55 @@ +# 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. + +"""Persistent identity groups across :func:`localize` calls. + +Every persistence read and write in ``localize`` goes through this store and +nothing else. A store owned by a long-lived caller makes verification +evidence cumulative over everything seen since the store was created; the +query window then bounds only new detection work, and frames a label has +already ingested are never re-segmented or re-lifted. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from dimos.perception.detection.identity import Identity + +if TYPE_CHECKING: + from collections.abc import Callable + + from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC + + +@dataclass +class LabelIdentity: + identity: Identity # groups + merged, cumulative + ingested: set[float] = field(default_factory=set) # frame ts already segmented+lifted+added + ungrounded: tuple[float, float] | None = None # best (score, ts) with no depth + + +@dataclass +class IdentityStore: + labels: dict[str, LabelIdentity] = field(default_factory=dict) + + def get_or_create( + self, label: str, is_same: Callable[[Detection3DPC, Detection3DPC], bool] + ) -> LabelIdentity: + entry = self.labels.get(label) + if entry is None: + entry = LabelIdentity(identity=Identity(is_same=is_same)) + self.labels[label] = entry + return entry diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 8f70ec109b..9fd2e5cd9e 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -53,6 +53,7 @@ from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC + from dimos.perception.memory.identity_store import IdentityStore logger = setup_logger() @@ -220,6 +221,7 @@ def localize( require_pose: bool = True, policy: LocalizePolicy | None = None, trace: LocalizeTrace | list[LocalizeTrace] | None = None, + identity_store: IdentityStore | None = None, ) -> list[Localization] | list[list[Localization]]: """Every verified 3D instance of *query*, latest-seen first. @@ -242,6 +244,12 @@ def localize( The window is the index's - build it with :func:`embed_index`. Without a ``rig`` the store's shape decides one, and without a ``policy`` the rig supplies scale-appropriate defaults. + + An ``identity_store`` makes evidence cumulative: each label's groups + persist across calls, frames the store already ingested for a label are + skipped entirely, and an object stays answerable after it leaves the + window, its position still following the latest sighting. Without one, + every call verifies from scratch inside its window. """ rig = rig or Rig.from_store(store) policy = policy or rig.default_localize_policy() @@ -281,27 +289,42 @@ def localize( ordered = sorted(frames.values(), key=lambda obs: obs.ts) logger.info(f"detection: {len(ordered)} candidate frames for {len(queries)} labels") - if not ordered: - empty: list[list[Localization]] = [[] for _ in queries] - return [] if isinstance(query, str) else empty - - from dimos.perception.memory.support_plane import fit_support_plane - - anchors = [peak.pose_tuple for label_peaks in peaks_per_label for peak in label_peaks] - mx = sum(t[0] for t in anchors) / len(anchors) - my = sum(t[1] for t in anchors) / len(anchors) - cell = (round(mx / 2.0), round(my / 2.0)) - plane = rig._plane_cache.get(cell) - if plane is None: - plane = fit_support_plane(rig, ordered) - if plane is not None: - rig._plane_cache[cell] = plane - identities = [Identity(is_same=spatial(policy.cluster_radius_m)) for _ in queries] - ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) + if ordered: + from dimos.perception.memory.support_plane import fit_support_plane + + anchors = [peak.pose_tuple for label_peaks in peaks_per_label for peak in label_peaks] + mx = sum(t[0] for t in anchors) / len(anchors) + my = sum(t[1] for t in anchors) / len(anchors) + cell = (round(mx / 2.0), round(my / 2.0)) + plane = rig._plane_cache.get(cell) + if plane is None: + plane = fit_support_plane(rig, ordered) + if plane is not None: + rig._plane_cache[cell] = plane + + if identity_store is None: + entries = None + identities = [Identity(is_same=spatial(policy.cluster_radius_m)) for _ in queries] + ingested: list[set[float]] = [set() for _ in queries] + ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) + else: + entries = [ + identity_store.get_or_create(q, spatial(policy.cluster_radius_m)) for q in queries + ] + identities = [entry.identity for entry in entries] + ingested = [entry.ingested for entry in entries] + ungrounded = [entry.ungrounded for entry in entries] floor = policy.candidate_floor cache = detector.score_cache - misses = [obs for obs in ordered if any((obs.ts, q, floor) not in cache for q in queries)] + misses = [ + obs + for obs in ordered + if any( + obs.ts not in ingested[j] and (obs.ts, q, floor) not in cache + for j, q in enumerate(queries) + ) + ] if misses: scored = detector.query_score_rows_batch( [obs.data for obs in misses], queries, threshold=floor @@ -314,16 +337,20 @@ def localize( cache.popitem(last=False) for obs in ordered: - rows_per_label: list[tuple[Any, Any]] = [] - for q in queries: - key = (obs.ts, q, floor) + active = [j for j in range(len(queries)) if obs.ts not in ingested[j]] + if not active: + continue + rows_per_label: list[tuple[int, tuple[Any, Any]]] = [] + for j in active: + key = (obs.ts, queries[j], floor) cache.move_to_end(key) - rows_per_label.append(cache[key]) - if not any(len(scores) for _boxes, scores in rows_per_label): + rows_per_label.append((j, cache[key])) + ingested[j].add(obs.ts) + if not any(len(scores) for _j, (_boxes, scores) in rows_per_label): continue img = obs.data candidates: list[Detection2DBBox] = [] - for j, (boxes, scores) in enumerate(rows_per_label): + for j, (boxes, scores) in rows_per_label: for box, score in zip(boxes, scores, strict=True): det = Detection2DBBox( bbox=(float(box[0]), float(box[1]), float(box[2]), float(box[3])), @@ -361,6 +388,10 @@ def localize( obs.derive(data=ImageDetections2D(image=obs.data, detections=label_dets)) ) + if entries is not None: + for entry, best in zip(entries, ungrounded, strict=True): + entry.ungrounded = best + results = [ _finalize( q, From 909e6ac6bb39291d92e9bcb373dc4b1b464ba3e0 Mon Sep 17 00:00:00 2001 From: bogwi Date: Tue, 25 Aug 2026 11:12:29 +0900 Subject: [PATCH 09/28] make fuse to compute the exact all-time mean --- dimos/perception/detection/identity.py | 45 +++++++++++++++++++++++ dimos/perception/memory/identity_store.py | 7 +++- dimos/perception/memory/localize.py | 12 ++++-- dimos/perception/memory/rig.py | 1 + dimos/perception/memory/types.py | 1 + 5 files changed, 61 insertions(+), 5 deletions(-) diff --git a/dimos/perception/detection/identity.py b/dimos/perception/detection/identity.py index b36652e080..f7f8b66bf3 100644 --- a/dimos/perception/detection/identity.py +++ b/dimos/perception/detection/identity.py @@ -36,7 +36,11 @@ import operator from typing import TYPE_CHECKING, Any +import numpy as np + from dimos.memory.transform import Transformer +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC from dimos.perception.detection.type.imageDetections import ImageDetections @@ -55,6 +59,47 @@ def is_same(a: Detection3DPC, b: Detection3DPC) -> bool: return is_same +def fused(voxel: float) -> Callable[[Detection3DPC, Detection3DPC], Detection3DPC]: + """Merges two detections by combining their point clouds. For each grid cell of size *voxel* meters, + it keeps the mean position of all points grouped in that cell. + The weight of each fused point is the number of raw points that contributed to it. + """ + weights: dict[int, np.ndarray[Any, Any]] = {} + + def merge(a: Detection3DPC, b: Detection3DPC) -> Detection3DPC: + union = a + b + wa = weights.pop(id(a), None) + na = float(wa.sum()) if wa is not None else float(len(a.pointcloud)) + nb = float(len(b.pointcloud)) + union.center = Vector3( + (na * a.center.x + nb * b.center.x) / (na + nb), + (na * a.center.y + nb * b.center.y) / (na + nb), + (na * a.center.z + nb * b.center.z) / (na + nb), + ) + if voxel <= 0: + return union + pts = np.asarray(union.pointcloud.pointcloud.points) + w = np.ones(len(pts)) + if wa is not None: + w[: len(wa)] = wa + cells, inverse = np.unique( + np.floor(pts / voxel).astype(np.int64), axis=0, return_inverse=True + ) + wsum = np.zeros(len(cells)) + np.add.at(wsum, inverse, w) + psum = np.zeros((len(cells), 3)) + np.add.at(psum, inverse, pts * w[:, None]) + union.pointcloud = PointCloud2.from_numpy( + psum / wsum[:, None], + frame_id=union.pointcloud.frame_id, + timestamp=union.pointcloud.ts, + ) + weights[id(union)] = wsum + return union + + return merge + + class Identity(Transformer[Any, Detection3DPC]): """One detection3D per object, aggregated from every sighting. diff --git a/dimos/perception/memory/identity_store.py b/dimos/perception/memory/identity_store.py index a0c9cddd4a..4f8c2ad0c4 100644 --- a/dimos/perception/memory/identity_store.py +++ b/dimos/perception/memory/identity_store.py @@ -46,10 +46,13 @@ class IdentityStore: labels: dict[str, LabelIdentity] = field(default_factory=dict) def get_or_create( - self, label: str, is_same: Callable[[Detection3DPC, Detection3DPC], bool] + self, + label: str, + is_same: Callable[[Detection3DPC, Detection3DPC], bool], + merge: Callable[[Detection3DPC, Detection3DPC], Detection3DPC], ) -> LabelIdentity: entry = self.labels.get(label) if entry is None: - entry = LabelIdentity(identity=Identity(is_same=is_same)) + entry = LabelIdentity(identity=Identity(is_same=is_same, merge=merge)) self.labels[label] = entry return entry diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 9fd2e5cd9e..57c3150645 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -40,7 +40,7 @@ from dimos.memory.embed import EmbedImages from dimos.memory.transform import QualityWindow, peaks -from dimos.perception.detection.identity import Identity, spatial +from dimos.perception.detection.identity import Identity, fused, spatial from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.memory.rig import Rig @@ -304,12 +304,18 @@ def localize( if identity_store is None: entries = None - identities = [Identity(is_same=spatial(policy.cluster_radius_m)) for _ in queries] + identities = [ + Identity(is_same=spatial(policy.cluster_radius_m), merge=fused(policy.fuse_voxel_m)) + for _ in queries + ] ingested: list[set[float]] = [set() for _ in queries] ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) else: entries = [ - identity_store.get_or_create(q, spatial(policy.cluster_radius_m)) for q in queries + identity_store.get_or_create( + q, spatial(policy.cluster_radius_m), fused(policy.fuse_voxel_m) + ) + for q in queries ] identities = [entry.identity for entry in entries] ingested = [entry.ingested for entry in entries] diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index f70d24a174..b6c023b405 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -102,6 +102,7 @@ candidate_floor=0.18, accept_score=0.32, cluster_radius_m=0.30, + fuse_voxel_m=0.03, min_depth_points=30, max_object_extent_m=2.0, min_camera_range_m=0.5, diff --git a/dimos/perception/memory/types.py b/dimos/perception/memory/types.py index fbbeaadc89..eedb99cfb4 100644 --- a/dimos/perception/memory/types.py +++ b/dimos/perception/memory/types.py @@ -120,6 +120,7 @@ class LocalizePolicy: min_views: int = 2 # a support seen from one pose only is unconfirmed cluster_radius_m: float = 0.08 # observations within this are the same support + fuse_voxel_m: float = 0.01 # union-cloud voxel at the identity merge; 0 concatenates min_depth_points: int = 60 max_object_extent_m: float = 0.60 min_camera_range_m: float = 0.28 From e717a42bb70af3b75db9fbeeef92301c287058cd Mon Sep 17 00:00:00 2001 From: bogwi Date: Tue, 25 Aug 2026 13:51:22 +0900 Subject: [PATCH 10/28] detection small fixes --- dimos/perception/memory/localize.py | 2 +- dimos/perception/memory/rig.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 57c3150645..8204cd3e46 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -367,7 +367,7 @@ def localize( ts=img.ts, image=img, ) - if det.is_valid() and det.bbox_2d_volume() > 3000: + if det.is_valid(): candidates.append(det) if not candidates: continue diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index b6c023b405..db3d1d29e1 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -172,7 +172,7 @@ def _lattice_quantum(points: np.ndarray) -> float | None: quantum = float(diffs.min()) if quantum < 1e-4: return None - scaled = sample / quantum + scaled = (sample - sample[0]) / quantum if float(np.abs(scaled - np.round(scaled)).max()) > 0.01: return None return quantum @@ -620,7 +620,7 @@ def cloud_at(self, ts: float) -> PointCloud2 | None: # A grid-quantized source repeats the same cell in every # snapshot it persists through; accumulation must not count # one voxel once per snapshot. Cell keys dedup in one pass. - cells = np.round(stacked / quantum).astype(np.int64) + _CELL_OFFSET + cells = np.round((stacked - stacked[0]) / quantum).astype(np.int64) + _CELL_OFFSET keys = (cells[:, 0] << 42) | (cells[:, 1] << 21) | cells[:, 2] _, index = np.unique(keys, return_index=True) points = stacked[index] From 04706c5dcda92c39e2db9fbd76511b6f0a29aba1 Mon Sep 17 00:00:00 2001 From: bogwi Date: Tue, 25 Aug 2026 23:31:24 +0900 Subject: [PATCH 11/28] project images; tool_localize update --- dimos/perception/memory/localize.py | 4 +- dimos/perception/memory/rig.py | 4 +- dimos/perception/memory/tool_localize.py | 157 +++++++++++++++++------ 3 files changed, 120 insertions(+), 45 deletions(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 8204cd3e46..2f4d07510d 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -107,7 +107,7 @@ class LocalizeTrace: detection_frames: list[Any] = field(default_factory=list) # Observation[ImageDetections2D] matched: list[tuple[float, Detection3DPC]] = field(default_factory=list) verified: list[tuple[float, Detection3DPC]] = field(default_factory=list) - answers: list[Detection3DPC] = field(default_factory=list) # merged union per instance + answers: list[list[Detection3DPC]] = field(default_factory=list) # sightings per instance backdrop_ts: float | None = None @@ -510,7 +510,7 @@ def _finalize( ) if trace is not None: - trace.answers.append(merged) + trace.answers.append(members) if k == 0: trace.backdrop_ts = latest.ts diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index db3d1d29e1..8d4c71a7b2 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -645,7 +645,7 @@ def lift(self, detections: ImageDetections2D) -> ImageDetections3DPC | None: detections, cloud, self.camera_info, transform, _CLOUD_LIFT_FILTERS ) - def backdrop(self, ts: float) -> PointCloud2 | None: + def backdrop(self, ts: float, depth_trunc: float = 1.5) -> PointCloud2 | None: """World-frame scene cloud around ts, for plane fits and rendering.""" if self.depth is None: return self.cloud_at(ts) @@ -660,7 +660,7 @@ def backdrop(self, ts: float) -> PointCloud2 | None: except LookupError: return None return PointCloud2.from_rgbd( - color, depth, self.camera_info, depth_scale=0.001, depth_trunc=1.5 + color, depth, self.camera_info, depth_scale=0.001, depth_trunc=depth_trunc ).transform(-transform) # predicates diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index c5ff11e3c5..664ee83eae 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -35,6 +35,7 @@ import json from pathlib import Path import sys +from typing import Any from dimos.memory.store.sqlite import SqliteStore from dimos.memory.transform import throttle @@ -51,14 +52,19 @@ def render( ) -> None: """Write the .rrd - rerun stays an inline import. - Entity contract (the acceptance color cheat sheet): ``map`` backdrop, - then one subtree per query - ``detections//matched/*`` green, - ``detections//verified/*`` red, ``detections//answer/`` - always blue, one per verified instance. + Entity contract: the ``map`` backdrop is what was NOT detected - its + cells are carved free of every answer's cells - and each + ``detections//answer/`` is one verified instance textured + with each sighting's own image, wrapped in a labeled wireframe box. + Mobile rigs show the robot as a translucent box carrying the live + frustum; stationary rigs keep a frozen frustum per detection frame. """ + import numpy as np import rerun as rr import rerun.blueprint as rrb + from dimos.memory.vis.color import Color + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, _get_colormap_lut from dimos.visualization.rerun.init import rerun_init rerun_init("memory-localize") @@ -73,46 +79,98 @@ def render( ) ) - GREEN, RED, BLUE = [46, 204, 113], [231, 76, 60], [52, 120, 246] - point_size = 0.005 if rig.depth is not None else 0.015 + answer_parts = [ + det.pointcloud.points_f32() + for _, trace in traces + for members in trace.answers + for det in members + ] def at(ts: float) -> None: rr.set_time("ts", timestamp=ts) - # scene backdrop: for depth rigs the answer frame's RGBD cloud, for - # pointcloud rigs the window's scans merged into one map - if rig.depth is not None: + def carve(points: np.ndarray, spacing: float) -> np.ndarray: + """Drop map cells claimed by a detected object, plus one cell of halo.""" + if not answer_parts: + return points + + def keys(cells: np.ndarray) -> np.ndarray: + return (cells[:, 0] << 42) | (cells[:, 1] << 21) | cells[:, 2] + + offset = 1 << 20 + claimed = np.floor(np.vstack(answer_parts) / spacing).astype(np.int64) + offset + claimed = np.unique(claimed, axis=0) + steps = np.array([-1, 0, 1]) + neighbors = np.stack(np.meshgrid(steps, steps, steps), -1).reshape(-1, 3) + dilated = (claimed[:, None, :] + neighbors[None, :, :]).reshape(-1, 3) + cells = np.floor(points / spacing).astype(np.int64) + offset + return points[~np.isin(keys(cells), np.unique(keys(dilated)))] + + def height_points(points: np.ndarray, radius: float, blend: float, scale: float) -> Any: + """Height colormap, grayed by ``blend`` and dimmed by ``scale`` so + textured detections pop.""" + z = points[:, 2] + t = (z - z.min()) / (z.max() - z.min() + 1e-8) + turbo = _get_colormap_lut("turbo")[(t * 255).astype(np.uint8)] + colors = ((turbo * (1 - blend) + blend * 205) * scale).astype(np.uint8) + return rr.Points3D(positions=points, colors=colors, radii=radius) + + def image_colors(det: Any) -> np.ndarray: + """Sample the sighting's own image at each cloud point's reprojection.""" + points = det.pointcloud.points_f32() + matrix = det.transform.to_matrix() + cam = points @ matrix[:3, :3].T + matrix[:3, 3] + K = rig.camera_info.K + cols = np.round(cam[:, 0] / cam[:, 2] * K[0] + K[2]).astype(int) + rows = np.round(cam[:, 1] / cam[:, 2] * K[4] + K[5]).astype(int) + rgb = det.image.to_rgb().data + height, width = rgb.shape[:2] + return rgb[np.clip(rows, 0, height - 1), np.clip(cols, 0, width - 1)] + + # scene backdrop: stationary depth rigs get the answer frame's RGBD + # cloud, mobile depth rigs the window's frames merged, pointcloud rigs + # the window's scans merged; point size follows the source's spacing + point_size = 0.005 if rig.depth is not None else 0.015 + if rig.depth is not None and rig.mobile: + parts = [] + for obs in rig.color.after(t0).before(t1).transform(throttle(1.0)): + cloud = rig.backdrop(obs.ts, depth_trunc=4.0) + if cloud is not None: + parts.append(cloud.voxel_downsample(0.03).points_f32()) + if parts: + merged = PointCloud2.from_numpy(np.vstack(parts), frame_id=rig.world_frame) + carved = carve(merged.voxel_downsample(0.03).points_f32(), 0.03) + rr.log("map", height_points(carved, 0.013, 0.7, 0.5), static=True) + elif rig.depth is not None: backdrop_ts = next((t.backdrop_ts for _, t in traces if t.backdrop_ts is not None), None) if backdrop_ts is None: backdrop_ts = next((t.matched[0][0] for _, t in traces if t.matched), None) if backdrop_ts is not None: backdrop = rig.backdrop(backdrop_ts) if backdrop is not None: - rr.log( - "map", - backdrop.voxel_downsample(0.01).to_rerun(voxel_size=point_size), - static=True, - ) + carved = carve(backdrop.voxel_downsample(0.01).points_f32(), 0.01) + rr.log("map", height_points(carved, point_size / 2, 0.0, 1.0), static=True) else: - import numpy as np - - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - scans = [ points for obs in rig.cloud.after(t0).before(t1).transform(throttle(2.0)) if (points := rig.registered_scan(obs)) is not None ] if scans: + point_size = rig._cloud_quantum(scans[0]) or point_size merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) - rr.log( - "map", - merged.voxel_downsample(0.05).to_rerun(voxel_size=0.01), - static=True, - ) + carved = carve(merged.voxel_downsample(0.05).points_f32(), 0.05) + rr.log("map", height_points(carved, 0.022, 0.7, 0.5), static=True) - # live camera feed + frustum tracking the capture pose along the timeline + # live camera feed + frustum tracking the capture pose along the + # timeline; on mobile rigs a translucent box marks the robot rr.log("camera", rig.camera_info.to_rerun(), static=True) + if rig.mobile: + rr.log( + "robot", + rr.Boxes3D(half_sizes=[[0.2, 0.2, 0.2]], colors=[[120, 120, 120, 70]]), + static=True, + ) feed_throttle = 0.1 if (t1 - t0) <= 160 else 0.4 feed = rig.color.after(t0).before(t1).transform(throttle(feed_throttle)) for obs in feed: @@ -122,11 +180,21 @@ def at(ts: float) -> None: at(obs.ts) rr.log("camera/image", obs.data.to_rerun()) rr.log("camera", pose.to_rerun()) + if rig.mobile: + rr.log("robot", pose.to_rerun()) + + n_instances = sum(len(trace.answers) for _, trace in traces) + box_colors = [ + list(Color.from_cmap("turbo", k / max(n_instances - 1, 1)).rgb_u8()) + for k in range(n_instances) + ] + box_index = 0 for query, trace in traces: root = f"detections/{query.replace(' ', '_')}" - # marked frames: into the live feed, plus a frozen frustum at the capture pose + # marked frames into the live feed; only stationary rigs also drop + # a frozen frustum at the capture pose for i, obs in enumerate(trace.detection_frames): pose = rig.camera_pose(obs.ts) if pose is None: @@ -134,30 +202,37 @@ def at(ts: float) -> None: at(obs.ts) annotated = obs.data.annotated_image() rr.log("camera/image", annotated.to_rerun()) + if rig.mobile: + continue frame = f"{root}/frames/{i}" rr.log(frame, pose.to_rerun()) rr.log(frame, rig.camera_info.to_rerun()) rr.log(f"{frame}/image", annotated.to_rerun()) - # 3d detections: green = matched candidates, red = cross-view re-detections - for tag, entries, rgb in [ - ("matched", trace.matched, GREEN), - ("verified", trace.verified, RED), - ]: - for i, (ts, det) in enumerate(entries): - at(ts) - rr.log( - f"{root}/{tag}/{i}_{det.name.replace(' ', '_')}", - det.pointcloud.to_rerun(voxel_size=point_size, colors=rgb), - ) - - # the answers: one blue union cloud per verified instance - for i, answer in enumerate(trace.answers): - at(answer.ts) + # the answers: per verified instance, every sighting's cloud colored + # by its own image, wrapped in a static labeled wireframe box + for i, members in enumerate(trace.answers): + at(max(det.ts for det in members)) + positions = np.vstack([det.pointcloud.points_f32() for det in members]) + colors = np.vstack([image_colors(det) for det in members]) rr.log( f"{root}/answer/{i}", - answer.pointcloud.to_rerun(voxel_size=point_size, colors=BLUE), + rr.Points3D(positions=positions, colors=colors, radii=point_size / 2), + ) + aabb_min, aabb_max = positions.min(axis=0), positions.max(axis=0) + score = max(det.confidence for det in members) + rr.log( + f"{root}/answer/{i}/box", + rr.Boxes3D( + centers=[(aabb_min + aabb_max) / 2], + half_sizes=[(aabb_max - aabb_min) / 2], + colors=[box_colors[box_index]], + labels=[f"{query} {score:.2f}"], + fill_mode=rr.components.FillMode.MajorWireframe, + ), + static=True, ) + box_index += 1 def report(query: str, hits: list[Localization], lo: float) -> bool: From 429c8a8b54ab5fd5c3588d3006bc6edcc683be15 Mon Sep 17 00:00:00 2001 From: bogwi Date: Thu, 27 Aug 2026 01:54:18 +0900 Subject: [PATCH 12/28] improve detection --- .../type/detection3d/imageDetections3DPC.py | 7 +- .../detection/type/detection3d/pointcloud.py | 131 +++++- dimos/perception/memory/localize.py | 9 +- dimos/perception/memory/rig.py | 406 +++++++++++++++--- dimos/perception/memory/tool_localize.py | 9 +- 5 files changed, 475 insertions(+), 87 deletions(-) diff --git a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py index 50e0735b9a..310b2651da 100644 --- a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py +++ b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING -from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC +from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC, lattice_quantum from dimos.perception.detection.type.imageDetections import ImageDetections if TYPE_CHECKING: @@ -44,11 +44,13 @@ def from_2d( """Project every 2D detection into 3D, dropping any that yield no valid points. The cloud is projected through the camera once; each detection then - selects its points from that shared projection. + selects its points from that shared projection, with the mask splat + radius taken from the cloud's own lattice pitch. """ world_points, points_2d = Detection3DPC.project_cloud( world_pointcloud, camera_info, world_to_optical_transform ) + quantum = lattice_quantum(world_points) detections_3d = [ d3d for det in detections_2d @@ -62,6 +64,7 @@ def from_2d( world_pointcloud.frame_id, world_pointcloud.ts, filters, + splat_m=quantum / 2 if quantum is not None else None, ) ) is not None diff --git a/dimos/perception/detection/type/detection3d/pointcloud.py b/dimos/perception/detection/type/detection3d/pointcloud.py index 786f2152ae..6464c720e7 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -39,6 +39,30 @@ from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +def lattice_quantum(points: np.ndarray) -> float | None: + """The grid pitch when coordinates lie on a lattice; None for continuous scans. + + Grid-quantized sources (an occupancy map streamed as clouds) carry their + pitch in the data itself; it decides cell dedup, snapshot merging and the + projection splat. Continuous scans never collide and skip all of it. + """ + sample = points[:2048] + x = np.unique(sample[:, 0]) + if len(x) < 8: + return None + diffs = np.diff(x) + diffs = diffs[diffs > 1e-9] + if len(diffs) == 0: + return None + quantum = float(diffs.min()) + if quantum < 1e-4: + return None + scaled = (sample - sample[0]) / quantum + if float(np.abs(scaled - np.round(scaled)).max()) > 0.01: + return None + return quantum + + @dataclass class Detection3DPC(Detection3D): pointcloud: PointCloud2 = field(default_factory=PointCloud2) @@ -220,6 +244,73 @@ def from_depth( frame_id=detection_pc.frame_id, ) + @staticmethod + def project_pixels(points_camera: np.ndarray, camera_info: CameraInfo) -> np.ndarray: + """Pixel coordinates of camera-frame points under the camera's own model. + + Applies the recorded distortion (equidistant fisheye or the radtan + family) so cloud pixels land where the image's pixels actually are; + an undistorted calibration falls through to the pinhole projection. + Points beyond the model's monotonic radius get out-of-image pixels. + """ + fx, fy = camera_info.K[0], camera_info.K[4] + cx, cy = camera_info.K[2], camera_info.K[5] + coefficients = np.asarray( + camera_info.D if camera_info.D is not None else (), dtype=np.float64 + ) + xy = points_camera[:, :2] / points_camera[:, 2:3] + if coefficients.size == 0 or not np.any(coefficients): + return np.column_stack((xy[:, 0] * fx + cx, xy[:, 1] * fy + cy)) + + import cv2 + + fisheye = camera_info.distortion_model == "equidistant" + # The polynomial is calibrated only out to the image corner; beyond + # it the projection extrapolates or folds back into the image, so + # points past the corner's angle (or past the first fold) are + # unmappable. + corners = np.array( + [[0, 0], [camera_info.width, 0], [0, camera_info.height], + [camera_info.width, camera_info.height]], + dtype=np.float64, + ) + corner_limit = float( + np.hypot((corners[:, 0] - cx) / fx, (corners[:, 1] - cy) / fy).max() + ) + theta = np.linspace(0.0, np.pi / 2 * 0.99, 2048) + if fisheye: + k = np.zeros(4) + k[: min(4, coefficients.size)] = coefficients[:4] + distorted = theta * ( + 1 + k[0] * theta**2 + k[1] * theta**4 + k[2] * theta**6 + k[3] * theta**8 + ) + radius = np.tan(theta) + else: + k = np.zeros(3) + radial = coefficients[[0, 1]].tolist() + ( + [coefficients[4]] if coefficients.size > 4 else [] + ) + k[: len(radial)] = radial + radius = np.tan(theta) + distorted = radius * (1 + k[0] * radius**2 + k[1] * radius**4 + k[2] * radius**6) + beyond = np.nonzero((np.diff(distorted) <= 0) | (distorted[1:] > corner_limit))[0] + r_max = radius[beyond[0]] if len(beyond) else radius[-1] + + pixels = np.full((len(points_camera), 2), -1.0) + mappable = (xy**2).sum(axis=1) <= r_max**2 + if mappable.any(): + pts = np.ascontiguousarray(points_camera[mappable, :3], dtype=np.float64) + camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64) + zero = np.zeros(3) + if fisheye: + projected, _ = cv2.fisheye.projectPoints( + pts.reshape(-1, 1, 3), zero, zero, camera_matrix, coefficients[:4] + ) + else: + projected, _ = cv2.projectPoints(pts, zero, zero, camera_matrix, coefficients) + pixels[mappable] = projected.reshape(-1, 2) + return pixels + @staticmethod def project_cloud( world_pointcloud: PointCloud2, @@ -232,10 +323,6 @@ def project_cloud( coordinates - the detection-independent half of ``from_2d``, shared by every detection of one frame via ``from_projection``. """ - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] - camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) - world_points, _ = world_pointcloud.as_numpy() # Project points to camera frame @@ -251,9 +338,7 @@ def project_cloud( if len(world_points) == 0: return world_points, np.empty((0, 2)) - # Project to 2D - points_2d_homogeneous = (camera_matrix @ points_camera[:, :3].T).T - points_2d = points_2d_homogeneous[:, :2] / points_2d_homogeneous[:, 2:3] + points_2d = Detection3DPC.project_pixels(points_camera, camera_info) # Filter points within image bounds in_image_mask = ( @@ -275,12 +360,16 @@ def from_projection( frame_id: str, timestamp: float, filters: list[PointCloudFilter] | None = None, + splat_m: float | None = None, ) -> Detection3DPC | None: """Create a Detection3D by selecting from a shared frame projection. ``world_points`` and ``points_2d`` come from ``project_cloud`` for this detection's frame; only the mask selection and the per-detection - filters run here. + filters run here. ``splat_m`` is the half-size of a source cell: a + point then selects when any of its projected footprint touches the + mask, not only its center pixel - center-only sampling starves masks + a few cells wide. """ # Set default filters if none provided if filters is None: @@ -298,9 +387,23 @@ def from_projection( # (Detection2DSeg), else bbox with a small margin seg_mask = getattr(det, "mask", None) if seg_mask is not None: - cols = np.minimum(points_2d[:, 0].astype(int), seg_mask.shape[1] - 1) - rows = np.minimum(points_2d[:, 1].astype(int), seg_mask.shape[0] - 1) - in_det_mask = seg_mask[rows, cols] > 0 + height, width = seg_mask.shape[:2] + + def mask_hit(cols_f: np.ndarray, rows_f: np.ndarray) -> np.ndarray: + cols = np.clip(cols_f.astype(int), 0, width - 1) + rows = np.clip(rows_f.astype(int), 0, height - 1) + hit: np.ndarray = seg_mask[rows, cols] > 0 + return hit + + in_det_mask = mask_hit(points_2d[:, 0], points_2d[:, 1]) + if splat_m is not None: + camera = world_to_optical_transform.inverse().translation.to_numpy() + ranges = np.linalg.norm(world_points - camera, axis=1) + radius = splat_m * camera_info.K[0] / np.maximum(ranges, 1e-6) + for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)): + in_det_mask |= mask_hit( + points_2d[:, 0] + dx * radius, points_2d[:, 1] + dy * radius + ) else: x_min, y_min, x_max, y_max = det.bbox margin = 5 # pixels @@ -364,11 +467,14 @@ def from_2d( # type: ignore[override] One-detection convenience over ``project_cloud`` + ``from_projection``; callers lifting several detections of one frame should project once - and call ``from_projection`` per detection instead. + and call ``from_projection`` per detection instead. The mask splat + radius comes from the cloud's own lattice pitch, so every caller + samples a grid-quantized source the same way. """ world_points, points_2d = cls.project_cloud( world_pointcloud, camera_info, world_to_optical_transform ) + quantum = lattice_quantum(world_points) return cls.from_projection( det, world_points, @@ -378,4 +484,5 @@ def from_2d( # type: ignore[override] world_pointcloud.frame_id, world_pointcloud.ts, filters, + splat_m=quantum / 2 if quantum is not None else None, ) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 2f4d07510d..c02d66a991 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -43,7 +43,7 @@ from dimos.perception.detection.identity import Identity, fused, spatial from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D -from dimos.perception.memory.rig import Rig +from dimos.perception.memory.rig import CLOUD_MIN_POINTS, Rig from dimos.perception.memory.types import Localization, LocalizePolicy, Support from dimos.utils.logging_config import setup_logger @@ -157,14 +157,15 @@ def _lift( if pose is None: return [] camera = np.array([pose.position.x, pose.position.y, pose.position.z]) - lifted = rig.lift(detections) + lifted = rig.lift(detections, plane) if lifted is None: return [] + floor = policy.min_depth_points if rig.cloud is None else CLOUD_MIN_POINTS valid: list[Detection3DPC] = [] for det3d in lifted: points = np.asarray(det3d.pointcloud.pointcloud.points) - if len(points) < policy.min_depth_points: + if len(points) < floor: continue extent = points.max(axis=0) - points.min(axis=0) if float(extent.max()) > policy.max_object_extent_m: @@ -173,7 +174,7 @@ def _lift( if float(np.median(ranges)) < policy.min_camera_range_m: continue if plane is not None: - heights = plane.height_above(points) + heights = rig.support_heights(detections.ts, plane, points) low = float(np.quantile(heights, 0.05)) high = float(np.quantile(heights, 0.95)) if low > policy.surface_patch_min_drop_m and high < policy.surface_patch_max_rise_m: diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 8d4c71a7b2..cc2e1a1394 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -22,9 +22,9 @@ * **Geometry source.** 3D geometry comes from an aligned ``depth`` stream, unprojected per detection mask, or from a world-frame pointcloud stream (a registered lidar), projected through the camera per detection mask. - Registered scans are sparse, so the cloud at a timestamp is the - concatenation of the scans in a short window around it - the scene is - static in world frame, which is what makes accumulation valid. + The cloud at a timestamp merges the scans in a short window around it: + sparse continuous scans concatenate whole, rolling-map snapshots merge + nearest-first with cleared cells honored (see ``Rig.cloud_at``). ``Rig.from_store`` recognizes both recording shapes; every field can also be supplied directly for live stores whose streams are still filling. @@ -45,12 +45,14 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.perception.detection.project import sees as project_sees from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC +from dimos.perception.detection.type.detection3d.pointcloud import lattice_quantum from dimos.perception.detection.type.detection3d.pointcloud_filters import ( range_cluster, statistical, ) from dimos.perception.memory import gates from dimos.perception.memory.gates import SPEED_MAX, STILL_ENVELOPE, TF_TOLERANCE +from dimos.perception.memory.support_plane import PLANE_DISTANCE_CLOUD from dimos.perception.memory.types import InventoryPolicy, LocalizePolicy from dimos.utils.logging_config import setup_logger @@ -82,7 +84,6 @@ CLOUD_ACCUM_S = 4.0 _SCAN_CACHE_MAX = 256 # registered scans held per rig; a window needs a few dozen -_CELL_OFFSET = 1 << 20 # shifts lattice cell indices positive for 21-bit key packing EMBED_HZ = 1.0 # index density for a wrist camera parked over a workspace # A walking robot changes viewpoint every frame and its frames blur @@ -90,10 +91,39 @@ # sharp sightings. WALK_EMBED_HZ = 3.0 +# Color-stream delay estimation: image timestamps stamped at receive lag the +# pose source by a constant the recording itself reveals - the lag that best +# correlates optical-flow yaw rate with the camera's heading rate. The grid +# is the searched span and its resolution; a true delay outside the span +# lands on an edge and is refused rather than clamped. +DELAY_LAGS = np.arange(-0.5, 0.5001, 0.01) +# Flow is sampled in short windows spread over the recording - a compute +# budget, like the pose samples of _camera_span. +DELAY_WINDOWS = 12 +DELAY_WINDOW_S = 2.5 + # Sparse projected clouds: split off background seen through the mask, then a # loose outlier trim. The dense-cloud defaults (raycast + radius) assume a # density registered lidar does not have. -_CLOUD_LIFT_FILTERS = [range_cluster(), statistical(nb_neighbors=12, std_ratio=2.0)] +_CLOUD_TRIM_NEIGHBORS = 12 +_CLOUD_LIFT_FILTERS = [ + range_cluster(), + statistical(nb_neighbors=_CLOUD_TRIM_NEIGHBORS, std_ratio=2.0), +] +# The projected lift's evidence floor. Depth-pixel counts and lattice-cell +# counts do not compare - a bottle-sized object can never cover thirty 5 cm +# cells - so a policy's depth-scale floor does not apply to a projected +# cloud; the floor there is the smallest cloud the trim can vet, its own +# neighborhood. +CLOUD_MIN_POINTS = _CLOUD_TRIM_NEIGHBORS + 1 + + +def _column_keys(points: np.ndarray, quantum: float, anchor: np.ndarray) -> np.ndarray: + """Packed XY lattice-cell key per point, anchored so any grid phase maps exactly.""" + cells = np.round((points[:, :2] - anchor) / quantum).astype(np.int64) + (1 << 20) + keys: np.ndarray = (cells[:, 0] << 21) | cells[:, 1] + return keys + # Room-scale policies for mobile-robot rigs: objects are furniture-sized, # viewpoints meters apart, odometry drifts centimeters between passes, and @@ -154,28 +184,126 @@ def _camera_span(rig: Rig) -> float: return float(np.linalg.norm(spread.max(axis=0) - spread.min(axis=0))) -def _lattice_quantum(points: np.ndarray) -> float | None: - """The grid pitch when coordinates lie on a lattice; None for continuous scans. +def _heading_series(rig: Rig, spans: list[tuple[float, float]]) -> tuple[np.ndarray, np.ndarray]: + """Camera heading rate over the given time spans: (rate midpoints, rates). - Grid-quantized sources (an occupancy map streamed as clouds) repeat the - same cell across snapshots and dedup by cell key; continuous scans never - collide and skip dedup entirely. + A poses stream is read directly - a rigid mount adds a constant offset, + so the base yaw rate is the camera heading rate. A tf rig has no pose + stream to iterate; the optical axis is sampled through tf on each span's + frame-period grid instead. + """ + times: list[float] = [] + headings: list[float] = [] + if rig.poses is not None: + for obs in rig.poses.after(spans[0][0]).before(spans[-1][1]): + q = obs.pose_stamped.orientation + times.append(obs.ts) + headings.append(np.arctan2(2 * (q.w * q.z + q.x * q.y), 1 - 2 * (q.y * q.y + q.z * q.z))) + else: + for lo, hi in spans: + frame_ts = [obs.ts for obs in rig.color.after(lo).before(hi)] + if len(frame_ts) < 2: + continue + step = float(np.median(np.diff(frame_ts))) + for t in np.arange(lo, hi, step): + transform = rig.world_to_optical(float(t)) + if transform is None: + continue + axis = (-transform).to_matrix()[:3, 2] + times.append(float(t)) + headings.append(np.arctan2(axis[1], axis[0])) + if len(times) < 3: + return np.empty(0), np.empty(0) + stamps = np.array(times) + yaw = np.unwrap(np.array(headings)) + return (stamps[1:] + stamps[:-1]) / 2, np.diff(yaw) / np.diff(stamps) + + +def estimate_color_delay(rig: Rig) -> float: + """Constant lag of color timestamps behind the rig's pose source, in seconds. + + Optical-flow yaw rate between consecutive frames is compared against the + camera heading rate at a grid of candidate lags; the lag maximizing + their absolute correlation is the delay (the sign of the relation is + mount-dependent). The estimate validates itself on the recording: the + full sample and its two interleaved halves must estimate mutually equal + lags to within one frame period - each flow sample integrates motion + over a frame period, so that is the measurement's own resolution - and + every peak must be interior to the searched span. A recording without + usable rotation fails that and keeps 0.0. """ - sample = points[:2048] - x = np.unique(sample[:, 0]) - if len(x) < 8: - return None - diffs = np.diff(x) - diffs = diffs[diffs > 1e-9] - if len(diffs) == 0: - return None - quantum = float(diffs.min()) - if quantum < 1e-4: - return None - scaled = (sample - sample[0]) / quantum - if float(np.abs(scaled - np.round(scaled)).max()) > 0.01: - return None - return quantum + import cv2 + + try: + t0, t1 = rig.color.get_time_range() + except LookupError: + return 0.0 # live store, nothing recorded yet + + fx = rig.camera_info.K[0] + flow_ts: list[float] = [] + flow_rate: list[float] = [] + flow_dt: list[float] = [] + span = max(t1 - t0 - DELAY_WINDOW_S, 0.0) + starts = [t0 + span * k / max(DELAY_WINDOWS - 1, 1) for k in range(DELAY_WINDOWS)] + windows = [(lo, lo + DELAY_WINDOW_S) for lo in starts] + for lo, hi in windows: + previous: np.ndarray | None = None + previous_ts = 0.0 + for obs in rig.color.after(lo).before(hi): + frame = obs.data.to_opencv() + scale = 640 / frame.shape[1] + gray = cv2.resize(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), (640, int(frame.shape[0] * scale))) + if previous is not None and obs.ts > previous_ts: + corners = cv2.goodFeaturesToTrack(previous, 150, 0.01, 12) + if corners is not None: + moved, status, _err = cv2.calcOpticalFlowPyrLK(previous, gray, corners, None) # type: ignore[call-overload] + tracked = status.ravel().astype(bool) + if tracked.any(): + dx = float(np.median(moved[tracked, 0, 0] - corners[tracked, 0, 0])) / scale + flow_ts.append(0.5 * (obs.ts + previous_ts)) + flow_rate.append(dx / fx / (obs.ts - previous_ts)) + flow_dt.append(obs.ts - previous_ts) + previous, previous_ts = gray, obs.ts + if len(flow_ts) < 4: + return 0.0 # each interleaved half needs a correlation of its own + + margin = float(DELAY_LAGS[-1]) + rate_t, rates = _heading_series(rig, [(lo - margin, hi + margin) for lo, hi in windows]) + if len(rates) == 0: + return 0.0 + + flow_t = np.array(flow_ts) + flow = np.array(flow_rate) + + def lag_of(sel: np.ndarray) -> float | None: + strength = np.abs( + np.array( + [ + np.corrcoef(flow[sel], np.interp(flow_t[sel] - lag, rate_t, rates))[0, 1] + for lag in DELAY_LAGS + ] + ) + ) + if not np.isfinite(strength).all(): + return None + peak = int(strength.argmax()) + if peak == 0 or peak == len(DELAY_LAGS) - 1: + return None # the true lag is outside the searched span + lag = float(DELAY_LAGS[peak]) + a, b, c = strength[peak - 1], strength[peak], strength[peak + 1] + denominator = a - 2 * b + c + if denominator < 0: + lag += 0.5 * (a - c) / denominator * float(DELAY_LAGS[1] - DELAY_LAGS[0]) + return lag + + everything = np.arange(len(flow)) + lags = [lag_of(everything), lag_of(everything[0::2]), lag_of(everything[1::2])] + if any(lag is None for lag in lags): + return 0.0 + estimates = cast("list[float]", lags) + if max(estimates) - min(estimates) > float(np.median(flow_dt)): + return 0.0 + return estimates[0] class RegisterScans(Transformer["PointCloud2", "PointCloud2"]): @@ -253,18 +381,24 @@ class Rig: tf_tolerance: float = TF_TOLERANCE cloud_accum_s: float = CLOUD_ACCUM_S speed_max: float = SPEED_MAX + color_delay: float = 0.0 # s - color timestamps lag the pose stream by this scene_gate: bool = True embed_hz: float = EMBED_HZ mobile: bool = False # camera rides a mobile base: room-scale policies - _cloud_memo: tuple[float, PointCloud2] | None = field(default=None, repr=False, init=False) + # (ts, merged cloud, the window's measured lattice quantum) + _cloud_memo: tuple[float, PointCloud2, float | None] | None = field( + default=None, repr=False, init=False + ) + # (ts, plane, per-column floor table of the frame's merged cloud) + _shell_memo: tuple[float, Any, tuple[np.ndarray, np.ndarray, float, np.ndarray]] | None = field( + default=None, repr=False, init=False + ) _scan_cache: OrderedDict[float, np.ndarray | None] = field( default_factory=OrderedDict, repr=False, init=False ) _plane_cache: dict[tuple[int, int], SupportPlane] = field( default_factory=dict, repr=False, init=False ) - _quantum: float | None = field(default=None, repr=False, init=False) - _quantum_known: bool = field(default=False, repr=False, init=False) @classmethod def from_store( @@ -473,15 +607,19 @@ def from_store( rig.speed_max = WALK_SPEED_MAX rig.scene_gate = False rig.embed_hz = WALK_EMBED_HZ + if camera_info is not None: + rig.color_delay = estimate_color_delay(rig) logger.info( f"rig: color={color_name!r} depth={depth_name!r} cloud={cloud_name!r} " - f"tf={tf_name!r} world={world_frame!r} span={span:.1f}m mobile={rig.mobile}" + f"tf={tf_name!r} world={world_frame!r} span={span:.1f}m mobile={rig.mobile} " + f"color_delay={rig.color_delay * 1000:.0f}ms" ) return rig # pose def world_to_optical(self, ts: float) -> Transform | None: + ts -= self.color_delay # color stamps lag; the capture instant is earlier if self.tf is not None: return self.tf.get(self.optical_frame, self.world_frame, ts, self.tf_tolerance) pose = self.pose_at(ts) @@ -491,18 +629,38 @@ def world_to_optical(self, ts: float) -> Transform | None: return -(Transform.from_pose("base_link", pose) + mount) def pose_at(self, ts: float) -> PoseStamped | None: - """World base pose nearest ts, from the poses stream. + """World base pose at ts, interpolated between the bracketing samples. - ``at().first()`` is window-earliest, not window-nearest; a walking - robot covers centimeters per pose period, so the nearest pose in the - window is what keeps the projection aligned. + A walking robot covers centimeters and degrees per pose period, so + snapping to a recorded sample misplaces the projection; interpolation + follows the motion. Outside the bracketed span the nearest sample in + the tolerance window stands. """ + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + candidates = list(self.poses.at(ts, self.tf_tolerance)) if not candidates: return None - obs = min(candidates, key=lambda o: abs(o.ts - ts)) - pose: PoseStamped | None = obs.pose_stamped - return pose + earlier = [o for o in candidates if o.ts <= ts] + later = [o for o in candidates if o.ts > ts] + if not earlier or not later: + pose: PoseStamped | None = min(candidates, key=lambda o: abs(o.ts - ts)).pose_stamped + return pose + a = max(earlier, key=lambda o: o.ts).pose_stamped + b = min(later, key=lambda o: o.ts).pose_stamped + alpha = (ts - a.ts) / (b.ts - a.ts) + qa = np.array([a.orientation.x, a.orientation.y, a.orientation.z, a.orientation.w]) + qb = np.array([b.orientation.x, b.orientation.y, b.orientation.z, b.orientation.w]) + if float(qa @ qb) < 0: + qb = -qb + q = (1 - alpha) * qa + alpha * qb + q /= np.linalg.norm(q) + return PoseStamped( + ts=ts, + frame_id=a.frame_id, + position=a.position + (b.position - a.position) * alpha, + orientation=(float(q[0]), float(q[1]), float(q[2]), float(q[3])), + ) def camera_pose(self, ts: float) -> PoseStamped | None: """World pose of the camera optical frame at ts.""" @@ -593,43 +751,158 @@ def registered_scan(self, scan: Observation[PointCloud2]) -> np.ndarray | None: self._scan_cache.popitem(last=False) return points - def _cloud_quantum(self, points: np.ndarray) -> float | None: - if not self._quantum_known: - self._quantum = _lattice_quantum(points) - self._quantum_known = True - return self._quantum - def cloud_at(self, ts: float) -> PointCloud2 | None: - """World-frame geometry at ts: the scans accumulated around it.""" + """World-frame geometry at ts: the window's scans merged. + + The merge rule follows the source's shape, measured per window from + the points themselves. A grid-quantized stream is a rolling + occupancy map: each snapshot is already temporally integrated over + the area it covers and the map clears what moved away, so snapshots + merge nearest-ts first, a farther snapshot contributing only cells + outside the XY coverage of every nearer one - plain accumulation + would resurrect every moved object's trail. Continuous scans are + sparse and the scene static in world frame, so they accumulate + whole. + """ if self._cloud_memo is not None and self._cloud_memo[0] == ts: return self._cloud_memo[1] from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 scans = self.cloud.after(ts - self.cloud_accum_s).before(ts + self.cloud_accum_s) - parts = [p for p in (self.registered_scan(scan) for scan in scans) if p is not None] - if not parts: + pairs = [ + (scan.ts, points) + for scan in scans + if (points := self.registered_scan(scan)) is not None + ] + if not pairs: return None - if len(parts) == 1: - points = parts[0] + nearest = min(pairs, key=lambda pair: abs(pair[0] - ts))[1] + quantum = lattice_quantum(nearest) + if len(pairs) == 1: + points = pairs[0][1] + elif quantum is None: + points = np.vstack([p for _t, p in pairs]) else: - stacked = np.vstack(parts) - quantum = self._cloud_quantum(parts[0]) - if quantum is None: - points = stacked - else: - # A grid-quantized source repeats the same cell in every - # snapshot it persists through; accumulation must not count - # one voxel once per snapshot. Cell keys dedup in one pass. - cells = np.round((stacked - stacked[0]) / quantum).astype(np.int64) + _CELL_OFFSET - keys = (cells[:, 0] << 42) | (cells[:, 1] << 21) | cells[:, 2] - _, index = np.unique(keys, return_index=True) - points = stacked[index] + anchor = nearest[0, :2] + kept: list[np.ndarray] = [] + lows: list[np.ndarray] = [] + highs: list[np.ndarray] = [] + for _t, part in sorted(pairs, key=lambda pair: abs(pair[0] - ts)): + cells = np.round((part[:, :2] - anchor) / quantum).astype(np.int64) + if lows: + lo, hi = np.stack(lows), np.stack(highs) + covered = ( + (cells[:, None, :] >= lo[None]) & (cells[:, None, :] <= hi[None]) + ).all(axis=2) + fresh = ~covered.any(axis=1) + if fresh.any(): + kept.append(part[fresh]) + else: + kept.append(part) + lows.append(cells.min(axis=0)) + highs.append(cells.max(axis=0)) + points = np.vstack(kept) merged = PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=ts) - self._cloud_memo = (ts, merged) + self._cloud_memo = (ts, merged, quantum) return merged - def lift(self, detections: ImageDetections2D) -> ImageDetections3DPC | None: - """2D detections to world-frame 3D clouds, or None without geometry/pose.""" + def _shell_table( + self, ts: float, plane: SupportPlane + ) -> tuple[np.ndarray, np.ndarray, float, np.ndarray] | None: + """Per-column floor of the frame's merged cloud; None for continuous sources. + + A rolling map registers each snapshot with its own odometry error, + so the support surface sits at a different absolute level per + region. The floor of a point's own XY column is the local reference + a single global plane cannot be. + """ + cloud = self.cloud_at(ts) + if cloud is None or self._cloud_memo[2] is None: # type: ignore[index] + return None + memo = self._shell_memo + if memo is not None and memo[0] == ts and memo[1] is plane: + return memo[2] + quantum = cast("float", self._cloud_memo[2]) # type: ignore[index] + points = cloud.as_numpy()[0] + anchor = points[0, :2] + keys = _column_keys(points, quantum, anchor) + heights = plane.height_above(points) + order = np.argsort(keys) + keys_sorted = keys[order] + starts = np.nonzero(np.concatenate(([True], np.diff(keys_sorted) != 0)))[0] + table = (keys_sorted[starts], np.minimum.reduceat(heights[order], starts), quantum, anchor) + self._shell_memo = (ts, plane, table) + return table + + def support_heights(self, ts: float, plane: SupportPlane, points: np.ndarray) -> np.ndarray: + """Signed heights of points above the frame's support shell. + + Depth rigs and continuous-scan sources measure against the fitted + plane directly. A rolling-map source measures against the local + column floor of the frame's own merged cloud (see ``_shell_table``); + ``points`` must come from that cloud, which is what every projected + lift produces. + """ + heights = plane.height_above(points) + if self.cloud is None: + return heights + table = self._shell_table(ts, plane) + if table is None: + return heights + col_keys, col_floor, quantum, anchor = table + idx = np.searchsorted(col_keys, _column_keys(points, quantum, anchor)) + local: np.ndarray = heights - col_floor[idx] + return local + + def _support_strip(self, ts: float, plane: SupportPlane, shell: float, gap: float = 0.3) -> Any: + """Filter dropping support-shell points outside the object's own stance. + + A misaligned mask row collects the support surface along the whole + view ray; shell points survive only within the camera-range span of + the detection's above-shell structure - under the object, not along + the approach. The stance is the dominant range cluster of the + above-shell points (the same ``gap`` split as ``range_cluster``), so + background caught on the mask rim cannot widen it. Points below the + shell (mirror returns through a glossy surface) never survive. A + detection with no structure above the shell passes untouched. + """ + + def filter_func(det: Any, pc: Any, ci: Any, tf: Any) -> Any: + points, _ = pc.as_numpy() + local = self.support_heights(ts, plane, points) + above = local > shell + if not above.any(): + return pc + camera = tf.inverse().translation.to_numpy() + ranges = np.linalg.norm(points - camera, axis=1) + stance = np.sort(ranges[above]) + splits = np.nonzero(np.diff(stance) > gap)[0] + starts = np.concatenate(([0], splits + 1)) + ends = np.concatenate((splits + 1, [len(stance)])) + median_idx = np.searchsorted(stance, np.median(stance)) + start, end = next( + (s, e) for s, e in zip(starts, ends, strict=True) if s <= median_idx < e + ) + in_stance = (ranges >= stance[start]) & (ranges <= stance[end - 1]) + keep = above | ((local >= -shell) & in_stance) + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + return PointCloud2.from_numpy(points[keep], frame_id=pc.frame_id, timestamp=pc.ts) + + return filter_func + + def lift( + self, detections: ImageDetections2D, plane: SupportPlane | None = None + ) -> ImageDetections3DPC | None: + """2D detections to world-frame 3D clouds, or None without geometry/pose. + + With a support ``plane``, a projected lift strips support-shell + points outside each detection's stance before the generic filters. + The shell is the support surface's own occupied band: a plane + crossing a lattice straddles at most two adjacent levels, so a + rolling map's shell ends halfway to the third; a continuous source's + shell is the plane fit's inlier distance. + """ transform = self.world_to_optical(detections.ts) if transform is None: return None @@ -641,9 +914,12 @@ def lift(self, detections: ImageDetections2D) -> ImageDetections3DPC | None: cloud = self.cloud_at(detections.ts) if cloud is None: return None - return ImageDetections3DPC.from_2d( - detections, cloud, self.camera_info, transform, _CLOUD_LIFT_FILTERS - ) + filters = _CLOUD_LIFT_FILTERS + if plane is not None: + quantum = cast("float | None", self._cloud_memo[2]) # type: ignore[index] + shell = 1.5 * quantum if quantum is not None else PLANE_DISTANCE_CLOUD + filters = [self._support_strip(detections.ts, plane, shell), *_CLOUD_LIFT_FILTERS] + return ImageDetections3DPC.from_2d(detections, cloud, self.camera_info, transform, filters) def backdrop(self, ts: float, depth_trunc: float = 1.5) -> PointCloud2 | None: """World-frame scene cloud around ts, for plane fits and rendering.""" diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index 664ee83eae..093cab0725 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -39,6 +39,7 @@ from dimos.memory.store.sqlite import SqliteStore from dimos.memory.transform import throttle +from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC, lattice_quantum from dimos.perception.memory.localize import LocalizeTrace from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Localization @@ -120,9 +121,9 @@ def image_colors(det: Any) -> np.ndarray: points = det.pointcloud.points_f32() matrix = det.transform.to_matrix() cam = points @ matrix[:3, :3].T + matrix[:3, 3] - K = rig.camera_info.K - cols = np.round(cam[:, 0] / cam[:, 2] * K[0] + K[2]).astype(int) - rows = np.round(cam[:, 1] / cam[:, 2] * K[4] + K[5]).astype(int) + pixels = Detection3DPC.project_pixels(cam, rig.camera_info) + cols = np.round(pixels[:, 0]).astype(int) + rows = np.round(pixels[:, 1]).astype(int) rgb = det.image.to_rgb().data height, width = rgb.shape[:2] return rgb[np.clip(rows, 0, height - 1), np.clip(cols, 0, width - 1)] @@ -157,7 +158,7 @@ def image_colors(det: Any) -> np.ndarray: if (points := rig.registered_scan(obs)) is not None ] if scans: - point_size = rig._cloud_quantum(scans[0]) or point_size + point_size = lattice_quantum(scans[0]) or point_size merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) carved = carve(merged.voxel_downsample(0.05).points_f32(), 0.05) rr.log("map", height_points(carved, 0.022, 0.7, 0.5), static=True) From 6935a6891a5a33d95bbb9da59ca5f97951d386bc Mon Sep 17 00:00:00 2001 From: bogwi Date: Thu, 27 Aug 2026 14:53:18 +0900 Subject: [PATCH 13/28] in clout_at: replace the broadcast with a boolean covered-grid over the window's cell range - one fancy-index lookup per snapshot's points and one rectangle memset after it --- dimos/perception/memory/rig.py | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index cc2e1a1394..96436101f1 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -784,23 +784,22 @@ def cloud_at(self, ts: float) -> PointCloud2 | None: points = np.vstack([p for _t, p in pairs]) else: anchor = nearest[0, :2] + ordered = sorted(pairs, key=lambda pair: abs(pair[0] - ts)) + cells = [ + np.round((part[:, :2] - anchor) / quantum).astype(np.int64) for _t, part in ordered + ] + mins = [c.min(axis=0) for c in cells] + maxs = [c.max(axis=0) for c in cells] + base = np.minimum.reduce(mins) + covered = np.zeros(tuple(np.maximum.reduce(maxs) - base + 1), dtype=bool) kept: list[np.ndarray] = [] - lows: list[np.ndarray] = [] - highs: list[np.ndarray] = [] - for _t, part in sorted(pairs, key=lambda pair: abs(pair[0] - ts)): - cells = np.round((part[:, :2] - anchor) / quantum).astype(np.int64) - if lows: - lo, hi = np.stack(lows), np.stack(highs) - covered = ( - (cells[:, None, :] >= lo[None]) & (cells[:, None, :] <= hi[None]) - ).all(axis=2) - fresh = ~covered.any(axis=1) - if fresh.any(): - kept.append(part[fresh]) - else: - kept.append(part) - lows.append(cells.min(axis=0)) - highs.append(cells.max(axis=0)) + for (_t, part), c, lo, hi in zip(ordered, cells, mins, maxs, strict=True): + fresh = ~covered[c[:, 0] - base[0], c[:, 1] - base[1]] + if fresh.any(): + kept.append(part[fresh]) + covered[ + lo[0] - base[0] : hi[0] + 1 - base[0], lo[1] - base[1] : hi[1] + 1 - base[1] + ] = True points = np.vstack(kept) merged = PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=ts) self._cloud_memo = (ts, merged, quantum) From 26e993f9d57a25d418e320dd71c241ed6ffcea6a Mon Sep 17 00:00:00 2001 From: bogwi Date: Thu, 27 Aug 2026 18:34:03 +0900 Subject: [PATCH 14/28] smal improv to lattice_quantum foo --- .../detection/type/detection3d/pointcloud.py | 6 +-- dimos/perception/memory/rig.py | 39 ++++++++++++------- 2 files changed, 27 insertions(+), 18 deletions(-) diff --git a/dimos/perception/detection/type/detection3d/pointcloud.py b/dimos/perception/detection/type/detection3d/pointcloud.py index 6464c720e7..2388c0e9e4 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -42,9 +42,9 @@ def lattice_quantum(points: np.ndarray) -> float | None: """The grid pitch when coordinates lie on a lattice; None for continuous scans. - Grid-quantized sources (an occupancy map streamed as clouds) carry their - pitch in the data itself; it decides cell dedup, snapshot merging and the - projection splat. Continuous scans never collide and skip all of it. + Grid-quantized sources carry their pitch in the data itself; it sizes + merge cells and the projection splat. Quantization does not classify the + source - a mm-integer wire format grids a scan without making it a map. """ sample = points[:2048] x = np.unique(sample[:, 0]) diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 96436101f1..55267a3bd9 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -23,8 +23,9 @@ unprojected per detection mask, or from a world-frame pointcloud stream (a registered lidar), projected through the camera per detection mask. The cloud at a timestamp merges the scans in a short window around it: - sparse continuous scans concatenate whole, rolling-map snapshots merge - nearest-first with cleared cells honored (see ``Rig.cloud_at``). + fresh scans concatenate whole; snapshots re-reporting each other's exact + points (a rolling map) merge nearest-first with cleared cells honored + (see ``Rig.cloud_at``). ``Rig.from_store`` recognizes both recording shapes; every field can also be supplied directly for live stores whose streams are still filling. @@ -385,7 +386,7 @@ class Rig: scene_gate: bool = True embed_hz: float = EMBED_HZ mobile: bool = False # camera rides a mobile base: room-scale policies - # (ts, merged cloud, the window's measured lattice quantum) + # (ts, merged cloud, the rolling map's pitch; None for scan sources) _cloud_memo: tuple[float, PointCloud2, float | None] | None = field( default=None, repr=False, init=False ) @@ -755,14 +756,16 @@ def cloud_at(self, ts: float) -> PointCloud2 | None: """World-frame geometry at ts: the window's scans merged. The merge rule follows the source's shape, measured per window from - the points themselves. A grid-quantized stream is a rolling - occupancy map: each snapshot is already temporally integrated over - the area it covers and the map clears what moved away, so snapshots - merge nearest-ts first, a farther snapshot contributing only cells + the points themselves. A stream whose next snapshot re-reports the + majority of the nearest one's exact points is a rolling occupancy + map: each snapshot is already temporally integrated over the area + it covers and the map clears what moved away, so snapshots merge + nearest-ts first, a farther snapshot contributing only cells outside the XY coverage of every nearer one - plain accumulation - would resurrect every moved object's trail. Continuous scans are - sparse and the scene static in world frame, so they accumulate - whole. + would resurrect every moved object's trail. Scans that never repeat + are fresh samples of a scene static in world frame, so they + accumulate whole; coordinate quantization alone proves nothing, a + mm-integer wire format grids a scan without making it a map. """ if self._cloud_memo is not None and self._cloud_memo[0] == ts: return self._cloud_memo[1] @@ -776,15 +779,21 @@ def cloud_at(self, ts: float) -> PointCloud2 | None: ] if not pairs: return None - nearest = min(pairs, key=lambda pair: abs(pair[0] - ts))[1] - quantum = lattice_quantum(nearest) - if len(pairs) == 1: - points = pairs[0][1] + ordered = sorted(pairs, key=lambda pair: abs(pair[0] - ts)) + nearest = ordered[0][1] + quantum = None + if len(ordered) > 1: + rows = np.dtype((np.void, nearest.dtype.itemsize * nearest.shape[1])) + near_rows = np.ascontiguousarray(nearest).view(rows).ravel() + next_rows = np.ascontiguousarray(ordered[1][1]).view(rows).ravel() + if np.isin(next_rows, near_rows).mean() > 0.5: + quantum = lattice_quantum(nearest) + if len(ordered) == 1: + points = nearest elif quantum is None: points = np.vstack([p for _t, p in pairs]) else: anchor = nearest[0, :2] - ordered = sorted(pairs, key=lambda pair: abs(pair[0] - ts)) cells = [ np.round((part[:, :2] - anchor) / quantum).astype(np.int64) for _t, part in ordered ] From ff031de66de99efa2220a939f99b78ef46d36836 Mon Sep 17 00:00:00 2001 From: bogwi Date: Thu, 27 Aug 2026 18:52:26 +0900 Subject: [PATCH 15/28] ship go2_short.db.rig.json inside the go2_short LFS archive --- data/.lfs/go2_short.db.tar.gz | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/data/.lfs/go2_short.db.tar.gz b/data/.lfs/go2_short.db.tar.gz index 1fc85989a5..c44813d3e5 100644 --- a/data/.lfs/go2_short.db.tar.gz +++ b/data/.lfs/go2_short.db.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8a19846a0adf5755815fd039492c0255e0bc282e9df75a06648d7585cae8d2d2 -size 83971952 +oid sha256:37949ca60f5048ea675188cb2d9d8b5bf1ea70da3cd045a2e91a30069eccdbc2 +size 83972455 From 8a7e58c04371fb79728923924f31318e512a5928 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:34:09 +0000 Subject: [PATCH 16/28] [autofix.ci] apply automated fixes --- dimos/memory/backend.py | 4 +--- dimos/perception/detection/identity.py | 2 +- .../detection/type/detection3d/pointcloud.py | 12 +++++++----- dimos/perception/memory/rig.py | 8 ++++++-- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/dimos/memory/backend.py b/dimos/memory/backend.py index 1194ba1ca0..3fa7852529 100644 --- a/dimos/memory/backend.py +++ b/dimos/memory/backend.py @@ -211,9 +211,7 @@ def _vector_search(self, query: StreamQuery) -> Iterator[Observation[T]]: continue if not all(f.matches(match) for f in query.filters): continue - ranked.append( - match.derive(data=match.data, embedding=query.search_vec, similarity=sim) - ) + ranked.append(match.derive(data=match.data, embedding=query.search_vec, similarity=sim)) # Apply remaining query ops (filters already applied; skip vector search) rest = replace(query, search_vec=None, search_k=None, filters=()) diff --git a/dimos/perception/detection/identity.py b/dimos/perception/detection/identity.py index f7f8b66bf3..5eb6c9e89d 100644 --- a/dimos/perception/detection/identity.py +++ b/dimos/perception/detection/identity.py @@ -61,7 +61,7 @@ def is_same(a: Detection3DPC, b: Detection3DPC) -> bool: def fused(voxel: float) -> Callable[[Detection3DPC, Detection3DPC], Detection3DPC]: """Merges two detections by combining their point clouds. For each grid cell of size *voxel* meters, - it keeps the mean position of all points grouped in that cell. + it keeps the mean position of all points grouped in that cell. The weight of each fused point is the number of raw points that contributed to it. """ weights: dict[int, np.ndarray[Any, Any]] = {} diff --git a/dimos/perception/detection/type/detection3d/pointcloud.py b/dimos/perception/detection/type/detection3d/pointcloud.py index 2388c0e9e4..d1d2d6eb75 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -270,13 +270,15 @@ def project_pixels(points_camera: np.ndarray, camera_info: CameraInfo) -> np.nda # points past the corner's angle (or past the first fold) are # unmappable. corners = np.array( - [[0, 0], [camera_info.width, 0], [0, camera_info.height], - [camera_info.width, camera_info.height]], + [ + [0, 0], + [camera_info.width, 0], + [0, camera_info.height], + [camera_info.width, camera_info.height], + ], dtype=np.float64, ) - corner_limit = float( - np.hypot((corners[:, 0] - cx) / fx, (corners[:, 1] - cy) / fy).max() - ) + corner_limit = float(np.hypot((corners[:, 0] - cx) / fx, (corners[:, 1] - cy) / fy).max()) theta = np.linspace(0.0, np.pi / 2 * 0.99, 2048) if fisheye: k = np.zeros(4) diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 55267a3bd9..95905535a4 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -199,7 +199,9 @@ def _heading_series(rig: Rig, spans: list[tuple[float, float]]) -> tuple[np.ndar for obs in rig.poses.after(spans[0][0]).before(spans[-1][1]): q = obs.pose_stamped.orientation times.append(obs.ts) - headings.append(np.arctan2(2 * (q.w * q.z + q.x * q.y), 1 - 2 * (q.y * q.y + q.z * q.z))) + headings.append( + np.arctan2(2 * (q.w * q.z + q.x * q.y), 1 - 2 * (q.y * q.y + q.z * q.z)) + ) else: for lo, hi in spans: frame_ts = [obs.ts for obs in rig.color.after(lo).before(hi)] @@ -253,7 +255,9 @@ def estimate_color_delay(rig: Rig) -> float: for obs in rig.color.after(lo).before(hi): frame = obs.data.to_opencv() scale = 640 / frame.shape[1] - gray = cv2.resize(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), (640, int(frame.shape[0] * scale))) + gray = cv2.resize( + cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), (640, int(frame.shape[0] * scale)) + ) if previous is not None and obs.ts > previous_ts: corners = cv2.goodFeaturesToTrack(previous, 150, 0.01, 12) if corners is not None: From 3d151d5fe39e24dcc6527547d1a22364293a97c8 Mon Sep 17 00:00:00 2001 From: bogwi Date: Fri, 28 Aug 2026 11:27:45 +0900 Subject: [PATCH 17/28] fix mypy --- dimos/perception/memory/localize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index c02d66a991..3e757ee0ab 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -147,7 +147,7 @@ def _axes_observed(members: list[Detection3DPC], center: np.ndarray) -> tuple[bo def _lift( - detections: ImageDetections2D, + detections: ImageDetections2D[Any], rig: Rig, policy: LocalizePolicy, plane: Any | None = None, From 55355cd10bb1b262b5cc3935dda7919d8cb8e7b9 Mon Sep 17 00:00:00 2001 From: bogwi Date: Sun, 30 Aug 2026 16:32:20 +0900 Subject: [PATCH 18/28] move to mem api --- dimos/perception/memory/localize.py | 239 +++++++++++++++++----------- 1 file changed, 144 insertions(+), 95 deletions(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 3e757ee0ab..badf05f0f9 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -34,7 +34,7 @@ from dataclasses import dataclass, field import math -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -48,7 +48,10 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from collections.abc import Iterator + from dimos.memory.stream import Stream + from dimos.memory.type.observation import PoseTuple from dimos.models.embedding.siglip import SigLIPModel from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector @@ -68,6 +71,31 @@ def _similarity(obs: Any) -> float: return float(obs.similarity) +def _settled(index: Stream[Any, Any], spacing: float) -> set[int]: + """Ids of the index frames left once sub-spacing duplicates are dropped. + + The index emits the sharpest frame per window from a fixed phase, so a + camera that moves between windows leaves both the settled frame and the + blurred one taken while it was still moving. Their poses differ, so the + blurred one passes as a second viewpoint and lifts to a displaced cloud. + A pair closer than half the index window is one window's content split + by that phase; only the sharper of the two survives here. + """ + ids: set[int] = set() + last_id = -1 + last_ts = -math.inf + last_sharpness = -1.0 + for obs in index.order_by("ts"): + sharpness = float(obs.data.sharpness) + if obs.ts - last_ts < spacing: + if sharpness <= last_sharpness: + continue + ids.discard(last_id) + ids.add(obs.id) + last_id, last_ts, last_sharpness = obs.id, obs.ts, sharpness + return ids + + def _centroid(det: Detection3DPC) -> np.ndarray: centroid: np.ndarray = np.asarray(det.pointcloud.pointcloud.points).mean(axis=0) return centroid @@ -235,9 +263,11 @@ def localize( is flagged below ``refusal_margin`` - never a silent guess. A list *query* shares one detection pass: every label's semantic peaks - select frames, and each unique frame is scored against every label in a - single OWLv2 forward, segmented and lifted once. One instance list per - label, in input order. ``trace`` then takes a list of the same length. + mark its sightings, each peak takes the frames adjacent to it in time + until the verification policy's viewpoints are covered, and each unique + selected frame is scored against every label, segmented and lifted once. + One instance list per label, in input order. ``trace`` then takes a list + of the same length. The index, the rig and the three models belong to the caller: nothing here is loaded or stopped, so one process can call this repeatedly on @@ -260,46 +290,66 @@ def localize( list(trace) if isinstance(trace, list) else [trace] * len(queries) ) - peaks_per_label: list[Stream[Any, Any]] = [] + index_count = index.count() + settled = _settled(index, 0.5 / rig.embed_hz) + source = index.filter(lambda obs: obs.id in settled) + candidate_ids: set[int] = set() + expanded: set[float] = set() + anchor_x = 0.0 + anchor_y = 0.0 + anchor_count = 0 + for q in queries: query_embedding = siglip.embed_text(q) - label_peaks: Stream[Any, Any] = ( - index.search(query_embedding) + sightings = list( + source.search(query_embedding) .order_by("ts") .transform(peaks(key=_similarity, distance=1.0)) - .materialize() ) - logger.info( - f"localize {q!r}: {label_peaks.count()} semantic peaks of {index.count()} embedded" + # A peak is never the window's last sample, and an instance's position + # follows its latest sighting: the tail's best frame is one too. + sightings.extend( + source.after(sightings[-1].ts if sightings else 0.0).search(query_embedding, k=1) ) - peaks_per_label.append(label_peaks) - - frames: dict[float, Any] = {} - expanded: set[float] = set() - for label_peaks in peaks_per_label: - for peak in label_peaks: - frames.setdefault(peak.ts, peak) + peak_count = 0 + for peak in sightings: + peak_pose = cast("PoseTuple", peak.pose_tuple) + peak_count += 1 + anchor_count += 1 + anchor_x += peak_pose[0] + anchor_y += peak_pose[1] + candidate_ids.add(peak.id) if peak.ts in expanded: continue expanded.add(peak.ts) - nearby: Stream[Any, Any] = index.near( + gathered: Stream[Any, Any] = source.near( peak.pose_stamped, radius=policy.verify_radius_m ).transform(QualityWindow(lambda img: img.sharpness, window=0.5)) - for obs in nearby: - frames.setdefault(obs.ts, obs) - ordered = sorted(frames.values(), key=lambda obs: obs.ts) - logger.info(f"detection: {len(ordered)} candidate frames for {len(queries)} labels") + for obs in gathered: + candidate_ids.add(obs.id) + logger.info(f"localize {q!r}: {peak_count} semantic peaks of {index_count} embedded") - if ordered: + candidates = index.filter(lambda obs: obs.id in candidate_ids).order_by("ts") + candidate_count = len(candidate_ids) + logger.info(f"detection: {candidate_count} candidate frames for {len(queries)} labels") + + plane = None + if candidate_count: from dimos.perception.memory.support_plane import fit_support_plane - anchors = [peak.pose_tuple for label_peaks in peaks_per_label for peak in label_peaks] - mx = sum(t[0] for t in anchors) / len(anchors) - my = sum(t[1] for t in anchors) / len(anchors) + mx = anchor_x / anchor_count + my = anchor_y / anchor_count cell = (round(mx / 2.0), round(my / 2.0)) plane = rig._plane_cache.get(cell) if plane is None: - plane = fit_support_plane(rig, ordered) + stride = max(1, candidate_count // 5) + keyframes = [] + for i, obs in enumerate(candidates): + if i % stride == 0: + keyframes.append(obs) + if len(keyframes) == 5: + break + plane = fit_support_plane(rig, keyframes) if plane is not None: rig._plane_cache[cell] = plane @@ -324,76 +374,75 @@ def localize( floor = policy.candidate_floor cache = detector.score_cache - misses = [ - obs - for obs in ordered - if any( - obs.ts not in ingested[j] and (obs.ts, q, floor) not in cache - for j, q in enumerate(queries) - ) - ] - if misses: - scored = detector.query_score_rows_batch( - [obs.data for obs in misses], queries, threshold=floor - ) - for obs, (boxes, rows) in zip(misses, scored, strict=True): - for j, q in enumerate(queries): - keep = rows[:, j] >= floor - cache[(obs.ts, q, floor)] = (boxes[keep], rows[keep, j]) - if len(cache) > _SCORE_CACHE_MAX: - cache.popitem(last=False) - - for obs in ordered: - active = [j for j in range(len(queries)) if obs.ts not in ingested[j]] - if not active: - continue - rows_per_label: list[tuple[int, tuple[Any, Any]]] = [] - for j in active: - key = (obs.ts, queries[j], floor) - cache.move_to_end(key) - rows_per_label.append((j, cache[key])) - ingested[j].add(obs.ts) - if not any(len(scores) for _j, (_boxes, scores) in rows_per_label): - continue - img = obs.data - candidates: list[Detection2DBBox] = [] - for j, (boxes, scores) in rows_per_label: - for box, score in zip(boxes, scores, strict=True): - det = Detection2DBBox( - bbox=(float(box[0]), float(box[1]), float(box[2]), float(box[3])), - track_id=len(candidates), - class_id=j, - confidence=float(score), - name=queries[j], - ts=img.ts, - image=img, - ) - if det.is_valid(): - candidates.append(det) - if not candidates: - continue - frame = segmenter.segment(ImageDetections2D(image=img, detections=candidates)) - lifted = _lift(frame, rig, policy, plane) - grounded = {det3d.track_id for det3d in lifted} - for det2d in frame: - j = det2d.class_id - best = ungrounded[j] - if det2d.track_id not in grounded and (best is None or det2d.confidence > best[0]): - ungrounded[j] = (det2d.confidence, det2d.ts) - for det3d in lifted: - label_trace = traces[det3d.class_id] - if label_trace is not None: - label_trace.matched.append((det3d.ts, det3d)) - identities[det3d.class_id].add(det3d) - for j, label_trace in enumerate(traces): - if label_trace is None: + def _detect(upstream: Iterator[Any]) -> Iterator[Any]: + for obs in upstream: + active = [j for j in range(len(queries)) if obs.ts not in ingested[j]] + if not active: continue - label_dets = [det for det in frame if det.class_id == j] - if label_dets: - label_trace.detection_frames.append( - obs.derive(data=ImageDetections2D(image=obs.data, detections=label_dets)) - ) + if any((obs.ts, queries[j], floor) not in cache for j in active): + boxes, rows = detector.query_score_rows_batch([obs.data], queries, threshold=floor)[ + 0 + ] + for j, q in enumerate(queries): + keep = rows[:, j] >= floor + cache[(obs.ts, q, floor)] = (boxes[keep], rows[keep, j]) + if len(cache) > _SCORE_CACHE_MAX: + cache.popitem(last=False) + + rows_per_label: list[tuple[int, tuple[Any, Any]]] = [] + for j in active: + key = (obs.ts, queries[j], floor) + cache.move_to_end(key) + rows_per_label.append((j, cache[key])) + ingested[j].add(obs.ts) + if not any(len(scores) for _j, (_boxes, scores) in rows_per_label): + continue + img = obs.data + detections: list[Detection2DBBox] = [] + for j, (boxes, scores) in rows_per_label: + for box, score in zip(boxes, scores, strict=True): + det = Detection2DBBox( + bbox=(float(box[0]), float(box[1]), float(box[2]), float(box[3])), + track_id=len(detections), + class_id=j, + confidence=float(score), + name=queries[j], + ts=img.ts, + image=img, + ) + if det.is_valid(): + detections.append(det) + if detections: + yield obs.derive(data=ImageDetections2D(image=img, detections=detections)) + + def _ingest(upstream: Iterator[Any]) -> Iterator[Any]: + for obs in upstream: + frame = obs.data + lifted = _lift(frame, rig, policy, plane) + grounded = {det3d.track_id for det3d in lifted} + for det2d in frame: + j = det2d.class_id + best = ungrounded[j] + if det2d.track_id not in grounded and (best is None or det2d.confidence > best[0]): + ungrounded[j] = (det2d.confidence, det2d.ts) + for det3d in lifted: + label_trace = traces[det3d.class_id] + if label_trace is not None: + label_trace.matched.append((det3d.ts, det3d)) + identities[det3d.class_id].add(det3d) + for j in range(len(queries)): + label_trace = traces[j] + if label_trace is None: + continue + label_dets = [det for det in frame if det.class_id == j] + if label_dets: + label_trace.detection_frames.append( + obs.derive(data=ImageDetections2D(image=frame.image, detections=label_dets)) + ) + yield obs + + candidates.transform(_detect).map(lambda obs: obs.derive(data=segmenter.segment(obs.data))).transform(_ingest).drain() if entries is not None: for entry, best in zip(entries, ungrounded, strict=True): From 091537d889a3c42a35bd5002d644df6db1774118 Mon Sep 17 00:00:00 2001 From: bogwi Date: Sun, 30 Aug 2026 18:15:06 +0900 Subject: [PATCH 19/28] test(detection): re-record OBB expectations for the corrected projection test_detection3dpc pinned obb.center and obb.extent to values from when from_2d projected through a bare pinhole matrix and ignored camera_info.D. The camera under test is go2_front_camera_720p, declared equidistant. project_pixels now applies that model, shifting projected pixels by a median of 20 px and moving the suitcase OBB center 18 mm in y, just past the 0.1 window. Tolerances and the file's other assertions are unchanged. --- .../detection/type/detection3d/test_pointcloud.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/dimos/perception/detection/type/detection3d/test_pointcloud.py b/dimos/perception/detection/type/detection3d/test_pointcloud.py index c9e62af6de..37d852a255 100644 --- a/dimos/perception/detection/type/detection3d/test_pointcloud.py +++ b/dimos/perception/detection/type/detection3d/test_pointcloud.py @@ -26,14 +26,14 @@ def test_detection3dpc(detection3dpc) -> None: assert obb is not None, "Oriented bounding box should not be None" # Verify OBB center values - assert obb.center[0] == pytest.approx(-3.36002, abs=0.1) - assert obb.center[1] == pytest.approx(-0.196446, abs=0.1) - assert obb.center[2] == pytest.approx(0.220184, abs=0.1) + assert obb.center[0] == pytest.approx(-3.316207, abs=0.1) + assert obb.center[1] == pytest.approx(-0.300175, abs=0.1) + assert obb.center[2] == pytest.approx(0.240114, abs=0.1) # Verify OBB extent values - assert obb.extent[0] == pytest.approx(0.531275, abs=0.12) - assert obb.extent[1] == pytest.approx(0.461054, abs=0.1) - assert obb.extent[2] == pytest.approx(0.155, abs=0.1) + assert obb.extent[0] == pytest.approx(0.593476, abs=0.12) + assert obb.extent[1] == pytest.approx(0.470315, abs=0.1) + assert obb.extent[2] == pytest.approx(0.164996, abs=0.1) # def test_bounding_box_dimensions(detection3dpc): """Test bounding box dimension calculation.""" From 533e0bcf1328b4d7ff56274fc42afd530cf4f115 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:52:10 +0000 Subject: [PATCH 20/28] [autofix.ci] apply automated fixes --- dimos/perception/memory/localize.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index badf05f0f9..886cd56b92 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -442,7 +442,9 @@ def _ingest(upstream: Iterator[Any]) -> Iterator[Any]: ) yield obs - candidates.transform(_detect).map(lambda obs: obs.derive(data=segmenter.segment(obs.data))).transform(_ingest).drain() + candidates.transform(_detect).map( + lambda obs: obs.derive(data=segmenter.segment(obs.data)) + ).transform(_ingest).drain() if entries is not None: for entry, best in zip(entries, ungrounded, strict=True): From 96ce3a6225afc87549f84b9f5b974a994817f29e Mon Sep 17 00:00:00 2001 From: bogwi Date: Mon, 31 Aug 2026 12:31:12 +0900 Subject: [PATCH 21/28] add LocalizePolicy for the localize caller. The caller can tune any localize call in runtime --- dimos/perception/memory/dandetect.py | 10 +++- dimos/perception/memory/localize.py | 23 ++++++--- dimos/perception/memory/rig.py | 21 +++++--- dimos/perception/memory/types.py | 77 +++++++++++++++++++++++----- 4 files changed, 104 insertions(+), 27 deletions(-) diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index 0553d7533d..77ec79b973 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -41,7 +41,7 @@ from dimos.models.embedding.siglip import SigLIPModel from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector - from dimos.perception.memory.types import Instance, Localization + from dimos.perception.memory.types import Instance, Localization, LocalizePolicy class DanDetector(Resource): @@ -140,9 +140,14 @@ def localize( query: str | list[str], *, index: Stream[Any, Any], + policy: LocalizePolicy | None = None, **kwargs: Any, ) -> list[Localization] | list[list[Localization]]: - """:func:`localize` on this resource's models.""" + """:func:`localize` on this resource's models. + + ``policy`` is the localize thresholds. ``None`` uses the rig's scale + defaults. + """ return localize( store, query, @@ -150,6 +155,7 @@ def localize( siglip=self.siglip, detector=self.detector, segmenter=self.segmenter, + policy=policy, **kwargs, ) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 886cd56b92..99e4df5524 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -291,7 +291,7 @@ def localize( ) index_count = index.count() - settled = _settled(index, 0.5 / rig.embed_hz) + settled = _settled(index, policy.settled_window_fraction / rig.embed_hz) source = index.filter(lambda obs: obs.id in settled) candidate_ids: set[int] = set() expanded: set[float] = set() @@ -304,12 +304,21 @@ def localize( sightings = list( source.search(query_embedding) .order_by("ts") - .transform(peaks(key=_similarity, distance=1.0)) + .transform( + peaks( + key=_similarity, + prominence=policy.peak_prominence, + distance=policy.peak_distance_s, + width=policy.peak_width_s, + ) + ) ) # A peak is never the window's last sample, and an instance's position # follows its latest sighting: the tail's best frame is one too. sightings.extend( - source.after(sightings[-1].ts if sightings else 0.0).search(query_embedding, k=1) + source.after(sightings[-1].ts if sightings else 0.0).search( + query_embedding, k=policy.tail_k + ) ) peak_count = 0 for peak in sightings: @@ -324,7 +333,7 @@ def localize( expanded.add(peak.ts) gathered: Stream[Any, Any] = source.near( peak.pose_stamped, radius=policy.verify_radius_m - ).transform(QualityWindow(lambda img: img.sharpness, window=0.5)) + ).transform(QualityWindow(lambda img: img.sharpness, window=policy.verify_window_s)) for obs in gathered: candidate_ids.add(obs.id) logger.info(f"localize {q!r}: {peak_count} semantic peaks of {index_count} embedded") @@ -339,15 +348,15 @@ def localize( mx = anchor_x / anchor_count my = anchor_y / anchor_count - cell = (round(mx / 2.0), round(my / 2.0)) + cell = (round(mx / policy.plane_cell_m), round(my / policy.plane_cell_m)) plane = rig._plane_cache.get(cell) if plane is None: - stride = max(1, candidate_count // 5) + stride = max(1, candidate_count // policy.plane_keyframes) keyframes = [] for i, obs in enumerate(candidates): if i % stride == 0: keyframes.append(obs) - if len(keyframes) == 5: + if len(keyframes) == policy.plane_keyframes: break plane = fit_support_plane(rig, keyframes) if plane is not None: diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 95905535a4..1b59cb2ece 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -133,13 +133,13 @@ def _column_keys(points: np.ndarray, quantum: float, anchor: np.ndarray) -> np.n candidate_floor=0.18, accept_score=0.32, cluster_radius_m=0.30, - fuse_voxel_m=0.03, - min_depth_points=30, + verify_radius_m=5.0, max_object_extent_m=2.0, - min_camera_range_m=0.5, surface_patch_max_rise_m=0.08, surface_patch_min_drop_m=-0.06, - verify_radius_m=5.0, + min_depth_points=30, + min_camera_range_m=0.5, + fuse_voxel_m=0.03, ) @@ -150,7 +150,7 @@ def _tf_root(store: Any, tf_name: str) -> str | None: """The tf tree's root frame: a parent that is never a child.""" parents: set[str] = set() children: set[str] = set() - for obs in store.stream(tf_name).limit(500): + for obs in store.stream(tf_name): for transform in obs.data.transforms: parents.add(transform.frame_id) children.add(transform.child_frame_id) @@ -464,12 +464,15 @@ def from_store( color_candidates: list[str] = [] depth_candidates: list[str] = [] empty_images: list[str] = [] + image_frames: dict[str, str] = {} for name in image_names: stream = store.stream(name) if stream.count() == 0: empty_images.append(name) continue - frame = stream.first().data.data + image = stream.first().data + image_frames[name] = image.frame_id + frame = image.data if frame.dtype == np.uint16 or frame.dtype.kind == "f": depth_candidates.append(name) elif ( @@ -480,6 +483,12 @@ def from_store( color_candidates.append(name) else: logger.info(f"rig: image stream {name!r} is neither color nor metric depth") + if color_name is None and depth_name is not None: + depth_frame = store.stream(depth_name).first().data.frame_id + matching = [name for name in color_candidates if image_frames[name] == depth_frame] + if len(matching) == 1: + color_name = matching[0] + logger.info(f"rig: paired depth {depth_name!r} with color {color_name!r}") if color_name is None: if len(color_candidates) > 1: color_name = max(color_candidates, key=lambda n: _stream_rate(store.stream(n))) diff --git a/dimos/perception/memory/types.py b/dimos/perception/memory/types.py index eedb99cfb4..0d7e30b8bb 100644 --- a/dimos/perception/memory/types.py +++ b/dimos/perception/memory/types.py @@ -112,24 +112,77 @@ class LocalizePolicy: and height distributions of one measured scene, so a different rig, object scale or detector vocabulary needs its own instance rather than the defaults. + + `candidate_floor`: OWLv2 score at which a box is formed for a query. + Boxes below this never lift. + + `accept_score`: A support group is returned when its highest member score meets this. + The RGB-only path uses the same floor. + + `min_views`: Unique camera positions, rounded to 1 cm, required before a group is confirmed. + A support seen from one pose only is dropped. + + `cluster_radius_m`: Two lifted detections are the same support when their cloud centers sit within this many meters. + + `peak_prominence`: Minimum SigLIP similarity rise for a semantic peak. + A local maximum below this is not a peak. + + `peak_distance_s`: Minimum seconds between two semantic peaks. + + `peak_width_s`: Minimum peak width in seconds at half prominence. + ``None`` disables the width gate. + + `verify_radius_m`: Index frames whose pose is within this many meters of a peak are gathered for OWLv2. + + `verify_window_s`: Keep the sharpest gathered frame per this many seconds. + + `settled_window_fraction`: Collapse index frames closer than this fraction of ``1/embed_hz`` to the sharper one, so one window does not count as two viewpoints. + + `tail_k`: After the last peak, take this many extra frames by query similarity so the window tail can be a latest sighting. + + `max_object_extent_m`: Drop a lift whose longest AABB edge exceeds this. + + `surface_patch_max_rise_m`: Drop a lift whose 95th-percentile height above the support is below this, when the drop test also holds. + A cloud that hugs the support is a patch of the surface, not an object. + + `surface_patch_min_drop_m`: Drop a lift whose 5th-percentile height is above this, when the rise test also holds. + + `min_depth_points`: Minimum points on a depth lift. + Projected-cloud rigs ignore this. + + `min_camera_range_m`: Drop a lift whose median point-to-camera range is below this. + + `fuse_voxel_m`: Voxel size of the union cloud at identity merge. + ``0`` concatenates. + + `plane_cell_m`: XY cell size for the support-plane cache. + + `plane_keyframes`: Candidate frames sampled to fit the support plane. + + `refusal_margin`: If this instance's score minus the best coexisting rival is below this, ``reason`` is set. + The instance is still returned. """ - candidate_floor: float = 0.25 # form a candidate at this score + candidate_floor: float = 0.25 accept_score: float = 0.40 - refusal_margin: float = 0.15 - min_views: int = 2 # a support seen from one pose only is unconfirmed - - cluster_radius_m: float = 0.08 # observations within this are the same support - fuse_voxel_m: float = 0.01 # union-cloud voxel at the identity merge; 0 concatenates - min_depth_points: int = 60 + min_views: int = 2 + cluster_radius_m: float = 0.08 + peak_prominence: float = 0.02 + peak_distance_s: float = 1.0 + peak_width_s: float | None = 0.5 + verify_radius_m: float = 1.6 + verify_window_s: float = 0.5 + settled_window_fraction: float = 0.5 + tail_k: int = 1 max_object_extent_m: float = 0.60 - min_camera_range_m: float = 0.28 - # A cloud that hugs the support surface is a patch of the surface, not an - # object: every real object rises above the plane, a surface patch does not. surface_patch_max_rise_m: float = 0.003 surface_patch_min_drop_m: float = -0.02 - # Images gathered around each semantic peak for the detection pass. - verify_radius_m: float = 1.6 + min_depth_points: int = 60 + min_camera_range_m: float = 0.28 + fuse_voxel_m: float = 0.01 + plane_cell_m: float = 2.0 + plane_keyframes: int = 5 + refusal_margin: float = 0.15 @dataclass(frozen=True) From 57d11cb4ded508577601563e3e796c9222932183 Mon Sep 17 00:00:00 2001 From: bogwi Date: Mon, 31 Aug 2026 16:27:26 +0900 Subject: [PATCH 22/28] enhanse rig.py --- dimos/perception/memory/rig.py | 59 +++++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 19 deletions(-) diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 1b59cb2ece..3288bf671b 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -146,15 +146,35 @@ def _column_keys(points: np.ndarray, quantum: float, anchor: np.ndarray) -> np.n MOBILE_SPAN_M = 3.0 # camera translation beyond this means a mobile base -def _tf_root(store: Any, tf_name: str) -> str | None: - """The tf tree's root frame: a parent that is never a child.""" - parents: set[str] = set() - children: set[str] = set() +ROOT_PROBES = 24 # instants a frame is probed at before it counts as unreachable + + +def _tf_root(store: Any, tf_name: str, tf: StreamTF, optical_frame: str) -> str | None: + """The tf tree's root frame: a parent that is never a child, among the + frames the camera actually reaches. + + A recording can carry an anchor edge published once at each end of the + run, above the frame every other transform is stamped in. It is a root the + camera never reaches through, and taking it strands every pose lookup, so + frames no probe resolves are dropped before the root is taken. Probes sit + at bin midpoints, where an anchor stamped at the ends cannot answer. + """ + edges: set[tuple[str, str]] = set() for obs in store.stream(tf_name): for transform in obs.data.transforms: - parents.add(transform.frame_id) - children.add(transform.child_frame_id) - roots = parents - children + edges.add((transform.frame_id, transform.child_frame_id)) + if not edges: + return None # live store, nothing recorded yet + frames = {frame for edge in edges for frame in edge} + t0, t1 = store.stream(tf_name).get_time_range() + reached = {optical_frame} + for k in range(ROOT_PROBES): + ts = t0 + (t1 - t0) * (k + 0.5) / ROOT_PROBES + for frame in frames - reached: + if tf.get(optical_frame, frame, ts, TF_TOLERANCE, warn=False) is not None: + reached.add(frame) + linked = [(p, c) for p, c in edges if p in reached and c in reached] + roots = {p for p, _ in linked} - {c for _, c in linked} return roots.pop() if len(roots) == 1 else None @@ -517,15 +537,6 @@ def from_store( if cloud_name is not None: claimed.add(cloud_name) - world_frame = gates.WORLD_FRAME - if tf is not None: - world_frame = _tf_root(store, cast("str", tf_name)) or world_frame - elif cloud is not None: - try: - world_frame = cloud.first().data.frame_id - except LookupError: - pass # live store, nothing recorded yet - # intrinsics: inline manifest dict, named stream, or discovery by # type with the color camera's frame deciding among several camera_info = None @@ -563,6 +574,19 @@ def from_store( if ci_name is not None: camera_info = store.stream(ci_name).first().data + # embed-only stores (no geometry) may carry no calibration at all; + # every geometry API raises on use, embedding never touches it + optical_frame = camera_info.frame_id if camera_info is not None else gates.OPTICAL_FRAME + + world_frame = gates.WORLD_FRAME + if tf is not None: + world_frame = _tf_root(store, cast("str", tf_name), tf, optical_frame) or world_frame + elif cloud is not None: + try: + world_frame = cloud.first().data.frame_id + except LookupError: + pass # live store, nothing recorded yet + base_to_optical = None mount = roles.get("base_to_optical") if isinstance(mount, dict): @@ -600,9 +624,6 @@ def from_store( "base_to_optical mount in the .rig.json manifest" ) - # embed-only stores (no geometry) may carry no calibration at all; - # every geometry API raises on use, embedding never touches it - optical_frame = camera_info.frame_id if camera_info is not None else gates.OPTICAL_FRAME rig = cls( camera_info=cast("CameraInfo", camera_info), color=color, From b32dd57dcdcc8515cd653fbfb5bf8c607796c114 Mon Sep 17 00:00:00 2001 From: bogwi Date: Tue, 1 Sep 2026 12:07:16 +0900 Subject: [PATCH 23/28] add localize live blueprint --- .../memory/blueprints/go2_localize_live.py | 566 ++++++++++++++++++ dimos/robot/all_blueprints.py | 3 + 2 files changed, 569 insertions(+) create mode 100644 dimos/perception/memory/blueprints/go2_localize_live.py diff --git a/dimos/perception/memory/blueprints/go2_localize_live.py b/dimos/perception/memory/blueprints/go2_localize_live.py new file mode 100644 index 0000000000..80fb23d11e --- /dev/null +++ b/dimos/perception/memory/blueprints/go2_localize_live.py @@ -0,0 +1,566 @@ +# 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. + +"""A go2 recording looped as a live feed, with localize behind an agent skill. + +This is the deployment shape of the perception memory stack: one module owns +a store and the models, ``DanDetector.embed(live=True)`` keeps a background +tail filling the index while the robot runs, and a ``@skill`` answers +``localize`` calls against whatever has been embedded so far. + +:class:`LoopFeeder` stands in for the sensors and the recorder. It replays a +recording into a fresh store forever, adding the recording's own span to +every timestamp each lap, so the store's clock only moves forward while the +scene repeats. The message stamp, the observation stamp and every tf +transform shift by the same lap offset, so pose interpolation, cloud +accumulation and the index window all see one continuous take. It writes the +tf chain and the camera intrinsics too, which is what lets +:class:`LocalizeModule` resolve its rig from the live store with no manifest. + +The canonical recording is opened read-only and never written. +""" + +from __future__ import annotations + +from dataclasses import replace +import heapq +import json +from pathlib import Path +import threading +import time +from typing import TYPE_CHECKING, Any +import zlib + +from dimos_lcm.geometry_msgs import Pose +from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose +import numpy as np + +from dimos.agents.annotation import skill +from dimos.agents.mcp.mcp_server import McpServer +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.core import rpc +from dimos.core.global_config import global_config +from dimos.core.stream import Out +from dimos.mapping.voxels.module import VoxelGridMapper +from dimos.memory.module import MemoryModule, MemoryModuleConfig +from dimos.memory.replay import resolve_db_path +from dimos.memory.store.sqlite import SqliteStore +from dimos.memory.tf import StreamTF +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.vision_msgs.Detection3D import Detection3D +from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.perception.detection.type.detection3d.pointcloud import ( + Detection3DPC, + lattice_quantum, +) +from dimos.perception.memory.dandetect import DanDetector +from dimos.perception.memory.identity_store import IdentityStore +from dimos.perception.memory.rig import Rig +from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL, GO2Connection +from dimos.utils.logging_config import setup_logger +from dimos.visualization.vis_module import vis_module + +if TYPE_CHECKING: + from collections.abc import Iterator + +logger = setup_logger() + +LAP_GAP_S = 0.2 # dead time between laps so consecutive laps stay disjoint +WARMUP_WINDOW_S = 5.0 # feed the warmup query a short window, only to load weights +UNSEEN_GREY = 0.25 # voxels outside the camera frustum carry no pixel +# The map is drawn achromatic so anything coloured - the camera texture, a +# detection box - reads as foreground against it. Lightness carries height, +# dark at the floor, the way shaded relief is read. +MAP_INK = (45, 215) +# Cube edge as a fraction of the source's own lattice pitch. Full pitch +# tiles into a solid block that hides the shape behind it; a gap between +# cubes reads as structure while the voxel count stays the same. +CUBE_FILL = 0.55 + + +class LoopFeederConfig(MemoryModuleConfig): + dataset: str = "go2_short" + db_path: str | Path = "recording_go2_live.db" + + +class LoopFeeder(MemoryModule): + """Replay ``dataset`` into this module's store forever, one lap at a time. + + Roles are read straight off the recording; tf and camera_info are derived + the way :class:`GO2Connection` derives them live, from each odom pose and + the static front-camera calibration. + """ + + config: LoopFeederConfig + + color_image: Out[Image] + lidar: Out[PointCloud2] + odom: Out[PoseStamped] + camera_info: Out[CameraInfo] + tf: Out[TFMessage] + + ROLES = ("color_image", "lidar", "odom", "tf", "camera_info", "color_image_embedded") + + @rpc + def start(self) -> None: + super().start() + # Drop what a previous run left, in place: the localize module holds + # the same file open, so the streams go rather than the inode. + for name in self.ROLES: + self.store.delete_stream(name) + + self._stop = threading.Event() + self._thread = threading.Thread(target=self._feed, name="loop-feeder", daemon=True) + self._thread.start() + + @rpc + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=5.0) + super().stop() + + @staticmethod + def _stamped(stream: Any, name: str) -> Iterator[tuple[float, str, Any]]: + """The stream's observations tagged with their role, for the lap merge.""" + for obs in stream: + yield (obs.ts, name, obs) + + def _feed(self) -> None: + source = SqliteStore(path=str(resolve_db_path(self.config.dataset)), must_exist=True) + source.start() + lo, hi = source.stream("color_image").get_time_range() + span = (hi - lo) + LAP_GAP_S + + live = self.store + targets: dict[str, Any] = { + name: live.stream(name, source.stream(name).data_type) + for name in ("color_image", "lidar", "odom") + } + tf_stream = live.stream("tf", TFMessage) + intrinsics = GO2Connection.camera_info_static + live.stream("camera_info", CameraInfo).append(intrinsics, ts=lo, pose=None) + + ports: dict[str, Any] = { + "color_image": self.color_image, + "lidar": self.lidar, + "odom": self.odom, + } + + # The embed tail runs here, in the process that writes the frames. + # SubjectNotifier fans out in-process only, so a tail subscribed from + # another worker backfills once and then never sees another append. + self._embedder = self.register_disposable(DanDetector()) + self._embedder.start() + self._embedder.embed(live, live=True, rig=_live_rig(Rig.from_store(source), live)) + logger.info( + f"loop feeder: {self.config.dataset} ({span - LAP_GAP_S:.1f}s) -> {self.config.db_path}" + ) + + lap = 0 + while not self._stop.is_set(): + offset = lap * span + wall_start = time.time() + self.camera_info.publish(intrinsics) + schedule: Any = heapq.merge( + *(self._stamped(source.stream(name), name) for name in targets), + key=lambda item: item[0], + ) + image: Image | None = None + base: PoseStamped | None = None + for ts, name, obs in schedule: + if self._stop.is_set(): + return + self._stop.wait(max(0.0, (ts - lo) - (time.time() - wall_start))) + data = obs.data + data.ts = ts + offset + if name == "odom": + data.frame_id = "world" + base = data + transforms = GO2Connection._odom_to_tf(data) + tf_stream.append(TFMessage(*transforms), ts=data.ts, pose=None) + self.tf.publish(TFMessage(*transforms)) + elif name == "color_image": + image = data + targets[name].append(data, ts=data.ts, pose=obs.pose) + # The store keeps the raw scan; only the viewer sees the texture. + if name == "lidar" and image is not None and base is not None: + self.lidar.publish(_textured(data, image, base)) + else: + ports[name].publish(data) + lap += 1 + logger.info(f"loop feeder: lap {lap} done, feed clock at +{lap * span:.0f}s") + + +def _textured(cloud: PointCloud2, image: Image, pose: PoseStamped) -> PointCloud2: + """The scan re-emitted with every voxel carrying the pixel it projects onto. + + The scan is already in the world frame, so it goes through the camera the + way a detection's cloud does: world to optical through the base pose and + the static mount, then the camera's own distortion model. Voxels behind + the camera or outside the image keep a neutral grey. + """ + import open3d as o3d + import open3d.core as o3c + + points = cloud.points_f32() + matrix = (-(Transform.from_pose("base_link", pose) + BASE_TO_OPTICAL)).to_matrix() + camera = points @ matrix[:3, :3].T + matrix[:3, 3] + rgb = image.to_rgb().data + height, width = rgb.shape[:2] + + colors = np.full((len(points), 3), UNSEEN_GREY, dtype=np.float32) + ahead = np.flatnonzero(camera[:, 2] > 0) + if len(ahead): + pixels = Detection3DPC.project_pixels(camera[ahead], GO2Connection.camera_info_static) + cols = np.round(pixels[:, 0]).astype(int) + rows = np.round(pixels[:, 1]).astype(int) + inside = (cols >= 0) & (cols < width) & (rows >= 0) & (rows < height) + colors[ahead[inside]] = rgb[rows[inside], cols[inside]] / 255.0 + + pcd = o3d.t.geometry.PointCloud() + pcd.point["positions"] = o3c.Tensor(points, dtype=o3c.float32) + pcd.point["colors"] = o3c.Tensor(colors, dtype=o3c.float32) + return PointCloud2(pointcloud=pcd, frame_id=cloud.frame_id, ts=cloud.ts) + + +class LocalizeModuleConfig(MemoryModuleConfig): + dataset: str = "go2_short" + db_path: str | Path = "recording_go2_live.db" + warmup_query: str = "zijnh" + + +def _live_rig(source: Rig, live: Any) -> Rig: + """The recording's rig, reading the live store's streams instead. + + A robot knows its rig from calibration; it does not rediscover it at + runtime. Taking the shape from the recording is also what lets startup + skip waiting for enough recorded motion for the mobile gate to settle, + and carries the measured color delay over instead of re-estimating it. + """ + return Rig( + camera_info=source.camera_info, + color=live.stream("color_image", source.color.data_type), + world_frame=source.world_frame, + optical_frame=source.optical_frame, + tf=StreamTF(live.stream("tf", TFMessage)) if source.tf is not None else None, + base_to_optical=source.base_to_optical, + poses=live.stream("odom", source.poses.data_type) if source.poses is not None else None, + cloud=live.stream("lidar", source.cloud.data_type) if source.cloud is not None else None, + tf_tolerance=source.tf_tolerance, + cloud_accum_s=source.cloud_accum_s, + speed_max=source.speed_max, + color_delay=source.color_delay, + scene_gate=source.scene_gate, + embed_hz=source.embed_hz, + mobile=source.mobile, + ) + + +class LocalizeModule(MemoryModule): + """Live object memory: a background embed tail plus a localize skill. + + Startup loads every weight against the recording itself, so it overlaps + the feed filling instead of following it, then opens the live tail and + answers as soon as the first frames are embedded. ``state`` reports where + it is; ``localize`` names the same stage when asked too early. + + Answers are published as a :class:`Detection3DArray`, which the rerun + bridge draws as labelled boxes. Nothing is written back to the recording. + """ + + config: LocalizeModuleConfig + + detections: Out[Detection3DArray] + + @rpc + def start(self) -> None: + super().start() + self._ready = threading.Event() + self._stage = "starting" + self._identity = IdentityStore() + self._thread = threading.Thread(target=self._warm, name="localize-warmup", daemon=True) + self._thread.start() + + @rpc + def stop(self) -> None: + self._ready.clear() + super().stop() + + def _warm(self) -> None: + source = SqliteStore(path=str(resolve_db_path(self.config.dataset)), must_exist=True) + source.start() + + self._stage = "loading SigLIP, OWLv2 and EdgeTAM weights" + logger.info(f"localize: {self._stage}") + recorded = Rig.from_store(source) + self.detector = self.register_disposable(DanDetector()) + self.detector.start() + + # Every model loads here, on this thread, against the recording, so + # none of it waits on the feed. DanDetector.start() only constructs + # the wrappers; the models themselves arrive on first use, and the + # tail would otherwise construct SigLIP at the same instant as the + # first query, leaving one of the two threads a half-loaded copy. + lo, _ = recorded.color.get_time_range() + warm = self.detector.embed(source, lo, lo + WARMUP_WINDOW_S, rig=recorded) + self.detector.localize(source, self.config.warmup_query, index=warm, rig=recorded) + + self._stage = "waiting for the first frames of feed" + logger.info(f"localize: {self._stage}") + self.rig = _live_rig(recorded, self.store) + # The feeder's process owns the embed tail; this one only reads it. + self.index = self.store.stream("color_image_embedded", Image) + while self.index.count() == 0: + time.sleep(0.5) + + self._stage = "ready" + self._ready.set() + logger.info(f"localize: ready on {self.index.count()} embedded frames") + + @skill + def state(self) -> str: + """Whether localize can answer yet, and what it is doing if not.""" + if self._ready.is_set(): + return f"ready: {self.index.count()} frames embedded, localize will answer" + return f"not ready: {self._stage}" + + @skill + def localize( + self, objects: str, start: float = -10.0, duration: float = 10.0, policy: str = "" + ) -> str: + """Locate objects in a window of the robot's memory. + + ``objects`` is one label, or several separated by commas. + + ``start`` and ``duration`` are seconds and name the window, the way + ``--from`` and ``--duration`` do on tool_localize. A positive + ``start`` counts forward from the beginning of the feed; a negative + one counts back from the newest frame, like a negative Python index, + so the default reads the last ten seconds. The window is what this + call examines; what earlier calls proved is remembered and still + answered, so coverage accumulates while the cost per call does not. + + ``policy`` is a JSON object of LocalizePolicy field overrides, + e.g. '{"accept_score": 0.4, "verify_radius_m": 2.0}'. + """ + if not self._ready.is_set(): + return f"localize cannot answer yet: {self._stage}. Poll state() until it reads ready." + + queries = [q.strip() for q in objects.split(",") if q.strip()] + first, head = self.index.get_time_range() + # A window reaching past the start of the feed begins at the start. + lo = max(first, head + start if start < 0 else first + start) + index = self.index.time_range(lo, lo + duration) + tuning = self.rig.default_localize_policy() + if policy: + tuning = replace(tuning, **json.loads(policy)) + results: Any = self.detector.localize( + self.store, + queries, + index=index, + rig=self.rig, + policy=tuning, + identity_store=self._identity, + ) + self.detections.publish(_as_detection_array(queries, results, self.rig.world_frame)) + + lines: list[str] = [ + f"window {lo - first:.1f}s to {lo + duration - first:.1f}s of " + f"{head - first:.1f}s of feed, {index.count()} frames" + ] + for query, hits in zip(queries, results, strict=True): + if not hits: + lines.append(f"no verified detection of {query!r}") + for hit in hits: + x, y, z = hit.position_world_xyz + lines.append( + f"{query!r} at ({x:.2f}, {y:.2f}, {z:.2f}) in {hit.frame_id} " + f"score={hit.semantic_score:.2f} views={hit.n_views}" + ) + return "\n".join(lines) + + +def _as_detection_array(queries: list[str], results: list[Any], frame_id: str) -> Detection3DArray: + """One labelled box per verified instance, for the rerun bridge.""" + boxes = [] + latest = 0.0 + for query, hits in zip(queries, results, strict=True): + for hit in hits: + if hit.point_cloud is None: + continue + points = hit.point_cloud.as_numpy()[0] + low, high = points.min(axis=0), points.max(axis=0) + middle, extent = (low + high) / 2, high - low + center = Vector3(*(float(v) for v in middle)) + size = Vector3(*(max(float(v), 1e-3) for v in extent)) + latest = max(latest, hit.last_seen_timestamp) + boxes.append( + Detection3D( + header=Header(hit.last_seen_timestamp, frame_id), + id=hit.instance_id, + results=[ + ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id=query, score=hit.semantic_score) + ) + ], + results_length=1, + bbox=BoundingBox3D( + center=Pose(position=center, orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + size=size, + ), + ) + ) + return Detection3DArray( + detections_length=len(boxes), + header=Header(latest, frame_id), + detections=boxes, + ) + + +def _static_robot_body(rr: Any) -> list[Any]: + return [ + rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), + rr.Transform3D(parent_frame="tf#/base_link"), + ] + + +def _cube_size(points: np.ndarray) -> dict[str, float]: + """``voxel_size`` for a grid source, or nothing for a continuous one.""" + quantum = lattice_quantum(points) + return {"voxel_size": quantum * CUBE_FILL} if quantum else {} + + +def _convert_global_map(grid: Any) -> Any: + """The accumulated lidar as cubes shaded dark-to-light by height. + + Whole map, floor included - a height cutoff at the world origin would + drop it, since the world frame is anchored at the robot's start pose and + the ground sits below it. Height is normalized over the map's own extent, + which only grows, so the ramp settles as the map fills. + """ + points = grid.points_f32() + if not len(points): + return grid.to_rerun(mode="boxes") + z = points[:, 2] + low, high = MAP_INK + lightness = low + (z - z.min()) / (z.max() - z.min() + 1e-8) * (high - low) + grey = np.repeat(lightness.astype(np.uint8)[:, None], 3, axis=1) + return grid.to_rerun(mode="boxes", colors=grey, **_cube_size(points)) + + +def _label_colour(label: str) -> tuple[int, int, int]: + """A saturated colour that follows the label, not the call. + + Derived from the text so one label keeps its colour across queries and + windows; two labels in view are two colours without reading either. + """ + import colorsys + + hue = (zlib.crc32(label.encode()) % 3600) / 3600.0 + return tuple(int(c * 255) for c in colorsys.hsv_to_rgb(hue, 0.75, 1.0)) # type: ignore[return-value] + + +def _detection_entities(array: Any) -> Any: + """One rerun entity per label, so each owns a row on the timeline.""" + import rerun as rr + + grouped: dict[str, list[Any]] = {} + for detection in array.detections[: array.detections_length]: + label = str(detection.results[0].hypothesis.class_id) + grouped.setdefault(label, []).append(detection) + + entities = [] + for label, members in grouped.items(): + boxes = [d.bbox for d in members] + entities.append( + ( + f"world/detections/{label.replace(' ', '_')}", + rr.Boxes3D( + centers=[ + (b.center.position.x, b.center.position.y, b.center.position.z) + for b in boxes + ], + half_sizes=[(b.size.x / 2, b.size.y / 2, b.size.z / 2) for b in boxes], + labels=[label] * len(boxes), + colors=[_label_colour(label)] * len(boxes), + ), + ) + ) + return entities + + +def _lidar_cubes(cloud: Any) -> Any: + """The lidar lattice as solid voxel cubes wearing the camera's pixels.""" + _, colors = cloud.as_numpy() + rgb = (colors * 255).astype(np.uint8) if colors is not None else None + return cloud.to_rerun(mode="boxes", colors=rgb, **_cube_size(cloud.points_f32())) + + +def _convert_camera_info(camera_info: Any) -> Any: + return camera_info.to_rerun( + image_topic="/world/color_image", + optical_frame="camera_optical", + ) + + +def _rerun_blueprint() -> Any: + """Camera feed, 3D world, and the localize answers on top of both.""" + import rerun as rr + import rerun.blueprint as rrb + + return rrb.Blueprint( + rrb.Horizontal( + rrb.Spatial2DView(origin="world/color_image", name="Camera"), + rrb.Spatial3DView( + origin="world", + name="3D", + background=rrb.Background(kind="SolidColor", color=[0, 0, 0]), + line_grid=rrb.LineGrid3D(plane=rr.components.Plane3D.XY.with_distance(0.5)), + ), + column_shares=[1, 2], + ), + rrb.TimePanel(state="expanded"), + rrb.SelectionPanel(state="hidden"), + ) + + +rerun_config: dict[str, Any] = { + "blueprint": _rerun_blueprint, + "visual_override": { + "world/camera_info": _convert_camera_info, + "world/lidar": _lidar_cubes, + "world/global_map": _convert_global_map, + "world/detections": _detection_entities, + }, + "max_hz": {"world/color_image": 0, "world/lidar": 1, "world/global_map": 0}, + "tf_axes": 0.5, + "static": {"world/robot_body": _static_robot_body}, +} + + +go2_localize_live = autoconnect( + vis_module(viewer_backend=global_config.viewer, rerun_config=rerun_config), + LoopFeeder.blueprint(), + VoxelGridMapper.blueprint(emit_every=5), + LocalizeModule.blueprint(), + McpServer.blueprint(), +).global_config(n_workers=7, robot_model="unitree_go2") diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 671a8baf57..b6dff5f55d 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -59,6 +59,7 @@ "drone-agentic": "dimos.robot.drone.blueprints.agentic.drone_agentic:drone_agentic", "drone-basic": "dimos.robot.drone.blueprints.basic.drone_basic:drone_basic", "dual-xarm6-planner-coordinator": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator", + "go2-localize-live": "dimos.perception.memory.blueprints.go2_localize_live:go2_localize_live", "go2-zenoh-basic": "dimos.robot.unitree.go2.zenoh.blueprints:go2_zenoh_basic", "go2-zenoh-htc": "dimos.robot.unitree.go2.zenoh.blueprints:go2_zenoh_htc", "go2-zenoh-nav": "dimos.robot.unitree.go2.zenoh.blueprints:go2_zenoh_nav", @@ -223,6 +224,8 @@ "joystick-module": "dimos.robot.unitree.b1.joystick_module.JoystickModule", "keyboard-teleop": "dimos.robot.unitree.keyboard_teleop.KeyboardTeleop", "keyboard-teleop-module": "dimos.teleop.keyboard.keyboard_teleop_module.KeyboardTeleopModule", + "localize-module": "dimos.perception.memory.blueprints.go2_localize_live.LocalizeModule", + "loop-feeder": "dimos.perception.memory.blueprints.go2_localize_live.LoopFeeder", "manipulation-module": "dimos.manipulation.manipulation_module.ManipulationModule", "manipulation-skills": "dimos.manipulation.manipulation_skills.ManipulationSkills", "map": "dimos.robot.unitree.type.map.Map", From 3e104869573f917c8117480bcbfe3125518dcef8 Mon Sep 17 00:00:00 2001 From: bogwi Date: Tue, 1 Sep 2026 13:58:27 +0900 Subject: [PATCH 24/28] profile memory for rerun --- dimos/perception/memory/blueprints/go2_localize_live.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/dimos/perception/memory/blueprints/go2_localize_live.py b/dimos/perception/memory/blueprints/go2_localize_live.py index 80fb23d11e..4c67216d7f 100644 --- a/dimos/perception/memory/blueprints/go2_localize_live.py +++ b/dimos/perception/memory/blueprints/go2_localize_live.py @@ -551,7 +551,13 @@ def _rerun_blueprint() -> Any: "world/global_map": _convert_global_map, "world/detections": _detection_entities, }, - "max_hz": {"world/color_image": 0, "world/lidar": 1, "world/global_map": 0}, + # Rerun keeps every frame it is logged, so an unthrottled entity is an + # unbounded allocation: the bridge grows with the feed, not with the + # scene. The map is a backdrop that only has to be current, not live. + "max_hz": {"world/color_image": 10, "world/lidar": 1, "world/global_map": 0.2}, + # A share of the machine rather than a number, and small enough that the + # viewer drops its own history long before the host runs out. + "memory_limit": "5%", "tf_axes": 0.5, "static": {"world/robot_body": _static_robot_body}, } From 886adbb4f66af05e9a55fbe5882ed3878587279b Mon Sep 17 00:00:00 2001 From: bogwi Date: Tue, 1 Sep 2026 15:02:30 +0900 Subject: [PATCH 25/28] make mcp call caller dependent; preserve the default 30 sec rule --- dimos/agents/mcp/mcp_adapter.py | 12 ++- dimos/cli/commands/mcp.py | 9 +- dimos/core/global_config.py | 4 + dimos/perception/memory/blueprints/README.md | 93 ++++++++++++++++++++ 4 files changed, 108 insertions(+), 10 deletions(-) create mode 100644 dimos/perception/memory/blueprints/README.md diff --git a/dimos/agents/mcp/mcp_adapter.py b/dimos/agents/mcp/mcp_adapter.py index 213bf71e23..16254b7374 100644 --- a/dimos/agents/mcp/mcp_adapter.py +++ b/dimos/agents/mcp/mcp_adapter.py @@ -41,8 +41,6 @@ logger = setup_logger() -DEFAULT_TIMEOUT = 30 - class McpError(Exception): """Raised when the MCP server returns a JSON-RPC error.""" @@ -55,13 +53,13 @@ def __init__(self, message: str, code: int | None = None) -> None: class McpAdapter: """Thin JSON-RPC client for a running MCP server.""" - def __init__(self, url: str | None = None, timeout: int = DEFAULT_TIMEOUT) -> None: - if url is None: - from dimos.core.global_config import global_config + def __init__(self, url: str | None = None, timeout: int | None = None) -> None: + from dimos.core.global_config import global_config + if url is None: url = f"http://localhost:{global_config.mcp_port}/mcp" self.url = url - self.timeout = timeout + self.timeout = global_config.mcp_timeout if timeout is None else timeout def call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: """Send a JSON-RPC request and return the parsed response. @@ -137,7 +135,7 @@ def wait_for_down(self, timeout: float = 10.0, interval: float = 0.5) -> bool: return False @classmethod - def from_run_entry(cls, entry: Any | None = None, timeout: int = DEFAULT_TIMEOUT) -> McpAdapter: + def from_run_entry(cls, entry: Any | None = None, timeout: int | None = None) -> McpAdapter: """Create an adapter from a RunEntry, or discover the latest one. Falls back to the default URL if no entry is found. diff --git a/dimos/cli/commands/mcp.py b/dimos/cli/commands/mcp.py index b190719c41..e82d6dd73e 100644 --- a/dimos/cli/commands/mcp.py +++ b/dimos/cli/commands/mcp.py @@ -27,9 +27,9 @@ mcp_app = typer.Typer(help="Interact with the running MCP server") -def _get_adapter() -> McpAdapter: +def _get_adapter(timeout: int | None = None) -> McpAdapter: """Get an McpAdapter from the latest RunEntry or default URL.""" - return McpAdapter.from_run_entry() + return McpAdapter.from_run_entry(timeout=timeout) @mcp_app.command("list-tools") @@ -72,6 +72,9 @@ def mcp_call_tool( [], "--arg", "-a", callback=_validate_key_value_args, help="Arguments as key=value" ), json_args: str = typer.Option("", "--json-args", "-j", help="Arguments as JSON string"), + timeout: int = typer.Option( + None, "--timeout", "-t", help="Seconds to wait for the tool (default: mcp_timeout)" + ), ) -> None: """Call an MCP tool by name.""" arguments: dict[str, Any] = {} @@ -89,7 +92,7 @@ def mcp_call_tool( raise typer.Exit(1) try: - result = _get_adapter().call_tool(tool_name, arguments) + result = _get_adapter(timeout).call_tool(tool_name, arguments) except requests.ConnectionError: typer.echo("Error: no running MCP server (is DimOS running?)", err=True) raise typer.Exit(1) diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index 8b5e448157..a9504972ce 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -104,6 +104,10 @@ class GlobalConfig(BaseSettings): robot_rotation_diameter: float = 0.6 nerf_speed: float = 1.0 mcp_port: int = 9990 + # Seconds an MCP client waits for a tool to answer. A skill that thinks + # for longer than this is cut off at the client, not the server, so the + # caller owns the number. + mcp_timeout: int = 30 # `DIMOS_TRANSPORT` (or `.env`) is the single switch read by every process # (dimos, humancli, agentspy, dtop). The `transport` alias keeps the bare # env name and the `--transport` CLI flag (which sets the field by name) working. diff --git a/dimos/perception/memory/blueprints/README.md b/dimos/perception/memory/blueprints/README.md new file mode 100644 index 0000000000..96a69845c6 --- /dev/null +++ b/dimos/perception/memory/blueprints/README.md @@ -0,0 +1,93 @@ +# go2-localize-live + +This blueprint runs the perception memory stack live. It takes the `go2_short` +recording and replays it into a fresh store forever, lap after lap, with the +clock always moving forward. A background tail embeds the frames as they +arrive, so the store keeps growing like a real robot's recorder would. You then +ask it where objects are, and it answers from what it has seen so far. The +answers show up in rerun as labelled boxes, on top of the lidar map. Nothing is +written back to the original recording. + +Start it in terminal 1: + + uv run dimos run go2-localize-live + +Wait for this line before you ask it anything: + + localize: ready on N embedded frames + +It takes about 15 seconds. Most of that is loading SigLIP, OWLv2 and EdgeTAM. + +## REPL + +Open terminal 2: + + uv run dimos shell + +Then: + + modules() # what is running + rpcs("LocalizeModule") # what you can call + describe("LocalizeModule.localize") # the signature and the docs + + app.LocalizeModule.state() # ready or not, and what it is doing + app.LocalizeModule.localize("chair") + app.LocalizeModule.localize("chair,table,green plant") + +`localize(objects, start, duration, policy)`. + +`objects` is one label, or several split by commas. Several labels share one +detection pass, so 16 labels cost about the same as one. + +`start` and `duration` are seconds, and they name the window this call looks +at. They work like `--from` and `--duration` on `tool_localize`. Positive +`start` counts forward from the beginning of the feed. Negative counts back +from the newest frame, like a negative Python index. The default is the last +ten seconds. + + app.LocalizeModule.localize("chair", -30.0, 30.0) # last 30 seconds + app.LocalizeModule.localize("chair", 5.0, 10.0) # 10 seconds, from 5s in + +The window is only what this call examines. What earlier calls proved is +remembered and still answered, so coverage builds up while the cost per call +does not. + +A wide window over time you have never asked about is the expensive case. Every +frame in it needs a detector pass. The same window asked twice is nearly free. + +`policy` is a JSON object. It overrides any field of `LocalizePolicy` for this +one call: + + app.LocalizeModule.localize("chair", -10.0, 10.0, + '{"accept_score": 0.45, "verify_radius_m": 2.0}') + +## Terminal + +From any terminal, while it runs: + + uv run dimos status # is it up + uv run dimos stop # stop it + uv run dimos log # its logs + +## MCP + +The same two calls are skills, so an agent can reach them over MCP: + + uv run dimos mcp list-tools + uv run dimos mcp call state + uv run dimos mcp call localize -a 'objects=chair,table' + uv run dimos mcp call localize -j '{"objects": "chair", "start": -30, "duration": 30}' + +The client waits 30 seconds for an answer and then gives up. A wide window can +take longer than that. The wait belongs to the caller, so raise it there. Per +call, or for every call you make: + + uv run dimos mcp call localize -a 'objects=chair' --timeout 300 + uv run dimos --mcp-timeout 300 mcp call localize -a 'objects=chair' + MCP_TIMEOUT=300 uv run dimos mcp call localize -a 'objects=chair' + +Setting it on the blueprint does nothing. The timeout is on the side that +waits, not the side that answers. + +Ask `state` first. Until it says ready, `localize` tells you which stage it is +in instead of answering. From fcd5b706d8f5a944a25126615922de1c241498c0 Mon Sep 17 00:00:00 2001 From: bogwi Date: Thu, 3 Sep 2026 18:00:31 +0900 Subject: [PATCH 26/28] remove inventory api --- dimos/perception/memory/dandetect.py | 27 +- dimos/perception/memory/inventory.py | 902 ---------------------- dimos/perception/memory/tool_inventory.py | 333 -------- 3 files changed, 5 insertions(+), 1257 deletions(-) delete mode 100644 dimos/perception/memory/inventory.py delete mode 100644 dimos/perception/memory/tool_inventory.py diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index 77ec79b973..c524761123 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -14,8 +14,8 @@ """One disposable resource wrapping the memory perception API. -``DanDetector`` owns the models behind :func:`embed_index`, :func:`localize`, -and :func:`inventory`: enter once, query many times on warm weights, and +``DanDetector`` owns the models behind :func:`embed_index` and +:func:`localize`: enter once, query many times on warm weights, and ``stop()`` (or leave the ``with`` block) releases whatever loaded. Every entry point takes an optional :class:`~dimos.perception.memory.rig.Rig` @@ -30,7 +30,6 @@ from dimos.core.resource import Resource from dimos.memory.embed import EmbedImages from dimos.memory.transform import QualityWindow -from dimos.perception.memory.inventory import DEFAULT_VOCABULARY, NamingVocabulary, inventory from dimos.perception.memory.localize import embed_index, localize from dimos.perception.memory.rig import Rig @@ -41,15 +40,15 @@ from dimos.models.embedding.siglip import SigLIPModel from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector - from dimos.perception.memory.types import Instance, Localization, LocalizePolicy + from dimos.perception.memory.types import Localization, LocalizePolicy class DanDetector(Resource): """The perception models as one resource. ``start()`` constructs SigLIP, OWLv2, and EdgeTAM. The two - HuggingFace models load lazily on first use, so an inventory-only - caller never pays for SigLIP; ``stop()`` releases whatever loaded. + HuggingFace models load lazily on first use. ``stop()`` releases + whatever loaded. """ siglip: SigLIPModel @@ -158,19 +157,3 @@ def localize( policy=policy, **kwargs, ) - - def inventory( - self, - store: Any, - *, - naming_vocabulary: NamingVocabulary = DEFAULT_VOCABULARY, - **kwargs: Any, - ) -> list[Instance]: - """:func:`inventory` on this resource's models.""" - return inventory( - store, - segmenter=self.segmenter, - detector=self.detector, - naming_vocabulary=naming_vocabulary, - **kwargs, - ) diff --git a/dimos/perception/memory/inventory.py b/dimos/perception/memory/inventory.py deleted file mode 100644 index 59a988bf02..0000000000 --- a/dimos/perception/memory/inventory.py +++ /dev/null @@ -1,902 +0,0 @@ -# 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. - -"""Query-time scene inventory: prompt-free discovery plus geometric dedup. - -``inventory()`` answers "what instances are on the table" for a time window, -computed at query time over the recording - no ingest pass, no persisted -instance table. Existence is decoupled from naming and the ordering is a -constraint, not a preference: propose (EdgeTAM automatic masks), lift -(masked depth to world supports), associate (hard constraints before any -score), and only then name (OWLv2, labels as metadata). Labels and -appearance never enter association; position and same-frame co-occurrence -decide everything, which is what keeps two identical objects two instances. - -Naming is passed by a caller and abstains: it reads the full per-box score row, -groups surface strings under a canonical label, and reports that label only -when it beats the runner-up group by a margin, at the frame and again over -the frames of a track. The word list is not what makes the output safe, so -it can be domain-specific, generic, or empty. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from time import perf_counter -from typing import TYPE_CHECKING, Any - -import numpy as np - -from dimos.perception.memory.gates import MOTION_THRESHOLD -from dimos.perception.memory.rig import Rig -from dimos.perception.memory.support_plane import SupportPlane, fit_support_plane -from dimos.perception.memory.types import ( - Instance, - InventoryPolicy, - Support, - SupportObservation, - aabb_overlap, -) -from dimos.utils.logging_config import setup_logger - -if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo - - from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter - from dimos.perception.detection.detectors.owlv2 import Owlv2Detector - from dimos.perception.detection.type.detection2d.seg import Detection2DSeg - -logger = setup_logger() - -MAX_PROPOSALS_PER_FRAME = 40 -NAME_FRAMES_PER_INSTANCE = 5 -SUPPRESS_SCORE = 0.25 -SUPPRESS_OVERLAP = 0.35 -UNGROUNDED_TRACK_IOU = 0.40 - -# Groups of surface strings for one thing, canonical label first. Only groups -# compete, so near-synonyms reinforce instead of splitting a box's score. A -# string in two groups pins their margin at zero and both refuse. -NamingVocabulary = tuple[tuple[str, ...], ...] - -# Candidate names for callers with no domain list of their own. -DEFAULT_VOCABULARY: NamingVocabulary = ( - ("pen", "ballpoint pen", "ink pen"), - ("pencil", "wooden pencil", "mechanical pencil"), - ("marker", "marker pen", "felt-tip pen", "permanent marker"), - ("highlighter", "highlighter pen"), - ("eraser", "rubber eraser"), - ("book", "hardcover book", "paperback book", "textbook"), - ("notebook", "spiral notebook", "notepad", "writing pad"), - ("sticky notes", "post-it notes", "sticky note pad", "pad of sticky notes"), - ("sheet of paper", "piece of paper", "printed page", "document"), - ("roll of tape", "adhesive tape", "sticky tape", "roll of duct tape"), - ("scissors", "pair of scissors", "shears"), - ("stapler", "desk stapler"), - ("ruler", "measuring ruler", "straightedge"), - ("laptop", "laptop computer", "notebook computer"), - ("computer keyboard", "keyboard", "laptop keyboard"), - ("computer mouse", "mouse", "wireless mouse"), - ("mobile phone", "smartphone", "cell phone"), - ("cup", "mug", "coffee mug", "drinking cup"), - ("bottle", "water bottle", "plastic bottle"), - ("drink can", "soda can", "aluminum can"), - ("bowl", "small bowl", "cereal bowl"), - ("cardboard box", "carton", "small box"), - ("cable", "power cable", "usb cable", "cord"), - ("remote control", "tv remote", "remote"), - ("glasses", "eyeglasses", "pair of glasses", "spectacles"), - ("headphones", "earphones", "headset"), - ("wallet", "billfold", "leather wallet"), - ("toy block", "building block", "foam block"), -) -# An existence policy, not a candidate name: its own request, its own accept. -SUPPRESS_QUERIES = ["person", "human hand", "human arm"] - - -@dataclass -class _Track: - members: list[SupportObservation] = field(default_factory=list) - frame_ts: set[float] = field(default_factory=set) - labels: dict[str, float] = field(default_factory=dict) - support_pts: np.ndarray = field(default_factory=lambda: np.empty((0, 3))) - - def add(self, obs: SupportObservation, frame_key: float) -> None: - self.members.append(obs) - self.frame_ts.add(frame_key) - points = np.asarray(obs.cloud.pointcloud.points) - self.support_pts = np.vstack([self.support_pts, points[:: max(1, len(points) // 400)]]) - - @property - def centroid(self) -> np.ndarray: - median: np.ndarray = np.median(np.stack([m.centroid for m in self.members]), axis=0) - return median - - @property - def aabb(self) -> tuple[np.ndarray, np.ndarray]: - lo = np.median(np.stack([m.aabb_min for m in self.members]), axis=0) - hi = np.median(np.stack([m.aabb_max for m in self.members]), axis=0) - return lo, hi - - @property - def latest(self) -> SupportObservation: - return max(self.members, key=lambda m: m.ts) - - -@dataclass -class _Track2D: - """Ungrounded track: RGB detections with no valid depth, 2D identity only.""" - - members: list[Detection2DSeg] = field(default_factory=list) - frame_ts: set[float] = field(default_factory=set) - labels: dict[str, float] = field(default_factory=dict) - - -def _bbox_iou(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float: - ax1, ay1, ax2, ay2 = a - bx1, by1, bx2, by2 = b - ix = max(0.0, min(ax2, bx2) - max(ax1, bx1)) - iy = max(0.0, min(ay2, by2) - max(ay1, by1)) - inter = ix * iy - union = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - inter - return inter / union if union > 0 else 0.0 - - -def _proposal_passes_2d(det: Detection2DSeg, image_area: float, policy: InventoryPolicy) -> bool: - area = float((det.mask > 0).sum()) - if area < policy.min_mask_area_px: - return False - if area > policy.max_mask_area_fraction * image_area: - return False - return True - - -def _split_oversized( - points: np.ndarray, plane: SupportPlane | None, policy: InventoryPolicy -) -> list[np.ndarray]: - """Re-segment a mask-bled cloud by 3D connectivity. - - Automatic masks occasionally bleed across an object onto its support - surface and its neighbors; the lifted cloud then violates single-object - bounds. The repair is geometric: strip the support-surface points, then - split by spatial connectivity - distinct objects are separated by more - than the cluster gap, one object's surface is not. - """ - extent = points.max(axis=0) - points.min(axis=0) - if float(extent.max()) <= policy.split_extent_m and float(extent[2]) <= policy.split_height_m: - return [points] - if plane is None: - return [points] - heights = plane.height_above(points) - if float((np.abs(heights) <= policy.min_height_above_plane_m).mean()) < 0.15: - # No appreciable support-surface content: this is one oversized body, - # not a mask that bled across the surface. Leave it to the extent cap. - return [points] - - above = heights > policy.min_height_above_plane_m * 2 / 3 - body = points[above] if above.sum() >= policy.min_depth_points else points - - import open3d as o3d - - cloud = o3d.geometry.PointCloud() - cloud.points = o3d.utility.Vector3dVector(body) - labels = np.asarray(cloud.cluster_dbscan(eps=policy.split_eps_m, min_points=20)) - clusters = [ - body[labels == label] - for label in range(labels.max() + 1) - if (labels == label).sum() >= policy.min_depth_points - ] - return clusters if clusters else [body] - - -def _pixel_bbox( - points: np.ndarray, camera_info: CameraInfo, transform: Any -) -> tuple[float, float, float, float]: - """Project world points back into the frame for a sub-observation's bbox.""" - matrix = transform.to_matrix() - optical = (matrix[:3, :3] @ points.T).T + matrix[:3, 3] - z = np.maximum(optical[:, 2], 1e-6) - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] - u = fx * optical[:, 0] / z + cx - v = fy * optical[:, 1] / z + cy - return (float(u.min()), float(v.min()), float(u.max()), float(v.max())) - - -def _lift_frame( - detections_2d: Any, - rig: Rig, - obs_ts: float, - camera_position: np.ndarray, - policy: InventoryPolicy, - plane: SupportPlane | None = None, -) -> tuple[list[SupportObservation], list[Detection2DSeg]]: - """Lift accepted proposals of one frame through the rig; (grounded, ungrounded).""" - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - - lifted = rig.lift(detections_2d) - transform = rig.world_to_optical(obs_ts) - if lifted is None or transform is None: - return [], list(detections_2d) - camera_info = rig.camera_info - - grounded: list[SupportObservation] = [] - ungrounded: list[Detection2DSeg] = [] - lifted_by_track = {det3d.track_id: det3d for det3d in lifted} - for det2d in detections_2d: - det3d = lifted_by_track.get(det2d.track_id) - if det3d is None or len(det3d.pointcloud) < policy.min_depth_points: - ungrounded.append(det2d) - continue - points = np.asarray(det3d.pointcloud.pointcloud.points) - mask_area = int((det2d.mask > 0).sum()) - for piece in _split_oversized(points, plane, policy): - aabb_min, aabb_max = piece.min(axis=0), piece.max(axis=0) - extent = aabb_max - aabb_min - if float(extent.max()) > policy.max_object_extent_m: - continue - ranges = np.linalg.norm(piece - camera_position, axis=1) - if float(np.median(ranges)) < policy.min_camera_range_m: - continue - whole = len(piece) == len(points) - grounded.append( - SupportObservation( - ts=obs_ts, - cloud=det3d.pointcloud - if whole - else PointCloud2.from_numpy(piece, frame_id=rig.world_frame, timestamp=obs_ts), - centroid=piece.mean(axis=0), - aabb_min=aabb_min, - aabb_max=aabb_max, - n_points=len(piece), - mask_area_px=mask_area if whole else int(mask_area * len(piece) / len(points)), - camera_position=camera_position, - bbox=det2d.bbox if whole else _pixel_bbox(piece, camera_info, transform), - ) - ) - return grounded, ungrounded - - -def _in_scope(obs: SupportObservation, plane: SupportPlane | None, policy: InventoryPolicy) -> bool: - """Support-plane scope: in a band above the plane, footprint on the workspace.""" - if policy.in_scope is not None: - points = np.asarray(obs.cloud.pointcloud.points) - return bool(policy.in_scope(points)) - if plane is None: - return True - points = np.asarray(obs.cloud.pointcloud.points) - heights = plane.height_above(points) - low, high = float(np.quantile(heights, 0.05)), float(np.quantile(heights, 0.95)) - band_lo, band_hi = policy.band_above_plane_m - if low < band_lo or high > band_hi: - return False - if not policy.include_surfaces and high < policy.min_height_above_plane_m: - # A patch of the surface itself: no volume above the plane. - return False - inside = plane.footprint_contains(points[:, :2]) - return bool(inside.mean() >= 0.3) - - -def _aabb_gap(a: SupportObservation, b: SupportObservation) -> float: - """Largest per-axis separation between two world AABBs (0 when touching).""" - gap = np.maximum(a.aabb_min - b.aabb_max, b.aabb_min - a.aabb_max) - return float(gap.max()) - - -def _cloud_gap(a: SupportObservation, b: SupportObservation, cut: float) -> float: - """Minimum point-to-point distance between two observation clouds. - - The AABB gap is a poor contact test for diagonal objects - an - axis-aligned box overhangs its object's true footprint and "touches" - neighbors that are centimeters of clear table away. Actual cloud - distance is the physical claim. The AABB test remains as a cheap - prefilter: beyond ``cut`` of box separation the exact distance cannot - matter to any caller. - """ - if _aabb_gap(a, b) > cut: - return np.inf - from scipy.spatial import cKDTree - - pa = np.asarray(a.cloud.pointcloud.points) - pb = np.asarray(b.cloud.pointcloud.points) - pa = pa[:: max(1, len(pa) // 800)] - pb = pb[:: max(1, len(pb) // 800)] - distances, _ = cKDTree(pa).query(pb, k=1) - return float(distances.min()) - - -def _support_explained(points: np.ndarray, support: np.ndarray, pad: float) -> float: - """Fraction of ``points`` lying within ``pad`` of the accumulated support.""" - from scipy.spatial import cKDTree - - sample = points[:: max(1, len(points) // 800)] - tree = cKDTree(support[:: max(1, len(support) // 4000)]) - distances, _ = tree.query(sample, k=1) - return float((distances <= pad).mean()) - - -def _absorb_into(target: SupportObservation, obs: SupportObservation) -> None: - target.aabb_min = np.minimum(target.aabb_min, obs.aabb_min) - target.aabb_max = np.maximum(target.aabb_max, obs.aabb_max) - target.cloud = target.cloud + obs.cloud - points = np.asarray(target.cloud.pointcloud.points) - target.centroid = points.mean(axis=0) - target.n_points = len(points) - target.mask_area_px = max(target.mask_area_px, obs.mask_area_px) - - -def _merge_same_frame( - observations: list[SupportObservation], policy: InventoryPolicy -) -> list[SupportObservation]: - """Fuse same-frame proposals that are one physical support. - - The criterion is contact: two same-frame observations whose clouds - touch (minimum cloud distance within the gap) are one rigid body - - duplicates, whole-and-part pairs, and split halves of one object all - satisfy it, while distinct objects on the workspace, identical twins - included, sit farther apart than the gap. Runs to a fixed point so - chains of touching pieces collapse into one support. - """ - if policy.include_object_parts: - return observations - items = sorted(observations, key=lambda o: -o.n_points) - changed = True - while changed: - changed = False - for i in range(len(items)): - for j in range(i + 1, len(items)): - if ( - _cloud_gap(items[i], items[j], 3 * policy.same_frame_merge_gap_m) - <= policy.same_frame_merge_gap_m - ): - _absorb_into(items[i], items[j]) - items.pop(j) - changed = True - break - if changed: - break - return items - - -def _associate( - frames: list[tuple[float, list[SupportObservation]]], policy: InventoryPolicy -) -> list[_Track]: - """Hard constraints first, geometric score second, Hungarian for the residual. - - Per frame, observations assign one-to-one to existing tracks - the - same-frame constraint is structural, no score overrides it. A pair is - forbidden outright (infinite cost) when the supports are farther apart - than the search radius, their envelopes do not overlap enough, their - sizes are incompatible beyond measurement error, or the observation is - not majority-explained by the track's accumulated support. - """ - from scipy.optimize import linear_sum_assignment - - forbidden = 1e6 - tracks: list[_Track] = [] - for frame_key, observations in frames: - if not observations: - continue - if not tracks: - for obs in observations: - track = _Track() - track.add(obs, frame_key) - tracks.append(track) - continue - - cost = np.full((len(observations), len(tracks)), forbidden) - for i, obs in enumerate(observations): - obs_points = np.asarray(obs.cloud.pointcloud.points) - for j, track in enumerate(tracks): - distance = float(np.linalg.norm(obs.centroid - track.centroid)) - if distance > policy.search_radius_m: - continue - t_lo, t_hi = track.aabb - size_gap = np.abs((t_hi - t_lo) - (obs.aabb_max - obs.aabb_min)) - if float(size_gap.max()) > policy.size_gap_max_m: - continue - overlap = aabb_overlap( - obs.aabb_min, obs.aabb_max, t_lo, t_hi, pad=policy.envelope_pad_m - ) - if overlap < policy.overlap_accept: - continue - explained = _support_explained(obs_points, track.support_pts, policy.envelope_pad_m) - if explained < policy.support_explained: - continue - cost[i, j] = 1.0 - overlap - - rows, cols = linear_sum_assignment(cost) - assigned: dict[int, int] = {} - for i, j in zip(rows, cols, strict=False): - if cost[i, j] < forbidden: - assigned[i] = j - for i, obs in enumerate(observations): - match = assigned.get(i) - if match is not None: - tracks[match].add(obs, frame_key) - else: - track = _Track() - track.add(obs, frame_key) - tracks.append(track) - return tracks - - -def _tracks_are_fragments(a: _Track, b: _Track, policy: InventoryPolicy) -> bool: - """Co-observed tracks that were touching whenever seen together. - - The same-frame veto keeps coexisting objects apart, but a mask split can - put two pieces of one object into the same frame with a cloud gap just - over the per-frame merge threshold, locking a permanent duplicate. Two - rigid objects cannot occupy one volume: when the supports interpenetrate - at containment level and every shared frame shows the pair in contact, - they are pieces of one body. Identical twins never satisfy this - their - supports do not overlap at all. - """ - a_lo, a_hi = a.aabb - b_lo, b_hi = b.aabb - overlap = aabb_overlap(a_lo, a_hi, b_lo, b_hi, pad=policy.envelope_pad_m) - if overlap < 0.5: - return False - shared = a.frame_ts & b.frame_ts - for ts in shared: - pairs_gap = min( - _cloud_gap(ma, mb, 3 * policy.same_frame_merge_gap_m) - for ma in a.members - if ma.ts == ts - for mb in b.members - if mb.ts == ts - ) - if pairs_gap > 1.5 * policy.same_frame_merge_gap_m: - return False - return True - - -def _merge_tracks(tracks: list[_Track], policy: InventoryPolicy) -> list[_Track]: - """Collapse fragmented tracks of one support. - - Tracks merge when they never share a frame (the same-frame veto at - instance level), their supports overlap within the envelope and either - accumulated support majority-explains the other - or when they do share - frames but were demonstrably pieces of one body in every one of them. - Runs to a fixed point. - """ - changed = True - while changed: - changed = False - for i in range(len(tracks)): - for j in range(i + 1, len(tracks)): - a, b = tracks[i], tracks[j] - if float(np.linalg.norm(a.centroid - b.centroid)) > policy.search_radius_m: - continue - if a.frame_ts & b.frame_ts: - if not _tracks_are_fragments(a, b, policy): - continue - else: - a_lo, a_hi = a.aabb - b_lo, b_hi = b.aabb - overlap = aabb_overlap(a_lo, a_hi, b_lo, b_hi, pad=policy.envelope_pad_m) - if overlap < policy.overlap_accept: - continue - explained = max( - _support_explained(a.support_pts, b.support_pts, policy.envelope_pad_m), - _support_explained(b.support_pts, a.support_pts, policy.envelope_pad_m), - ) - if explained < policy.support_explained: - continue - for obs in b.members: - a.add(obs, obs.ts) - tracks.pop(j) - changed = True - break - if changed: - break - return tracks - - -def _track_ungrounded( - frames: list[tuple[float, list[Detection2DSeg]]], -) -> list[_Track2D]: - """Greedy 2D IoU association for detections that never produced depth.""" - tracks: list[_Track2D] = [] - for frame_key, detections in frames: - for det in detections: - best, best_iou = None, UNGROUNDED_TRACK_IOU - for track in tracks: - if frame_key in track.frame_ts: - continue - iou = _bbox_iou(det.bbox, track.members[-1].bbox) - if iou > best_iou: - best, best_iou = track, iou - if best is None: - best = _Track2D() - tracks.append(best) - best.members.append(det) - best.frame_ts.add(frame_key) - return [t for t in tracks if len(t.members) >= 2] - - -def _view_coverage(members: list[SupportObservation]) -> tuple[float, tuple[bool, bool, bool]]: - """Azimuth-bin coverage of the viewpoints and which world axes were observed.""" - if not members: - return 0.0, (False, False, False) - centroid = np.median(np.stack([m.centroid for m in members]), axis=0) - directions = [] - for m in members: - v = m.camera_position - centroid - norm = np.linalg.norm(v) - if norm > 1e-6: - directions.append(v / norm) - if not directions: - return 0.0, (False, False, False) - dirs = np.stack(directions) - azimuth = np.arctan2(dirs[:, 1], dirs[:, 0]) - bins = set(((azimuth + np.pi) / (2 * np.pi) * 8).astype(int) % 8) - coverage = len(bins) / 8.0 - observed = tuple(bool((np.abs(dirs[:, i]) > 0.3).any()) for i in range(3)) - return coverage, observed # type: ignore[return-value] - - -def _aggregated_label(labels: tuple[tuple[str, float], ...], policy: InventoryPolicy) -> str | None: - """The track's name: its best canonical group, if the margin holds again. - - The per-frame margin referees only candidates competing inside one view. - Two views can each accept a different group cleanly, and that - disagreement would otherwise reach ``primary_label`` unrefereed. - """ - if not labels: - return None - if len(labels) > 1 and labels[0][1] - labels[1][1] < policy.name_refusal_margin: - return None - return labels[0][0] - - -def _build_instance( - index: int, track: _Track, policy: InventoryPolicy, frame_id: str, grounded: bool = True -) -> Instance: - labels = tuple(sorted(track.labels.items(), key=lambda kv: -kv[1])) - primary = _aggregated_label(labels, policy) - latest = track.latest - coverage, axes_observed = _view_coverage(track.members) - lo, hi = track.aabb - center = (lo + hi) / 2 - extent = np.maximum(hi - lo, 0.005) - centroids = np.stack([m.centroid for m in track.members]) - sigma = centroids.std(axis=0) if len(track.members) > 1 else np.full(3, 0.01) - support = Support( - center_xyz=(float(center[0]), float(center[1]), float(center[2])), - extent_xyz_m=(float(extent[0]), float(extent[1]), float(extent[2])), - orientation_xyzw=(0.0, 0.0, 0.0, 1.0), - sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), - coverage=coverage, - axes_observed=axes_observed, - frame_id=frame_id, - ) - distinct_views = len({tuple(np.round(m.camera_position, 2)) for m in track.members}) - return Instance( - instance_id=f"obj-{index:02d}", - grounded=grounded, - primary_label=primary, - labels=labels, - state="active", - identity_confidence=min(1.0, distinct_views / 3.0), - support=support, - latest_position_xyz=( - float(latest.centroid[0]), - float(latest.centroid[1]), - float(latest.centroid[2]), - ), - latest_seen_ts=latest.ts, - members=track.members, - ) - - -def _naming_picks(track: _Track) -> list[SupportObservation]: - """Members to name on: the largest view, then maximal viewpoint spread. - - Picking by mask area alone selects near-duplicate views when the sweep - keeps returning to one vantage; the label then hinges on a single - viewing angle. Greedy farthest-point selection over camera positions - guarantees the close-up passes participate. - """ - if len(track.members) <= NAME_FRAMES_PER_INSTANCE: - return list(track.members) - picks = [max(track.members, key=lambda m: m.mask_area_px)] - remaining = [m for m in track.members if m is not picks[0]] - while len(picks) < NAME_FRAMES_PER_INSTANCE and remaining: - best = max( - remaining, - key=lambda m: min( - float(np.linalg.norm(m.camera_position - p.camera_position)) for p in picks - ), - ) - picks.append(best) - remaining.remove(best) - return picks - - -def _flatten(vocabulary: NamingVocabulary) -> tuple[list[str], np.ndarray, list[str]]: - """Group table to a query list, the group start offsets, and the canonicals.""" - queries = [surface for group in vocabulary for surface in group] - starts = np.cumsum([0] + [len(group) for group in vocabulary[:-1]]) - return queries, starts, [group[0] for group in vocabulary] - - -def _accepted_groups( - scores: np.ndarray, starts: np.ndarray, policy: InventoryPolicy -) -> tuple[np.ndarray, np.ndarray]: - """Per box, the winning group and its score - or -1 where the margin refuses. - - A group's claim is its best surface string, since strings of one group - have unequal detector affinity. The accept floor is already applied: it - is the threshold the request was made with. - """ - groups = np.maximum.reduceat(scores, starts, axis=1) - best = groups.argmax(axis=1) - top = groups[np.arange(len(groups)), best] - if groups.shape[1] == 1: - return best, top - runner_up = np.partition(groups, -2, axis=1)[:, -2] - return np.where(top - runner_up >= policy.name_refusal_margin, best, -1), top - - -def _name_and_suppress( - tracks: list[_Track], - tracks_2d: list[_Track2D], - color: Any, - detector: Owlv2Detector, - vocabulary: NamingVocabulary, - policy: InventoryPolicy, -) -> None: - """OWLv2 naming per instance on keyframes, person/hand suppressing observations. - - Runs after association by construction: association consumed unnamed - supports, so per-view label instability cannot starve existence or split - an instance. A naming failure degrades names, never counts. - - Naming reads whole score rows, so one request per frame carries every - candidate name. Suppression is an existence decision and keeps its own - request, which runs whether or not there is a vocabulary. - """ - frame_members: dict[float, list[tuple[_Track, SupportObservation]]] = {} - for track in tracks: - for member in _naming_picks(track): - frame_members.setdefault(member.ts, []).append((track, member)) - frame_members_2d: dict[float, list[tuple[_Track2D, Detection2DSeg]]] = {} - for track2d in tracks_2d: - for det in sorted(track2d.members, key=lambda d: -(d.mask > 0).sum())[:2]: - frame_members_2d.setdefault(det.ts, []).append((track2d, det)) - - if not frame_members and not frame_members_2d: - return - - queries, starts, canonical = _flatten(vocabulary) - all_ts = sorted(set(frame_members) | set(frame_members_2d)) - logger.info( - f"naming: OWLv2 over {len(all_ts)} keyframes, " - f"{len(canonical)} groups of {len(queries)} queries" - ) - for ts in all_ts: - try: - image = color.at(ts, 0.05).first().data - except LookupError: - continue - - if queries: - boxes, scores = detector.query_score_rows( - image, queries, threshold=policy.name_accept_score - ) - winners, top = _accepted_groups(scores, starts, policy) - for box, group, score in zip(boxes, winners, top, strict=True): - if group < 0: - continue - bbox = (float(box[0]), float(box[1]), float(box[2]), float(box[3])) - best_target: Any = None - best_iou = policy.name_attach_iou - for track, member in frame_members.get(ts, []): - if member.bbox is None: - continue - iou = _bbox_iou(member.bbox, bbox) - if iou > best_iou: - best_target, best_iou = track, iou - for track2d, det2d in frame_members_2d.get(ts, []): - iou = _bbox_iou(det2d.bbox, bbox) - if iou > best_iou: - best_target, best_iou = track2d, iou - if best_target is not None: - label = canonical[group] - best_target.labels[label] = max( - best_target.labels.get(label, 0.0), float(score) - ) - - for det in detector.query_detections(image, SUPPRESS_QUERIES, threshold=SUPPRESS_SCORE): - for track, member in frame_members.get(ts, []): - if member.bbox is None: - continue - inside = _mask_overlap_fraction_bbox(member.bbox, det.bbox) - if inside >= SUPPRESS_OVERLAP and member in track.members: - track.members.remove(member) - - -def _mask_overlap_fraction_bbox( - member_box: tuple[float, float, float, float], region: tuple[float, float, float, float] -) -> float: - """Fraction of the member box inside the region box.""" - mx1, my1, mx2, my2 = member_box - rx1, ry1, rx2, ry2 = region - ix = max(0.0, min(mx2, rx2) - max(mx1, rx1)) - iy = max(0.0, min(my2, ry2) - max(my1, ry1)) - area = (mx2 - mx1) * (my2 - my1) - return (ix * iy) / area if area > 0 else 0.0 - - -def inventory( - store: Any, - *, - segmenter: EdgeTAMImageSegmenter, - detector: Owlv2Detector, - naming_vocabulary: NamingVocabulary, - after: float | None = None, - before: float | None = None, - include_ungrounded: bool = False, - policy: InventoryPolicy | None = None, - motion_threshold: float = MOTION_THRESHOLD, - log_progress: bool = False, - rig: Rig | None = None, -) -> list[Instance]: - """Deduplicated object instances for the window, computed at query time. - - Reports the scene as of the window's end: an instance's position and - timestamp come from its latest member observation in this call. An - object that moved between rest positions inside the window registers - once per rest position; linking rest positions of one object is - cross-time identity and out of scope here. - - ``naming_vocabulary`` supplies candidate names and nothing else: an - instance whose best canonical group misses the refusal margin keeps its - ``unknown-N`` name. The empty tuple is a supported mode - discovery and - suppression run and every instance is ``unknown-N``. - - ``log_progress`` enables per-keyframe discovery lines - (``discovery: i/n ts_offset=… prop=… scope=… …s``). Off by default. - - Both models belong to the caller: nothing here is loaded or stopped, so - one process can call this repeatedly on warm weights, over as many - windows as it wants. Without a ``rig`` the store's shape decides one, - and without a ``policy`` the rig supplies scale-appropriate defaults. - """ - rig = rig or Rig.from_store(store) - policy = policy or rig.default_inventory_policy() - lo, hi = rig.color.get_time_range() - t0 = after if after is not None else lo - t1 = before if before is not None else hi - logger.info(f"inventory window: {t0 - lo:.1f}s to {t1 - lo:.1f}s ({t1 - t0:.1f}s)") - - keyframes = rig.keyframes(t0, t1, policy.keyframe_stride_s, motion_threshold) - logger.info(f"gates: {len(keyframes)} keyframes pass the capture gates") - if not keyframes: - return [] - - plane = fit_support_plane(rig, keyframes) - if plane is not None: - logger.info(f"support plane: {plane.inlier_count} inliers") - - frames_grounded: list[tuple[float, list[SupportObservation]]] = [] - frames_ungrounded: list[tuple[float, list[Detection2DSeg]]] = [] - image_area = float(rig.camera_info.width * rig.camera_info.height) - - from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D - - n_kf = len(keyframes) - for i, obs in enumerate(keyframes, start=1): - t_frame = perf_counter() if log_progress else 0.0 - proposals = segmenter.propose_all(obs.data) - accepted = [det for det in proposals if _proposal_passes_2d(det, image_area, policy)] - accepted = sorted(accepted, key=lambda d: -(d.mask > 0).sum())[:MAX_PROPOSALS_PER_FRAME] - if not accepted: - if log_progress: - logger.info( - f"discovery: {i}/{n_kf} ts_offset={obs.ts - lo:.1f}s " - f"prop={len(proposals)}->0 scope=0 {perf_counter() - t_frame:.1f}s" - ) - continue - - pose = rig.camera_pose(obs.ts) - if pose is None: - if log_progress: - logger.info( - f"discovery: {i}/{n_kf} ts_offset={obs.ts - lo:.1f}s " - f"skip=no_pose {perf_counter() - t_frame:.1f}s" - ) - continue - camera_position = np.array([pose.position.x, pose.position.y, pose.position.z]) - - for j, det in enumerate(accepted): - det.track_id = j - grounded, ungrounded = _lift_frame( - ImageDetections2D(obs.data, accepted), - rig, - obs.ts, - camera_position, - policy, - plane, - ) - grounded = [o for o in grounded if _in_scope(o, plane, policy)] - grounded = _merge_same_frame(grounded, policy) - frames_grounded.append((obs.ts, grounded)) - frames_ungrounded.append((obs.ts, ungrounded)) - if log_progress: - logger.info( - f"discovery: {i}/{n_kf} ts_offset={obs.ts - lo:.1f}s " - f"prop={len(proposals)}->{len(accepted)} scope={len(grounded)} " - f"{perf_counter() - t_frame:.1f}s" - ) - - _free_accelerator() - - total = sum(len(g) for _, g in frames_grounded) - logger.info(f"discovery: {total} in-scope supports across {len(frames_grounded)} keyframes") - - tracks = _associate(frames_grounded, policy) - tracks = _merge_tracks(tracks, policy) - tracks_2d = _track_ungrounded(frames_ungrounded) if include_ungrounded else [] - logger.info(f"association: {len(tracks)} grounded instances") - - _name_and_suppress(tracks, tracks_2d, rig.color, detector, naming_vocabulary, policy) - tracks = [t for t in tracks if len(t.members) >= policy.min_member_observations] - - tracks.sort(key=lambda t: min(m.ts for m in t.members)) - instances: list[Instance] = [] - unknown = 0 - for index, track in enumerate(tracks): - instance = _build_instance(index, track, policy, rig.world_frame) - if instance.primary_label is None: - instance.primary_label = f"unknown-{unknown}" - unknown += 1 - instances.append(instance) - - if include_ungrounded: - for track2d in tracks_2d: - labels = tuple(sorted(track2d.labels.items(), key=lambda kv: -kv[1])) - primary = _aggregated_label(labels, policy) - if primary is None: - primary = f"unknown-{unknown}" - unknown += 1 - latest = max(track2d.members, key=lambda d: d.ts) - instances.append( - Instance( - instance_id=f"obj-{len(instances):02d}", - grounded=False, - primary_label=primary, - labels=labels, - state="active", - identity_confidence=0.3, - support=None, - latest_position_xyz=None, - latest_seen_ts=latest.ts, - members=[], - ) - ) - return instances - - -def _free_accelerator() -> None: - import gc - - import torch - - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif torch.backends.mps.is_available(): - torch.mps.empty_cache() diff --git a/dimos/perception/memory/tool_inventory.py b/dimos/perception/memory/tool_inventory.py deleted file mode 100644 index 850d3150e1..0000000000 --- a/dimos/perception/memory/tool_inventory.py +++ /dev/null @@ -1,333 +0,0 @@ -# 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. - -"""Enumerate deduplicated object instances in a recording window. - -Run: uv run python -m dimos.perception.memory.tool_inventory [out.rrd] - [--from ] [--duration ] [--labels ...] [--no-vocabulary] - [--include-ungrounded] [--log-progress] - -``--labels`` replaces the candidate names with the caller's own. Each token -is one group: the first phrase is the canonical label, and ``|`` separates -synonyms (``"pen|ballpoint pen|ink pen"``). With none given the run uses -``DEFAULT_VOCABULARY``, and ``--no-vocabulary`` names nothing. Naming -abstains either way: a name is reported only when it beats its runner-up by -``InventoryPolicy.name_refusal_margin``, else the instance is ``unknown-N``. - -Stdout contract: a summary line ``instances: N`` followed by one line per -instance: `` id= name= xyz=(x,y,z) ts_offset= -members= extent=(x,y,z) sigma=(x,y,z) coverage=``. Exit code 0 -whenever the call completes; an empty scene is ``instances: 0``, not a -failure. - -``extent`` is the instance's bounding size in meters and ``sigma`` the -spread of its member centroids, so a caller can check an object against a -gripper without a second query. Both are what the views actually saw: -``coverage`` is the fraction of viewing azimuth covered, and an instance -seen from one side reports the extent of that side. - -The instance list reports the scene as of the window's end: position and -timestamp come from each instance's latest member observation. An object -moved between rest positions inside the window registers once per rest -position - linking rest positions of one object across time is -re-identification, which this tool does not do. - -The .rrd holds the same instances the stdout lines report, no second -perception pass: the member clouds on the timeline, one labeled box per -instance in 3D, and each member's pixel box on the camera view at the -keyframe it came from. -""" - -import argparse -import json -from pathlib import Path -import sys -from typing import cast - -from dimos.memory.store.sqlite import SqliteStore -from dimos.memory.transform import throttle -from dimos.perception.memory.inventory import DEFAULT_VOCABULARY, NamingVocabulary -from dimos.perception.memory.rig import Rig -from dimos.perception.memory.types import Instance, SupportObservation -from dimos.utils.data import get_data - - -def labels_to_vocabulary(tokens: list[str]) -> NamingVocabulary: - """Parse ``--labels`` tokens into synonym groups. - - Each token is one group. Phrases inside a token are separated by ``|``. - The first non-empty phrase is the canonical label. - """ - groups: list[tuple[str, ...]] = [] - for token in tokens: - surfaces = tuple(part.strip() for part in token.split("|") if part.strip()) - if not surfaces: - raise ValueError(f"empty --labels group: {token!r}") - groups.append(surfaces) - return tuple(groups) - - -def instance_label(instance: Instance) -> str: - """``obj-NN name score`` - the score only when that name won the instance.""" - if instance.labels and instance.labels[0][0] == instance.primary_label: - return f"{instance.instance_id} {instance.primary_label} {instance.labels[0][1]:.2f}" - return f"{instance.instance_id} {instance.primary_label}" - - -def render(out: str, rig: Rig, instances: list[Instance], t0: float, t1: float) -> None: - """Write the .rrd - rerun stays an inline import. - - Entity contract: ``map`` backdrop, ``camera/image`` the live feed carrying - the per-keyframe pixel boxes, ``instances/_`` the member clouds - with a static labeled box. One color per instance across both views. - """ - import rerun as rr - import rerun.blueprint as rrb - - from dimos.memory.vis.color import Color - from dimos.visualization.rerun.init import rerun_init - - rerun_init("memory-inventory") - rr.save(out) - rr.send_blueprint( - rrb.Blueprint( - rrb.Horizontal( - rrb.Spatial3DView(origin="/", name="Scene"), - rrb.Spatial2DView(origin="camera", name="Live"), - column_shares=[2, 1], - ) - ) - ) - - point_size = 0.005 if rig.depth is not None else 0.015 - - def at(ts: float) -> None: - rr.set_time("ts", timestamp=ts) - - grounded = [instance for instance in instances if instance.support is not None] - colors = [ - list(Color.from_cmap("turbo", i / max(len(grounded) - 1, 1)).rgb_u8()) - for i in range(len(grounded)) - ] - labels = [instance_label(instance) for instance in grounded] - paths = [ - f"instances/{instance.instance_id}_{cast('str', instance.primary_label).replace(' ', '_')}" - for instance in grounded - ] - - frames: dict[float, list[tuple[int, SupportObservation]]] = {} - for i, instance in enumerate(grounded): - for member in instance.members: - frames.setdefault(member.ts, []).append((i, member)) - - # scene backdrop: for depth rigs the last instance keyframe's RGBD cloud, - # for pointcloud rigs the window's scans merged into one map - if rig.depth is not None: - backdrop_ts = max(frames, default=None) - if backdrop_ts is not None: - backdrop = rig.backdrop(backdrop_ts) - if backdrop is not None: - rr.log( - "map", - backdrop.voxel_downsample(0.01).to_rerun(voxel_size=point_size), - static=True, - ) - else: - import numpy as np - - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - - scans = [ - points - for obs in rig.cloud.after(t0).before(t1).transform(throttle(2.0)) - if (points := rig.registered_scan(obs)) is not None - ] - if scans: - merged = PointCloud2.from_numpy(np.vstack(scans), frame_id=rig.world_frame) - rr.log("map", merged.voxel_downsample(0.05).to_rerun(voxel_size=0.01), static=True) - - # live camera feed + frustum; the empty box clears the overlay off non-keyframes - rr.log("camera", rig.camera_info.to_rerun(), static=True) - feed_throttle = 0.1 if (t1 - t0) <= 160 else 0.4 - feed = rig.color.after(t0).before(t1).transform(throttle(feed_throttle)) - for obs in feed: - pose = rig.camera_pose(obs.ts) - if pose is None: - continue - at(obs.ts) - rr.log("camera/image", obs.data.to_rerun()) - rr.log("camera", pose.to_rerun()) - rr.log("camera/image/instances", rr.Boxes2D(array=[], array_format=rr.Box2DFormat.XYXY)) - - # keyframes, logged after the feed so their boxes win the shared timestamps - for ts, entries in sorted(frames.items()): - at(ts) - keyframe_pose = rig.camera_pose(ts) - assert keyframe_pose is not None - rr.log("camera/image", rig.color.at(ts, 0.05).first().data.to_rerun()) - rr.log("camera", keyframe_pose.to_rerun()) - rr.log( - "camera/image/instances", - rr.Boxes2D( - array=[ - cast("tuple[float, float, float, float]", member.bbox) for _, member in entries - ], - array_format=rr.Box2DFormat.XYXY, - labels=[labels[i] for i, _ in entries], - colors=[colors[i] for i, _ in entries], - ), - ) - for i, member in entries: - rr.log(paths[i], member.cloud.to_rerun(voxel_size=point_size, colors=colors[i])) - - # the reported instance: one labeled box, static so it holds over the whole timeline - for i, instance in enumerate(grounded): - support = instance.support - assert support is not None - cx, cy, cz = support.center_xyz - ex, ey, ez = support.extent_xyz_m - rr.log( - f"{paths[i]}/box", - rr.Boxes3D( - centers=[(cx, cy, cz)], - half_sizes=[(ex / 2, ey / 2, ez / 2)], - colors=[colors[i]], - labels=[labels[i]], - fill_mode=rr.components.FillMode.MajorWireframe, - ), - static=True, - ) - - -def main() -> int: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument( - "out", nargs="?", default=None, help="rerun recording to write; omitted writes none" - ) - parser.add_argument("--dataset", type=Path, help="memory recording database") - parser.add_argument("--manifest", type=Path, help="rig manifest json (default: .rig.json)") - parser.add_argument("--color", help="stream name override for the color role") - parser.add_argument("--depth", help="stream name override for the depth role") - parser.add_argument("--cloud", help="stream name override for the pointcloud role") - parser.add_argument("--odom", help="stream name override for the poses role") - parser.add_argument( - "--from", dest="start", type=float, default=0.0, help="start offset into the recording (s)" - ) - parser.add_argument("--duration", type=float, default=None, help="how much to parse (s)") - vocab = parser.add_mutually_exclusive_group() - vocab.add_argument( - "--labels", - nargs="+", - metavar="GROUP", - help=( - "candidate name groups; each token is one group, " - "'|' separates synonyms (first phrase is canonical); " - "default is DEFAULT_VOCABULARY" - ), - ) - vocab.add_argument( - "--no-vocabulary", - action="store_true", - help="name nothing: every instance stays unknown-N", - ) - parser.add_argument( - "--include-ungrounded", - action="store_true", - help="also list RGB-only instances that never produced valid depth", - ) - parser.add_argument( - "--log-progress", - action="store_true", - help="per-keyframe discovery progress lines (off by default)", - ) - args = parser.parse_args() - - out = args.out - if args.labels and args.labels[-1].endswith(".rrd"): - out = args.labels.pop() - - dataset = args.dataset or get_data( - "xarm6_worldbelief_realsense_d435i_stationery_calibrated/" - "xarm6_worldbelief_20260729_203624_161992.db" - ) - manifest = json.loads(args.manifest.read_text()) if args.manifest else None - overrides = { - role: name - for role, name in [ - ("color", args.color), - ("depth", args.depth), - ("cloud", args.cloud), - ("poses", args.odom), - ] - if name - } - store = SqliteStore(path=dataset) - rig = Rig.from_store(store, manifest=manifest, overrides=overrides) - lo, hi = rig.color.get_time_range() - after = lo + args.start - before = lo + args.start + args.duration if args.duration is not None else None - - if args.no_vocabulary: - naming_vocabulary: NamingVocabulary = () - elif args.labels: - naming_vocabulary = labels_to_vocabulary(args.labels) - else: - naming_vocabulary = DEFAULT_VOCABULARY - - from dimos.perception.memory.dandetect import DanDetector - - with DanDetector() as dan: - instances = dan.inventory( - store, - naming_vocabulary=naming_vocabulary, - after=after, - before=before, - include_ungrounded=args.include_ungrounded, - log_progress=args.log_progress, - rig=rig, - ) - - print(f"instances: {len(instances)}") - for i, instance in enumerate(instances): - if instance.latest_position_xyz is not None: - x, y, z = instance.latest_position_xyz - xyz = f"({x:.3f},{y:.3f},{z:.3f})" - else: - xyz = "None" - if instance.support is not None: - ex, ey, ez = instance.support.extent_xyz_m - sx, sy, sz = instance.support.sigma_xyz_m - geometry = ( - f" extent=({ex:.3f},{ey:.3f},{ez:.3f})" - f" sigma=({sx:.3f},{sy:.3f},{sz:.3f})" - f" coverage={instance.support.coverage:.2f}" - ) - else: - geometry = "" - print( - f"{i} id={instance.instance_id} name={instance.primary_label} " - f"xyz={xyz} ts_offset={instance.latest_seen_ts - lo:.1f} " - f"members={len(instance.members)}{geometry}" - ) - - if out is not None: - render(out, rig, instances, after, before if before is not None else hi) - print(f"saved {out}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) From de79d19933149c3e811a18d747803f689b449511 Mon Sep 17 00:00:00 2001 From: bogwi Date: Fri, 4 Sep 2026 01:52:18 +0900 Subject: [PATCH 27/28] refactor the stack --- dimos/perception/detection/identity.py | 146 --- .../memory/blueprints/go2_localize_live.py | 17 +- dimos/perception/memory/dandetect.py | 77 +- dimos/perception/memory/gates.py | 134 --- dimos/perception/memory/identity_store.py | 58 -- dimos/perception/memory/localize.py | 612 ++++++------ dimos/perception/memory/rig.py | 874 +++++++----------- dimos/perception/memory/support_plane.py | 62 +- dimos/perception/memory/tool_localize.py | 8 +- dimos/perception/memory/types.py | 230 +---- 10 files changed, 704 insertions(+), 1514 deletions(-) delete mode 100644 dimos/perception/detection/identity.py delete mode 100644 dimos/perception/memory/gates.py delete mode 100644 dimos/perception/memory/identity_store.py diff --git a/dimos/perception/detection/identity.py b/dimos/perception/detection/identity.py deleted file mode 100644 index 5eb6c9e89d..0000000000 --- a/dimos/perception/detection/identity.py +++ /dev/null @@ -1,146 +0,0 @@ -# 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. - -"""Object identity over detection streams: many sightings in, one per object out. - -``Identity`` is the aggregation stage of a search pipeline:: - - detections_2d.transform(ProjectTo3D(cloud, ...)).transform(Identity()) - -As a stream transformer it is a batch search processor: it consumes the full -upstream of 3D detections, groups the sightings of one physical object, and -then emits one merged :class:`Detection3DPC` per object - the union cloud of -every viewpoint that saw it. What counts as "the same object" is a pluggable -``is_same(a, b)`` strategy; v0 is spatial only ("is it roughly at the same -spot"), so there is no permanence: an object that moved registers as a new -object at its new rest position. - -The grouping core (``add`` plus ``groups``) is usable directly for callers -that need the members of each identity rather than the merged stream output; -``localize`` forms its support candidates with it. -""" - -from __future__ import annotations - -import operator -from typing import TYPE_CHECKING, Any - -import numpy as np - -from dimos.memory.transform import Transformer -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC -from dimos.perception.detection.type.imageDetections import ImageDetections - -if TYPE_CHECKING: - from collections.abc import Callable, Iterator - - from dimos.memory.type.observation import Observation - - -def spatial(radius: float = 0.1) -> Callable[[Detection3DPC, Detection3DPC], bool]: - """Same object when the cloud centers sit within *radius* meters.""" - - def is_same(a: Detection3DPC, b: Detection3DPC) -> bool: - return float((a.center - b.center).magnitude()) <= radius - - return is_same - - -def fused(voxel: float) -> Callable[[Detection3DPC, Detection3DPC], Detection3DPC]: - """Merges two detections by combining their point clouds. For each grid cell of size *voxel* meters, - it keeps the mean position of all points grouped in that cell. - The weight of each fused point is the number of raw points that contributed to it. - """ - weights: dict[int, np.ndarray[Any, Any]] = {} - - def merge(a: Detection3DPC, b: Detection3DPC) -> Detection3DPC: - union = a + b - wa = weights.pop(id(a), None) - na = float(wa.sum()) if wa is not None else float(len(a.pointcloud)) - nb = float(len(b.pointcloud)) - union.center = Vector3( - (na * a.center.x + nb * b.center.x) / (na + nb), - (na * a.center.y + nb * b.center.y) / (na + nb), - (na * a.center.z + nb * b.center.z) / (na + nb), - ) - if voxel <= 0: - return union - pts = np.asarray(union.pointcloud.pointcloud.points) - w = np.ones(len(pts)) - if wa is not None: - w[: len(wa)] = wa - cells, inverse = np.unique( - np.floor(pts / voxel).astype(np.int64), axis=0, return_inverse=True - ) - wsum = np.zeros(len(cells)) - np.add.at(wsum, inverse, w) - psum = np.zeros((len(cells), 3)) - np.add.at(psum, inverse, pts * w[:, None]) - union.pointcloud = PointCloud2.from_numpy( - psum / wsum[:, None], - frame_id=union.pointcloud.frame_id, - timestamp=union.pointcloud.ts, - ) - weights[id(union)] = wsum - return union - - return merge - - -class Identity(Transformer[Any, Detection3DPC]): - """One detection3D per object, aggregated from every sighting. - - Each incoming detection is matched against the running merged - representative of every known object with ``is_same``; a match joins - that object and folds into its representative with ``merge``, otherwise - it founds a new object. Upstream observations may carry a single - :class:`Detection3DPC` or a per-frame :class:`ImageDetections` batch. - """ - - def __init__( - self, - is_same: Callable[[Detection3DPC, Detection3DPC], bool] | None = None, - merge: Callable[[Detection3DPC, Detection3DPC], Detection3DPC] | None = None, - ) -> None: - self.is_same = is_same or spatial() - self.merge = merge or operator.add - self.groups: list[list[Detection3DPC]] = [] - self.merged: list[Detection3DPC] = [] - - def add(self, detection: Detection3DPC) -> int: - """Assign one sighting to its object; returns the object's index.""" - for index, representative in enumerate(self.merged): - if self.is_same(representative, detection): - self.groups[index].append(detection) - self.merged[index] = self.merge(representative, detection) - return index - self.groups.append([detection]) - self.merged.append(detection) - return len(self.groups) - 1 - - def __call__( - self, upstream: Iterator[Observation[Any]] - ) -> Iterator[Observation[Detection3DPC]]: - template: Observation[Any] | None = None - for obs in upstream: - template = obs - data = obs.data - for detection in data if isinstance(data, ImageDetections) else [data]: - self.add(detection) - if template is None: - return - for merged in self.merged: - yield template.derive(data=merged, ts=merged.ts, pose=merged.pose) diff --git a/dimos/perception/memory/blueprints/go2_localize_live.py b/dimos/perception/memory/blueprints/go2_localize_live.py index 4c67216d7f..3c99178e8e 100644 --- a/dimos/perception/memory/blueprints/go2_localize_live.py +++ b/dimos/perception/memory/blueprints/go2_localize_live.py @@ -15,7 +15,7 @@ """A go2 recording looped as a live feed, with localize behind an agent skill. This is the deployment shape of the perception memory stack: one module owns -a store and the models, ``DanDetector.embed(live=True)`` keeps a background +a store and the models, ``DanDetector.embed_live`` keeps a background tail filling the index while the robot runs, and a ``@skill`` answers ``localize`` calls against whatever has been embedded so far. @@ -73,7 +73,7 @@ lattice_quantum, ) from dimos.perception.memory.dandetect import DanDetector -from dimos.perception.memory.identity_store import IdentityStore +from dimos.perception.memory.localize import Groups from dimos.perception.memory.rig import Rig from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL, GO2Connection from dimos.utils.logging_config import setup_logger @@ -170,7 +170,7 @@ def _feed(self) -> None: # another worker backfills once and then never sees another append. self._embedder = self.register_disposable(DanDetector()) self._embedder.start() - self._embedder.embed(live, live=True, rig=_live_rig(Rig.from_store(source), live)) + self._embedder.embed_live(live, rig=_live_rig(Rig.from_store(source), live)) logger.info( f"loop feeder: {self.config.dataset} ({span - LAP_GAP_S:.1f}s) -> {self.config.db_path}" ) @@ -257,19 +257,16 @@ def _live_rig(source: Rig, live: Any) -> Rig: and carries the measured color delay over instead of re-estimating it. """ return Rig( - camera_info=source.camera_info, + cameras=source.cameras, color=live.stream("color_image", source.color.data_type), world_frame=source.world_frame, - optical_frame=source.optical_frame, tf=StreamTF(live.stream("tf", TFMessage)) if source.tf is not None else None, - base_to_optical=source.base_to_optical, + mounts=source.mounts, poses=live.stream("odom", source.poses.data_type) if source.poses is not None else None, cloud=live.stream("lidar", source.cloud.data_type) if source.cloud is not None else None, tf_tolerance=source.tf_tolerance, cloud_accum_s=source.cloud_accum_s, - speed_max=source.speed_max, color_delay=source.color_delay, - scene_gate=source.scene_gate, embed_hz=source.embed_hz, mobile=source.mobile, ) @@ -296,7 +293,7 @@ def start(self) -> None: super().start() self._ready = threading.Event() self._stage = "starting" - self._identity = IdentityStore() + self._groups: dict[str, Groups] = {} self._thread = threading.Thread(target=self._warm, name="localize-warmup", daemon=True) self._thread.start() @@ -379,7 +376,7 @@ def localize( index=index, rig=self.rig, policy=tuning, - identity_store=self._identity, + groups=self._groups, ) self.detections.publish(_as_detection_array(queries, results, self.rig.world_frame)) diff --git a/dimos/perception/memory/dandetect.py b/dimos/perception/memory/dandetect.py index c524761123..5f4bd008c7 100644 --- a/dimos/perception/memory/dandetect.py +++ b/dimos/perception/memory/dandetect.py @@ -12,20 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""One disposable resource wrapping the memory perception API. +"""The perception models as one disposable resource. -``DanDetector`` owns the models behind :func:`embed_index` and -:func:`localize`: enter once, query many times on warm weights, and -``stop()`` (or leave the ``with`` block) releases whatever loaded. - -Every entry point takes an optional :class:`~dimos.perception.memory.rig.Rig` -describing where poses and 3D geometry come from; without one the store's -shape decides. +Enter once, query many times on warm weights; ``stop()`` releases whatever +loaded. Every entry point takes an optional +:class:`~dimos.perception.memory.rig.Rig`; without one the store's shape +decides. """ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Literal, cast, overload +from typing import TYPE_CHECKING, Any from dimos.core.resource import Resource from dimos.memory.embed import EmbedImages @@ -44,12 +41,7 @@ class DanDetector(Resource): - """The perception models as one resource. - - ``start()`` constructs SigLIP, OWLv2, and EdgeTAM. The two - HuggingFace models load lazily on first use. ``stop()`` releases - whatever loaded. - """ + """SigLIP, OWLv2 and EdgeTAM; the HuggingFace two load on first use.""" siglip: SigLIPModel detector: Owlv2Detector @@ -74,52 +66,21 @@ def stop(self) -> None: self.detector.stop() del self.segmenter - @overload - def embed( - self, - store: Any, - after: float, - before: float, - *, - live: Literal[False] = False, - rig: Rig | None = ..., - ) -> Stream[Any, Any]: ... - @overload - def embed( - self, - store: Any, - *, - live: Literal[True], - rig: Rig | None = ..., - ) -> Stream[Any, Any]: ... def embed( - self, - store: Any, - after: float | None = None, - before: float | None = None, - *, - live: bool = False, - rig: Rig | None = None, + self, store: Any, after: float, before: float, *, rig: Rig | None = None ) -> Stream[Any, Any]: - """SigLIP-embedded, world-posed frame index for :meth:`localize`. + """SigLIP-embedded, world-posed index over ``[after, before]``.""" + return embed_index(store, self.siglip, after, before, rig=rig or Rig.from_store(store)) - Replay mode indexes ``[after, before]`` in memory and returns when - done. ``live=True`` instead tails the rig's color stream and keeps - saving into the store's named ``color_image_embedded`` stream on a - background thread; the returned stream is that named stream. - """ - rig = rig or Rig.from_store(store) - if not live: - return embed_index( - store, - self.siglip, - cast("float", after), - cast("float", before), - rig=rig, - ) + def embed_live(self, store: Any, *, rig: Rig | None = None) -> Stream[Any, Any]: + """Tail the colour stream into ``color_image_embedded`` on a background thread. + Returns that named stream, which :meth:`localize` reads like a replay + index; it keeps filling for as long as the resource is open. + """ from dimos.msgs.sensor_msgs.Image import Image + rig = rig or Rig.from_store(store) embedded: Stream[Any, Any] = store.stream("color_image_embedded", Image) pipeline = ( rig.color.live() @@ -142,11 +103,7 @@ def localize( policy: LocalizePolicy | None = None, **kwargs: Any, ) -> list[Localization] | list[list[Localization]]: - """:func:`localize` on this resource's models. - - ``policy`` is the localize thresholds. ``None`` uses the rig's scale - defaults. - """ + """:func:`localize` on this resource's models.""" return localize( store, query, diff --git a/dimos/perception/memory/gates.py b/dimos/perception/memory/gates.py deleted file mode 100644 index 145f25efdd..0000000000 --- a/dimos/perception/memory/gates.py +++ /dev/null @@ -1,134 +0,0 @@ -# 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. - -"""The scene-motion gate: image differencing conditioned on camera stillness. - -Pose-derived gates (camera pose, speed, stillness, keyframes) live on -:class:`~dimos.perception.memory.rig.Rig` - they depend on where the rig's -poses come from. What stays here is purely image-based: the scene-motion -gate differences frames of the color stream, conditioned on camera -stillness at both compared instants. It rejects frames captured while the -scene itself changes - a parked camera watching hands rearrange objects is -exactly the case poses cannot see. The reference frame is anchored at the -start of the surrounding camera-still interval, so a frame is trusted only -while the scene still matches the state it had when the camera parked. - -The stillness intervals and the grayscale memo are allocated by the caller -and passed in, one per query. -""" - -from __future__ import annotations - -from bisect import bisect_right -from typing import Any - -import numpy as np - -OPTICAL_FRAME = "camera_color_optical_frame" -WORLD_FRAME = "world" - -# One world-pose period plus margin. -TF_TOLERANCE = 0.12 - -SPEED_MAX = 0.02 # m/s - camera counts as still below this -STILL_ENVELOPE = 0.15 # s - stillness must hold over the whole capture envelope - -# Scene-motion gate: downscaled grayscale absolute difference. -DIFF_PIXEL_THRESHOLD = 28 # gray levels - per-pixel change floor -MOTION_THRESHOLD = 0.02 # fraction of changed pixels that flags scene motion -SHORT_DIFF_DT = 0.45 # s - bilateral diff span for active motion -DIFF_WIDTH = 212 # px - diff resolution (1/4 of 848) - - -def _interval_containing( - ts: float, intervals: list[tuple[float, float]] -) -> tuple[float, float] | None: - idx = bisect_right([a for a, _ in intervals], ts) - 1 - if idx < 0: - return None - a, b = intervals[idx] - return (a, b) if a - 0.25 <= ts <= b + 0.25 else None - - -def _gray_small(color: Any, ts: float, gray: dict[float, np.ndarray | None]) -> np.ndarray | None: - """Downscaled grayscale of the color frame nearest ts, memoized in *gray*.""" - key = round(ts, 2) - if key in gray: - return gray[key] - - import cv2 - - small_gray: np.ndarray | None = None - try: - frame = color.at(ts, 0.1).first().data - except LookupError: - frame = None - if frame is not None: - img = frame.to_opencv() - h = int(img.shape[0] * DIFF_WIDTH / img.shape[1]) - small = cv2.resize(img, (DIFF_WIDTH, h), interpolation=cv2.INTER_AREA) - small_gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY) - - gray[key] = small_gray - return small_gray - - -def _diff_fraction( - color: Any, ts_a: float, ts_b: float, gray: dict[float, np.ndarray | None] -) -> float | None: - """Fraction of pixels changed between the frames nearest the two instants.""" - a, b = _gray_small(color, ts_a, gray), _gray_small(color, ts_b, gray) - if a is None or b is None or a.shape != b.shape: - return None - delta = np.abs(a.astype(np.int16) - b.astype(np.int16)) - return float((delta > DIFF_PIXEL_THRESHOLD).mean()) - - -def scene_still( - color: Any, - ts: float, - intervals: list[tuple[float, float]], - gray: dict[float, np.ndarray | None], - motion_threshold: float = MOTION_THRESHOLD, -) -> bool: - """True when the scene around ts is static and unchanged since the camera parked. - - Requires camera stillness at every compared instant - unconditioned, - a moving camera changes every pixel and scans become indistinguishable - from manipulation. Three image terms, all below ``motion_threshold``: - - * anchored: against the start of the surrounding camera-still - interval, so anything the scene changed since the camera parked - (an object placed, moved, or removed) rejects every later frame of - that interval; - * bilateral: against frames a fixed short span before and after ts, - which catches hands actively moving through the view. - """ - interval = _interval_containing(ts, intervals) - if interval is None: - return False - a, b = interval - - anchor = min(a + 0.3, ts) - fraction = _diff_fraction(color, anchor, ts, gray) - if fraction is None or fraction > motion_threshold: - return False - - for other in (max(a, ts - SHORT_DIFF_DT), min(b, ts + SHORT_DIFF_DT)): - if abs(other - ts) < 0.05: - continue - fraction = _diff_fraction(color, other, ts, gray) - if fraction is None or fraction > motion_threshold: - return False - return True diff --git a/dimos/perception/memory/identity_store.py b/dimos/perception/memory/identity_store.py deleted file mode 100644 index 4f8c2ad0c4..0000000000 --- a/dimos/perception/memory/identity_store.py +++ /dev/null @@ -1,58 +0,0 @@ -# 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. - -"""Persistent identity groups across :func:`localize` calls. - -Every persistence read and write in ``localize`` goes through this store and -nothing else. A store owned by a long-lived caller makes verification -evidence cumulative over everything seen since the store was created; the -query window then bounds only new detection work, and frames a label has -already ingested are never re-segmented or re-lifted. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import TYPE_CHECKING - -from dimos.perception.detection.identity import Identity - -if TYPE_CHECKING: - from collections.abc import Callable - - from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC - - -@dataclass -class LabelIdentity: - identity: Identity # groups + merged, cumulative - ingested: set[float] = field(default_factory=set) # frame ts already segmented+lifted+added - ungrounded: tuple[float, float] | None = None # best (score, ts) with no depth - - -@dataclass -class IdentityStore: - labels: dict[str, LabelIdentity] = field(default_factory=dict) - - def get_or_create( - self, - label: str, - is_same: Callable[[Detection3DPC, Detection3DPC], bool], - merge: Callable[[Detection3DPC, Detection3DPC], Detection3DPC], - ) -> LabelIdentity: - entry = self.labels.get(label) - if entry is None: - entry = LabelIdentity(identity=Identity(is_same=is_same, merge=merge)) - self.labels[label] = entry - return entry diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 99e4df5524..49bf544a26 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -14,20 +14,11 @@ """Query-time object localization: text prompt to latest 3D pose and cloud. -Search memory with embeddings (SigLIP, -frame-level), open-vocabulary detection (OWLv2, calibrated per-box scores), -segmentation (EdgeTAM), projection to 3D through the rig's geometry - an -aligned depth stream or a registered pointcloud stream. Two -algorithm rules distinguish it from a best-crop search: - -* **Latest-pose semantics.** Every verified instance is returned, - latest-seen first, and each instance's position follows its latest - sighting: "where is it now", never "where did it match best". The - instance's cloud is the union of every viewpoint that saw it. -* **Calibrated refusal.** Every stage carries a score and the answer can be - empty: no accept-level detection, no multi-view confirmation. Coexisting - same-label instances below the refusal margin are flagged, never merged - or silently dropped. +SigLIP frame retrieval, OWLv2 detection, EdgeTAM segmentation, then a lift to +3D through the rig's geometry. Two rules distinguish it from a best-crop +search: every verified instance is returned latest-seen first and positioned +by its latest sighting ("where is it now"), and every stage carries a score so +the answer can be empty rather than a guess. """ from __future__ import annotations @@ -40,10 +31,10 @@ from dimos.memory.embed import EmbedImages from dimos.memory.transform import QualityWindow, peaks -from dimos.perception.detection.identity import Identity, fused, spatial +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D -from dimos.perception.memory.rig import CLOUD_MIN_POINTS, Rig +from dimos.perception.memory.rig import Rig from dimos.perception.memory.types import Localization, LocalizePolicy, Support from dimos.utils.logging_config import setup_logger @@ -56,30 +47,128 @@ from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter from dimos.perception.detection.detectors.owlv2 import Owlv2Detector from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC - from dimos.perception.memory.identity_store import IdentityStore logger = setup_logger() _SCORE_CACHE_MAX = 8192 +# One row per accepted sighting. Everything verification reports is a reduction +# over these columns, so no detection - and no image - is retained. +M_TS, M_SCORE, M_CX, M_CY, M_CZ, M_KX, M_KY, M_KZ = range(8) +M_WIDTH = 8 -# A support candidate is an identity group: the member sightings of one -# object. Everything a group reports is a plain function over its members. + +@dataclass +class IndexFrame: + """What frame selection needs. The pixels stay in the color stream.""" + + ts: float + frame_id: str + sharpness: float -def _similarity(obs: Any) -> float: - return float(obs.similarity) +@dataclass +class LocalizeTrace: + """Artifacts a renderer needs; collected only when one is passed in.""" + + detection_frames: list[Any] = field(default_factory=list) + answers: list[list[Detection3DPC]] = field(default_factory=list) + backdrop_ts: float | None = None + first_match_ts: float | None = None + + +@dataclass +class Groups: + """Sightings of one label grouped into objects, cumulative across calls. + + Parallel per-group arrays. The only geometry held is each group's fused + union and its latest member's own cloud, which the answer's orientation + comes from; per-sighting evidence is eight floats in a flat row. + """ + + centers: list[np.ndarray] = field(default_factory=list) # running weighted center + clouds: list[PointCloud2] = field(default_factory=list) # fused union + weights: list[np.ndarray | None] = field(default_factory=list) # per fused point + raw: list[float] = field(default_factory=list) # raw points folded in so far + latest: list[PointCloud2] = field(default_factory=list) + latest_ts: list[float] = field(default_factory=list) + members: list[list[float]] = field(default_factory=list) # flat, M_WIDTH per row + trace: list[list[Detection3DPC]] = field(default_factory=list) + ingested: set[float] = field(default_factory=set) # frame ts already lifted + ungrounded: tuple[float, float] | None = None # best (score, ts) with no depth + + def add(self, det: Detection3DPC, radius: float, voxel: float) -> int: + """Assign one sighting to its object, folding its cloud in.""" + points = np.asarray(det.pointcloud.pointcloud.points) + centroid = points.mean(axis=0) + camera = (-det.transform).translation + + hit = next( + (i for i, c in enumerate(self.centers) if np.linalg.norm(c - centroid) <= radius), -1 + ) + if hit < 0: + hit = len(self.centers) + self.centers.append(centroid) + self.clouds.append(det.pointcloud) + self.weights.append(None) + self.raw.append(float(len(points))) + self.latest.append(det.pointcloud) + self.latest_ts.append(det.ts) + self.members.append([]) + self.trace.append([]) + else: + self._fuse(hit, det.pointcloud.ts, points, centroid, voxel) + if det.ts > self.latest_ts[hit]: + self.latest[hit] = det.pointcloud + self.latest_ts[hit] = det.ts + + self.members[hit].extend( + (det.ts, float(det.confidence), *centroid, camera.x, camera.y, camera.z) + ) + return hit + + def _fuse(self, i: int, ts: float, points: np.ndarray, centroid: np.ndarray, voxel: float): + """Voxel-average a sighting into the group's union. + + Each fused point carries the raw points behind it, so the union is a + weighted mean over every viewpoint rather than a mean of means, and the + running center follows the same counts. + """ + held = np.asarray(self.clouds[i].pointcloud.points) + na, nb = self.raw[i], float(len(points)) + self.centers[i] = (na * self.centers[i] + nb * centroid) / (na + nb) + self.raw[i] = na + nb + merged = np.vstack([held, points]) + ts = max(self.clouds[i].ts, ts) + frame = self.clouds[i].frame_id + + if voxel > 0: + w = np.ones(len(merged)) + weights = self.weights[i] + if weights is not None: + w[: len(weights)] = weights + _, inverse = np.unique( + np.floor(merged / voxel).astype(np.int64), axis=0, return_inverse=True + ) + wsum = np.zeros(inverse.max() + 1) + np.add.at(wsum, inverse, w) + psum = np.zeros((len(wsum), 3)) + np.add.at(psum, inverse, merged * w[:, None]) + merged = psum / wsum[:, None] + self.weights[i] = wsum + self.clouds[i] = PointCloud2.from_numpy(merged, frame_id=frame, timestamp=ts) + + def rows(self, i: int) -> np.ndarray: + return np.asarray(self.members[i]).reshape(-1, M_WIDTH) def _settled(index: Stream[Any, Any], spacing: float) -> set[int]: - """Ids of the index frames left once sub-spacing duplicates are dropped. + """Ids left once sub-spacing duplicates are dropped. The index emits the sharpest frame per window from a fixed phase, so a camera that moves between windows leaves both the settled frame and the blurred one taken while it was still moving. Their poses differ, so the - blurred one passes as a second viewpoint and lifts to a displaced cloud. - A pair closer than half the index window is one window's content split - by that phase; only the sharper of the two survives here. + blurred one would pass as a second viewpoint and lift to a displaced cloud. """ ids: set[int] = set() last_id = -1 @@ -96,116 +185,32 @@ def _settled(index: Stream[Any, Any], spacing: float) -> set[int]: return ids -def _centroid(det: Detection3DPC) -> np.ndarray: - centroid: np.ndarray = np.asarray(det.pointcloud.pointcloud.points).mean(axis=0) - return centroid - - -def _camera_position(det: Detection3DPC) -> np.ndarray: - position = (-det.transform).translation - return np.array([position.x, position.y, position.z]) - - -def _group_center(members: list[Detection3DPC]) -> np.ndarray: - center: np.ndarray = np.mean(np.stack([_centroid(d) for d in members]), axis=0) - return center - - -def _max_score(members: list[Detection3DPC]) -> float: - return max(d.confidence for d in members) - - -def _latest(members: list[Detection3DPC]) -> Detection3DPC: - return max(members, key=lambda d: d.ts) - - -def _interval(members: list[Detection3DPC]) -> tuple[float, float]: - times = [d.ts for d in members] - return min(times), max(times) - - -def _n_views(members: list[Detection3DPC]) -> int: - return len({tuple(np.round(_camera_position(d), 2)) for d in members}) - - -@dataclass -class LocalizeTrace: - """Intermediate artifacts collected for rendering; filled when passed in.""" - - detection_frames: list[Any] = field(default_factory=list) # Observation[ImageDetections2D] - matched: list[tuple[float, Detection3DPC]] = field(default_factory=list) - verified: list[tuple[float, Detection3DPC]] = field(default_factory=list) - answers: list[list[Detection3DPC]] = field(default_factory=list) # sightings per instance - backdrop_ts: float | None = None - - -def _quaternion_from_matrix(rotation: np.ndarray) -> tuple[float, float, float, float]: - from scipy.spatial.transform import Rotation - - x, y, z, w = Rotation.from_matrix(rotation).as_quat() - return (float(x), float(y), float(z), float(w)) - - -def _azimuth_coverage(members: list[Detection3DPC], center: np.ndarray) -> float: - directions = [] - for det in members: - v = _camera_position(det) - center - norm = np.linalg.norm(v) - if norm > 1e-6: - directions.append(v / norm) - if not directions: - return 0.0 - dirs = np.stack(directions) - azimuth = np.arctan2(dirs[:, 1], dirs[:, 0]) - bins = set(((azimuth + math.pi) / (2 * math.pi) * 8).astype(int) % 8) - return len(bins) / 8.0 - - -def _axes_observed(members: list[Detection3DPC], center: np.ndarray) -> tuple[bool, bool, bool]: - directions = [] - for det in members: - v = _camera_position(det) - center - norm = np.linalg.norm(v) - if norm > 1e-6: - directions.append(v / norm) - if not directions: - return (False, False, False) - dirs = np.stack(directions) - return tuple(bool((np.abs(dirs[:, i]) > 0.3).any()) for i in range(3)) # type: ignore[return-value] - - def _lift( - detections: ImageDetections2D[Any], - rig: Rig, - policy: LocalizePolicy, - plane: Any | None = None, + detections: ImageDetections2D[Any], rig: Rig, policy: LocalizePolicy, plane: Any | None ) -> list[Detection3DPC]: """Gate a frame's lifted detections.""" - pose = rig.camera_pose(detections.ts) - if pose is None: - return [] - camera = np.array([pose.position.x, pose.position.y, pose.position.z]) - lifted = rig.lift(detections, plane) + pose = rig.camera_pose(detections.ts, detections.image.frame_id) + lifted = rig.lift(detections, plane) if pose is not None else None if lifted is None: return [] + camera = np.array([pose.position.x, pose.position.y, pose.position.z]) - floor = policy.min_depth_points if rig.cloud is None else CLOUD_MIN_POINTS valid: list[Detection3DPC] = [] for det3d in lifted: points = np.asarray(det3d.pointcloud.pointcloud.points) - if len(points) < floor: + if len(points) < policy.min_points: continue - extent = points.max(axis=0) - points.min(axis=0) - if float(extent.max()) > policy.max_object_extent_m: + if float((points.max(axis=0) - points.min(axis=0)).max()) > policy.max_object_extent_m: continue - ranges = np.linalg.norm(points - camera, axis=1) - if float(np.median(ranges)) < policy.min_camera_range_m: + if float(np.median(np.linalg.norm(points - camera, axis=1))) < policy.min_camera_range_m: continue if plane is not None: heights = rig.support_heights(detections.ts, plane, points) - low = float(np.quantile(heights, 0.05)) - high = float(np.quantile(heights, 0.95)) - if low > policy.surface_patch_min_drop_m and high < policy.surface_patch_max_rise_m: + # A cloud hugging the support is a patch of the surface, not an object. + if ( + float(np.quantile(heights, 0.05)) > policy.surface_patch_min_drop_m + and float(np.quantile(heights, 0.95)) < policy.surface_patch_max_rise_m + ): continue valid.append(det3d) return valid @@ -221,19 +226,27 @@ def embed_index( ) -> Stream[Any, Any]: """SigLIP-embedded, world-posed frame index at the rig's embed rate. - Built once per window and handed to every ``localize`` call on it: the - embed forwards are what a second query would otherwise repeat. + Built once per window and handed to every ``localize`` call on it. Only the + pose, the embedding and the sharpness are kept; the pixels are re-read from + the color stream for the handful of frames that reach the detector, so the + index costs bytes per frame instead of a decoded image. """ rig = rig or Rig.from_store(store) - posed = ( + embedded: Stream[Any, Any] = ( rig.color.after(t0) .before(t1) .filter(lambda obs: obs.data.brightness > 0.1) .transform(QualityWindow(lambda img: img.sharpness, window=1.0 / rig.embed_hz)) .map(lambda obs: obs.derive(data=obs.data, pose=rig.index_pose(obs))) .filter(lambda obs: obs.pose is not None) + .transform(EmbedImages(siglip)) + .map( + lambda obs: obs.derive( + data=IndexFrame(obs.ts, obs.data.frame_id, float(obs.data.sharpness)) + ) + ) + .materialize() ) - embedded: Stream[Any, Any] = posed.transform(EmbedImages(siglip)).materialize() logger.info(f"index: {embedded.count()} frames embedded over {t1 - t0:.1f}s") return embedded @@ -250,37 +263,20 @@ def localize( require_pose: bool = True, policy: LocalizePolicy | None = None, trace: LocalizeTrace | list[LocalizeTrace] | None = None, - identity_store: IdentityStore | None = None, + groups: dict[str, Groups] | None = None, ) -> list[Localization] | list[list[Localization]]: """Every verified 3D instance of *query*, latest-seen first. - Each instance's ``point_cloud`` is the union of every viewpoint that saw - it, and its position follows the latest sighting. An empty list is a - first-class answer: nothing reached the accept score, no support was - confirmed from a second viewpoint, or the best candidate had no valid - depth and ``require_pose`` holds. Coexisting instances of one label are - all returned; each carries ``ambiguity_margin`` against its rivals and - is flagged below ``refusal_margin`` - never a silent guess. - - A list *query* shares one detection pass: every label's semantic peaks - mark its sightings, each peak takes the frames adjacent to it in time - until the verification policy's viewpoints are covered, and each unique - selected frame is scored against every label, segmented and lifted once. - One instance list per label, in input order. ``trace`` then takes a list - of the same length. - - The index, the rig and the three models belong to the caller: nothing - here is loaded or stopped, so one process can call this repeatedly on - warm weights, and every query on one window reuses the same embeddings. - The window is the index's - build it with :func:`embed_index`. Without a - ``rig`` the store's shape decides one, and without a ``policy`` the rig - supplies scale-appropriate defaults. - - An ``identity_store`` makes evidence cumulative: each label's groups - persist across calls, frames the store already ingested for a label are - skipped entirely, and an object stays answerable after it leaves the - window, its position still following the latest sighting. Without one, - every call verifies from scratch inside its window. + An empty list is a first-class answer: nothing reached the accept score, no + support was confirmed from a second viewpoint, or the best candidate had no + valid depth and ``require_pose`` holds. Coexisting instances of one label + are all returned, each carrying ``ambiguity_margin`` against its rivals. + + A list *query* shares one detection pass and returns one list per label, in + input order; ``trace`` then takes a list of the same length. The index, the + rig and the models belong to the caller. A ``groups`` dict makes evidence + cumulative across calls: frames already ingested for a label are skipped, + and an object stays answerable after it leaves the window. """ rig = rig or Rig.from_store(store) policy = policy or rig.default_localize_policy() @@ -289,14 +285,14 @@ def localize( traces: list[LocalizeTrace | None] = ( list(trace) if isinstance(trace, list) else [trace] * len(queries) ) + state = [(groups if groups is not None else {}).setdefault(q, Groups()) for q in queries] index_count = index.count() settled = _settled(index, policy.settled_window_fraction / rig.embed_hz) source = index.filter(lambda obs: obs.id in settled) candidate_ids: set[int] = set() expanded: set[float] = set() - anchor_x = 0.0 - anchor_y = 0.0 + anchor = np.zeros(2) anchor_count = 0 for q in queries: @@ -306,7 +302,7 @@ def localize( .order_by("ts") .transform( peaks( - key=_similarity, + key=lambda obs: float(obs.similarity), prominence=policy.peak_prominence, distance=policy.peak_distance_s, width=policy.peak_width_s, @@ -320,23 +316,18 @@ def localize( query_embedding, k=policy.tail_k ) ) - peak_count = 0 for peak in sightings: - peak_pose = cast("PoseTuple", peak.pose_tuple) - peak_count += 1 + anchor += cast("PoseTuple", peak.pose_tuple)[:2] anchor_count += 1 - anchor_x += peak_pose[0] - anchor_y += peak_pose[1] candidate_ids.add(peak.id) if peak.ts in expanded: continue expanded.add(peak.ts) gathered: Stream[Any, Any] = source.near( peak.pose_stamped, radius=policy.verify_radius_m - ).transform(QualityWindow(lambda img: img.sharpness, window=policy.verify_window_s)) - for obs in gathered: - candidate_ids.add(obs.id) - logger.info(f"localize {q!r}: {peak_count} semantic peaks of {index_count} embedded") + ).transform(QualityWindow(lambda f: f.sharpness, window=policy.verify_window_s)) + candidate_ids.update(obs.id for obs in gathered) + logger.info(f"localize {q!r}: {len(sightings)} semantic peaks of {index_count} embedded") candidates = index.filter(lambda obs: obs.id in candidate_ids).order_by("ts") candidate_count = len(candidate_ids) @@ -346,71 +337,48 @@ def localize( if candidate_count: from dimos.perception.memory.support_plane import fit_support_plane - mx = anchor_x / anchor_count - my = anchor_y / anchor_count - cell = (round(mx / policy.plane_cell_m), round(my / policy.plane_cell_m)) - plane = rig._plane_cache.get(cell) + mean = anchor / anchor_count + cell = (round(mean[0] / policy.plane_cell_m), round(mean[1] / policy.plane_cell_m)) + plane = rig.plane_cache.get(cell) if plane is None: stride = max(1, candidate_count // policy.plane_keyframes) - keyframes = [] - for i, obs in enumerate(candidates): - if i % stride == 0: - keyframes.append(obs) - if len(keyframes) == policy.plane_keyframes: - break - plane = fit_support_plane(rig, keyframes) - if plane is not None: - rig._plane_cache[cell] = plane - - if identity_store is None: - entries = None - identities = [ - Identity(is_same=spatial(policy.cluster_radius_m), merge=fused(policy.fuse_voxel_m)) - for _ in queries - ] - ingested: list[set[float]] = [set() for _ in queries] - ungrounded: list[tuple[float, float] | None] = [None] * len(queries) # (score, ts) - else: - entries = [ - identity_store.get_or_create( - q, spatial(policy.cluster_radius_m), fused(policy.fuse_voxel_m) + plane = fit_support_plane( + rig, + [obs for i, obs in enumerate(candidates) if i % stride == 0][ + : policy.plane_keyframes + ], ) - for q in queries - ] - identities = [entry.identity for entry in entries] - ingested = [entry.ingested for entry in entries] - ungrounded = [entry.ungrounded for entry in entries] + if plane is not None: + rig.plane_cache[cell] = plane floor = policy.candidate_floor cache = detector.score_cache def _detect(upstream: Iterator[Any]) -> Iterator[Any]: for obs in upstream: - active = [j for j in range(len(queries)) if obs.ts not in ingested[j]] + active = [j for j in range(len(queries)) if obs.ts not in state[j].ingested] if not active: continue - if any((obs.ts, queries[j], floor) not in cache for j in active): - boxes, rows = detector.query_score_rows_batch([obs.data], queries, threshold=floor)[ - 0 - ] + keys = [(obs.ts, queries[j], floor) for j in active] + img = rig.frame_at(obs) if any(k not in cache for k in keys) else None + if img is not None: + boxes, rows = detector.query_score_rows_batch([img], queries, threshold=floor)[0] for j, q in enumerate(queries): keep = rows[:, j] >= floor cache[(obs.ts, q, floor)] = (boxes[keep], rows[keep, j]) if len(cache) > _SCORE_CACHE_MAX: cache.popitem(last=False) - - rows_per_label: list[tuple[int, tuple[Any, Any]]] = [] - for j in active: - key = (obs.ts, queries[j], floor) + for key in keys: cache.move_to_end(key) - rows_per_label.append((j, cache[key])) - ingested[j].add(obs.ts) - if not any(len(scores) for _j, (_boxes, scores) in rows_per_label): + for j in active: + state[j].ingested.add(obs.ts) + if not any(len(cache[key][1]) for key in keys): continue - img = obs.data + img = img if img is not None else rig.frame_at(obs) + detections: list[Detection2DBBox] = [] - for j, (boxes, scores) in rows_per_label: - for box, score in zip(boxes, scores, strict=True): + for j, key in zip(active, keys, strict=True): + for box, score in zip(*cache[key], strict=True): det = Detection2DBBox( bbox=(float(box[0]), float(box[1]), float(box[2]), float(box[3])), track_id=len(detections), @@ -431,21 +399,20 @@ def _ingest(upstream: Iterator[Any]) -> Iterator[Any]: lifted = _lift(frame, rig, policy, plane) grounded = {det3d.track_id for det3d in lifted} for det2d in frame: - j = det2d.class_id - best = ungrounded[j] + group = state[det2d.class_id] + best = group.ungrounded if det2d.track_id not in grounded and (best is None or det2d.confidence > best[0]): - ungrounded[j] = (det2d.confidence, det2d.ts) + group.ungrounded = (det2d.confidence, det2d.ts) for det3d in lifted: - label_trace = traces[det3d.class_id] - if label_trace is not None: - label_trace.matched.append((det3d.ts, det3d)) - identities[det3d.class_id].add(det3d) - for j in range(len(queries)): - label_trace = traces[j] - if label_trace is None: - continue + j = det3d.class_id + slot = state[j].add(det3d, policy.cluster_radius_m, policy.fuse_voxel_m) + if traces[j] is not None: + state[j].trace[slot].append(det3d) + if traces[j].first_match_ts is None: # type: ignore[union-attr] + traces[j].first_match_ts = det3d.ts # type: ignore[union-attr] + for j, label_trace in enumerate(traces): label_dets = [det for det in frame if det.class_id == j] - if label_dets: + if label_trace is not None and label_dets: label_trace.detection_frames.append( obs.derive(data=ImageDetections2D(image=frame.image, detections=label_dets)) ) @@ -455,148 +422,145 @@ def _ingest(upstream: Iterator[Any]) -> Iterator[Any]: lambda obs: obs.derive(data=segmenter.segment(obs.data)) ).transform(_ingest).drain() - if entries is not None: - for entry, best in zip(entries, ungrounded, strict=True): - entry.ungrounded = best - results = [ - _finalize( - q, - identity=identities[j], - ungrounded_best=ungrounded[j], - rig=rig, - require_pose=require_pose, - policy=policy, - trace=traces[j], - ) - for j, q in enumerate(queries) + _finalize(q, state[j], rig, require_pose, policy, traces[j]) for j, q in enumerate(queries) ] return results[0] if isinstance(query, str) else results +def _orientation(cloud: PointCloud2) -> tuple[float, float, float, float]: + from scipy.spatial.transform import Rotation + + try: + x, y, z, w = Rotation.from_matrix(np.asarray(cloud.oriented_bounding_box.R)).as_quat() + except Exception: + return (0.0, 0.0, 0.0, 1.0) + return (float(x), float(y), float(z), float(w)) + + def _finalize( query: str, - *, - identity: Identity, - ungrounded_best: tuple[float, float] | None, + group: Groups, rig: Rig, require_pose: bool, policy: LocalizePolicy, trace: LocalizeTrace | None, ) -> list[Localization]: + rows = [group.rows(i) for i in range(len(group.centers))] + views = [len(np.unique(np.round(r[:, M_KX : M_KZ + 1], 2), axis=0)) for r in rows] verified = [ - (merged, members) - for merged, members in zip(identity.merged, identity.groups, strict=True) - if _max_score(members) >= policy.accept_score and _n_views(members) >= policy.min_views + i + for i, r in enumerate(rows) + if r[:, M_SCORE].max() >= policy.accept_score and views[i] >= policy.min_views ] logger.info( f"verification {query!r}: " + ", ".join( - f"score={_max_score(g):.2f} views={_n_views(g)} obs={len(g)}" for g in identity.groups + f"score={r[:, M_SCORE].max():.2f} views={views[i]} obs={len(r)}" + for i, r in enumerate(rows) ) ) - if trace is not None: - for _merged, members in verified: - for det3d in members: - trace.verified.append((det3d.ts, det3d)) if not verified: - if ungrounded_best is not None and ungrounded_best[0] >= policy.accept_score: - if require_pose: - logger.info(f"{query!r}: best candidate has no valid depth and require_pose is set") - return [] - score, ts = ungrounded_best - return [ - Localization( - instance_id="query-0", - semantic_score=score, - identity_score=0.0, - ambiguity_margin=1.0, - position_world_xyz=None, - orientation_world_xyzw=None, - frame_id=rig.world_frame, - support=None, - pose_timestamp=ts, - geometry_timestamp=ts, - last_seen_timestamp=ts, - point_cloud=None, - coverage=0.0, - n_views=1, - reason="no_valid_depth", - ) - ] - return [] + best = group.ungrounded + if best is None or best[0] < policy.accept_score: + return [] + if require_pose: + logger.info(f"{query!r}: best candidate has no valid depth and require_pose is set") + return [] + score, ts = best + return [ + Localization( + instance_id="query-0", + semantic_score=score, + identity_score=0.0, + ambiguity_margin=1.0, + position_world_xyz=None, + orientation_world_xyzw=None, + frame_id=rig.world_frame, + support=None, + pose_timestamp=ts, + geometry_timestamp=ts, + last_seen_timestamp=ts, + point_cloud=None, + coverage=0.0, + n_views=1, + reason="no_valid_depth", + ) + ] - verified.sort(key=lambda pair: _latest(pair[1]).ts, reverse=True) + verified.sort(key=lambda i: rows[i][:, M_TS].max(), reverse=True) instances: list[Localization] = [] - for k, (merged, members) in enumerate(verified): - m_lo, m_hi = _interval(members) - rival_scores = [ - _max_score(others) - for _m, others in verified - if others is not members - and not (_interval(others)[1] < m_lo or _interval(others)[0] > m_hi) # coexisting + for k, i in enumerate(verified): + mine = rows[i] + score = float(mine[:, M_SCORE].max()) + lo, hi = float(mine[:, M_TS].min()), float(mine[:, M_TS].max()) + rivals = [ + float(rows[j][:, M_SCORE].max()) + for j in verified + if j != i + and not (rows[j][:, M_TS].max() < lo or rows[j][:, M_TS].min() > hi) # coexisting ] - margin = _max_score(members) - max(rival_scores) if rival_scores else 1.0 - reason = ( - "ambiguous_between_coexisting_candidates" if margin < policy.refusal_margin else None - ) + margin = score - max(rivals) if rivals else 1.0 - latest = _latest(members) - union = merged.pointcloud + union = group.clouds[i] points = np.asarray(union.pointcloud.points) aabb_min, aabb_max = points.min(axis=0), points.max(axis=0) - try: - orientation = _quaternion_from_matrix( - np.asarray(latest.pointcloud.oriented_bounding_box.R) + centroids = mine[:, M_CX : M_CZ + 1] + offsets = mine[:, M_KX : M_KZ + 1] - centroids.mean(axis=0) # center to camera + norms = np.linalg.norm(offsets, axis=1) + dirs = offsets[norms > 1e-6] / norms[norms > 1e-6, None] + if len(dirs): + octant = ((np.arctan2(dirs[:, 1], dirs[:, 0]) + math.pi) / (2 * math.pi) * 8).astype( + int ) - except Exception: - orientation = (0.0, 0.0, 0.0, 1.0) + coverage = len(set(octant % 8)) / 8.0 + axes = tuple(bool((np.abs(dirs[:, a]) > 0.3).any()) for a in range(3)) + else: + coverage, axes = 0.0, (False, False, False) + center = (aabb_min + aabb_max) / 2 extent = np.maximum(aabb_max - aabb_min, 0.005) - sigma = ( - np.stack([_centroid(d) for d in members]).std(axis=0) - if len(members) > 1 - else np.full(3, 0.01) - ) - group_center = _group_center(members) - support = Support( - center_xyz=(float(center[0]), float(center[1]), float(center[2])), - extent_xyz_m=(float(extent[0]), float(extent[1]), float(extent[2])), - orientation_xyzw=(0.0, 0.0, 0.0, 1.0), - sigma_xyz_m=(float(sigma[0]), float(sigma[1]), float(sigma[2])), - coverage=_azimuth_coverage(members, group_center), - axes_observed=_axes_observed(members, group_center), - frame_id=rig.world_frame, - ) - + sigma = centroids.std(axis=0) if len(mine) > 1 else np.full(3, 0.01) if trace is not None: - trace.answers.append(members) + trace.answers.append(group.trace[i]) if k == 0: - trace.backdrop_ts = latest.ts + trace.backdrop_ts = hi - latest_centroid = _centroid(latest) + newest = mine[mine[:, M_TS].argmax()] instances.append( Localization( instance_id=f"query-{k}", - semantic_score=_max_score(members), - identity_score=min(1.0, _n_views(members) / 4.0), + semantic_score=score, + identity_score=min(1.0, views[i] / 4.0), ambiguity_margin=margin, position_world_xyz=( - float(latest_centroid[0]), - float(latest_centroid[1]), - float(latest_centroid[2]), + float(newest[M_CX]), + float(newest[M_CY]), + float(newest[M_CZ]), ), - orientation_world_xyzw=orientation, + orientation_world_xyzw=_orientation(group.latest[i]), frame_id=rig.world_frame, - support=support, - pose_timestamp=latest.ts, - geometry_timestamp=latest.ts, - last_seen_timestamp=latest.ts, + support=Support( + center_xyz=tuple(float(v) for v in center), # type: ignore[arg-type] + extent_xyz_m=tuple(float(v) for v in extent), # type: ignore[arg-type] + orientation_xyzw=(0.0, 0.0, 0.0, 1.0), + sigma_xyz_m=tuple(float(v) for v in sigma), # type: ignore[arg-type] + coverage=coverage, + axes_observed=cast("tuple[bool, bool, bool]", axes), + frame_id=rig.world_frame, + ), + pose_timestamp=hi, + geometry_timestamp=hi, + last_seen_timestamp=hi, point_cloud=union, - coverage=support.coverage, - n_views=_n_views(members), - reason=reason, + coverage=coverage, + n_views=views[i], + reason=( + "ambiguous_between_coexisting_candidates" + if margin < policy.refusal_margin + else None + ), ) ) return instances diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index 3288bf671b..bef87de093 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -12,29 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The sensor rig behind a recording: intrinsics, poses, and 3D geometry. - -Two independent axes generalize the perception stack beyond one robot: - -* **Pose source.** The world-to-optical transform comes from a recorded - ``tf`` stream, or - for recordings without one - from the world pose - stamped on each observation plus a static base-to-optical mount. -* **Geometry source.** 3D geometry comes from an aligned ``depth`` stream, - unprojected per detection mask, or from a world-frame pointcloud stream - (a registered lidar), projected through the camera per detection mask. - The cloud at a timestamp merges the scans in a short window around it: - fresh scans concatenate whole; snapshots re-reporting each other's exact - points (a rolling map) merge nearest-first with cleared cells honored - (see ``Rig.cloud_at``). - -``Rig.from_store`` recognizes both recording shapes; every field can also -be supplied directly for live stores whose streams are still filling. +"""The sensor rig behind a run: intrinsics, poses, and 3D geometry. + +Three independent axes generalize the stack beyond one robot: pose from a +``tf`` stream or from stamped world poses plus a static mount; geometry from +an aligned ``depth`` stream or a world-frame pointcloud; and intrinsics and +mounts keyed by optical frame, each lookup taking the frame off the image it +resolves rather than off a single rig-wide field. + +:meth:`Rig.from_store` recognizes a recording's shape and self-calibrates what +it omits; a robot that knows its calibration builds a ``Rig`` field by field +and never touches it. """ from __future__ import annotations from collections import OrderedDict -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace import json from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -42,24 +36,20 @@ import numpy as np from dimos.memory.tf import StreamTF -from dimos.memory.transform import Transformer +from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.perception.detection.project import sees as project_sees +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC from dimos.perception.detection.type.detection3d.pointcloud import lattice_quantum from dimos.perception.detection.type.detection3d.pointcloud_filters import ( range_cluster, statistical, ) -from dimos.perception.memory import gates -from dimos.perception.memory.gates import SPEED_MAX, STILL_ENVELOPE, TF_TOLERANCE from dimos.perception.memory.support_plane import PLANE_DISTANCE_CLOUD -from dimos.perception.memory.types import InventoryPolicy, LocalizePolicy +from dimos.perception.memory.types import LocalizePolicy from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from collections.abc import Callable, Iterator - from dimos_lcm.sensor_msgs import CameraInfo from dimos.memory.type.observation import Observation @@ -73,62 +63,40 @@ logger = setup_logger() -# Pose-stamped rigs ride per-frame odometry, so walking does not stale the -# projection the way a sweeping wrist stales interpolated tf; the gate only -# drops speed glitches and sprints. -WALK_SPEED_MAX = 1.5 +OPTICAL_FRAME = "camera_color_optical_frame" # only when a store carries no calibration +WORLD_FRAME = "world" +TF_TOLERANCE = 0.12 # s - one world-pose period plus margin DEPTH_TOLERANCE = 0.06 # s - temporal join color->depth -# Scans within this of a frame form its geometry. Wide enough that a -# spinning lidar's near-floor blind ring is filled by scans taken from -# earlier and later poses - the scene is static in world frame. +# Scans within this of a frame form its geometry. Wide enough that a spinning +# lidar's near-floor blind ring is filled by scans taken from earlier and +# later poses - the scene is static in world frame. CLOUD_ACCUM_S = 4.0 +FRAME_TOLERANCE = 0.02 # s - an index timestamp back to its own color frame _SCAN_CACHE_MAX = 256 # registered scans held per rig; a window needs a few dozen +_CLOUD_CACHE_MAX = 8 # merged windows held; adjacent frames re-ask for the same one -EMBED_HZ = 1.0 # index density for a wrist camera parked over a workspace -# A walking robot changes viewpoint every frame and its frames blur -# unevenly, so the index must sample denser for retrieval to catch the -# sharp sightings. -WALK_EMBED_HZ = 3.0 - -# Color-stream delay estimation: image timestamps stamped at receive lag the -# pose source by a constant the recording itself reveals - the lag that best -# correlates optical-flow yaw rate with the camera's heading rate. The grid -# is the searched span and its resolution; a true delay outside the span -# lands on an edge and is refused rather than clamped. -DELAY_LAGS = np.arange(-0.5, 0.5001, 0.01) -# Flow is sampled in short windows spread over the recording - a compute -# budget, like the pose samples of _camera_span. -DELAY_WINDOWS = 12 -DELAY_WINDOW_S = 2.5 +EMBED_HZ = 1.0 # index density for a camera parked over a workspace +WALK_EMBED_HZ = 3.0 # a walking robot changes viewpoint every frame +MOBILE_SPAN_M = 3.0 # camera translation beyond this means a mobile base # Sparse projected clouds: split off background seen through the mask, then a -# loose outlier trim. The dense-cloud defaults (raycast + radius) assume a -# density registered lidar does not have. +# loose outlier trim. The dense-cloud defaults assume a density a registered +# lidar does not have. _CLOUD_TRIM_NEIGHBORS = 12 _CLOUD_LIFT_FILTERS = [ range_cluster(), statistical(nb_neighbors=_CLOUD_TRIM_NEIGHBORS, std_ratio=2.0), ] -# The projected lift's evidence floor. Depth-pixel counts and lattice-cell -# counts do not compare - a bottle-sized object can never cover thirty 5 cm -# cells - so a policy's depth-scale floor does not apply to a projected -# cloud; the floor there is the smallest cloud the trim can vet, its own -# neighborhood. +# A projected lift's evidence floor is the smallest cloud the trim can vet, +# its own neighborhood: lattice-cell counts and depth-pixel counts do not +# compare, so a depth-scale floor does not apply to one. CLOUD_MIN_POINTS = _CLOUD_TRIM_NEIGHBORS + 1 - -def _column_keys(points: np.ndarray, quantum: float, anchor: np.ndarray) -> np.ndarray: - """Packed XY lattice-cell key per point, anchored so any grid phase maps exactly.""" - cells = np.round((points[:, :2] - anchor) / quantum).astype(np.int64) + (1 << 20) - keys: np.ndarray = (cells[:, 0] << 21) | cells[:, 1] - return keys - - -# Room-scale policies for mobile-robot rigs: objects are furniture-sized, -# viewpoints meters apart, odometry drifts centimeters between passes, and -# registered lidar is noisier and sparser than wrist-camera depth. +# Room-scale defaults for mobile rigs: objects are furniture-sized, viewpoints +# meters apart, odometry drifts centimeters, and registered lidar is noisier +# and sparser than wrist-camera depth. ROOM_LOCALIZE_POLICY = LocalizePolicy( candidate_floor=0.18, accept_score=0.32, @@ -137,68 +105,73 @@ def _column_keys(points: np.ndarray, quantum: float, anchor: np.ndarray) -> np.n max_object_extent_m=2.0, surface_patch_max_rise_m=0.08, surface_patch_min_drop_m=-0.06, - min_depth_points=30, + min_points=30, min_camera_range_m=0.5, fuse_voxel_m=0.03, ) -MOBILE_SPAN_M = 3.0 # camera translation beyond this means a mobile base +def _column_keys(points: np.ndarray, quantum: float, anchor: np.ndarray) -> np.ndarray: + """Packed XY lattice-cell key per point, anchored so any grid phase maps exactly.""" + cells = np.round((points[:, :2] - anchor) / quantum).astype(np.int64) + (1 << 20) + keys: np.ndarray = (cells[:, 0] << 21) | cells[:, 1] + return keys ROOT_PROBES = 24 # instants a frame is probed at before it counts as unreachable +# Color timestamps stamped at receive lag the pose source by a constant the +# recording reveals: the lag that best correlates optical-flow yaw rate with +# the camera's heading rate. The grid is the searched span and its resolution; +# a true delay outside it lands on an edge and is refused rather than clamped. +DELAY_LAGS = np.arange(-0.5, 0.5001, 0.01) +DELAY_WINDOWS = 12 # flow is sampled in short windows spread over the recording +DELAY_WINDOW_S = 2.5 -def _tf_root(store: Any, tf_name: str, tf: StreamTF, optical_frame: str) -> str | None: - """The tf tree's root frame: a parent that is never a child, among the - frames the camera actually reaches. +def _rate(stream: Any) -> float: + count: int = stream.count() + if count < 2: + return float(count) + t0, t1 = stream.get_time_range() + return count / max(float(t1 - t0), 1e-6) - A recording can carry an anchor edge published once at each end of the - run, above the frame every other transform is stamped in. It is a root the - camera never reaches through, and taking it strands every pose lookup, so - frames no probe resolves are dropped before the root is taken. Probes sit - at bin midpoints, where an anchor stamped at the ends cannot answer. + +def _tf_root(store: Any, tf_name: str, tf: StreamTF, optical: str) -> str | None: + """The tf root: a parent that is never a child, among reachable frames. + + A recording can carry an anchor edge published once at each end of the run, + above the frame everything else is stamped in. Taking it strands every pose + lookup, so frames no probe resolves are dropped first; probes sit at bin + midpoints, where an anchor stamped at the ends cannot answer. """ - edges: set[tuple[str, str]] = set() - for obs in store.stream(tf_name): - for transform in obs.data.transforms: - edges.add((transform.frame_id, transform.child_frame_id)) + edges = { + (t.frame_id, t.child_frame_id) for obs in store.stream(tf_name) for t in obs.data.transforms + } if not edges: return None # live store, nothing recorded yet frames = {frame for edge in edges for frame in edge} t0, t1 = store.stream(tf_name).get_time_range() - reached = {optical_frame} + reached = {optical} for k in range(ROOT_PROBES): ts = t0 + (t1 - t0) * (k + 0.5) / ROOT_PROBES for frame in frames - reached: - if tf.get(optical_frame, frame, ts, TF_TOLERANCE, warn=False) is not None: + if tf.get(optical, frame, ts, TF_TOLERANCE, warn=False) is not None: reached.add(frame) linked = [(p, c) for p, c in edges if p in reached and c in reached] roots = {p for p, _ in linked} - {c for _, c in linked} return roots.pop() if len(roots) == 1 else None -def _stream_rate(stream: Any) -> float: - count: int = stream.count() - if count < 2: - return float(count) - t0, t1 = stream.get_time_range() - return count / max(float(t1 - t0), 1e-6) - - def _camera_span(rig: Rig) -> float: """Diagonal of the camera positions' bounding box over the recording.""" if rig.tf is None and rig.poses is None: - return 0.0 # no pose source at all: an embed-only store being seeded + return 0.0 try: t0, t1 = rig.color.get_time_range() except LookupError: return 0.0 # live store, nothing recorded yet - positions = [] - for k in range(12): - pose = rig.camera_pose(t0 + (t1 - t0) * k / 11) - if pose is not None: - positions.append([pose.position.x, pose.position.y, pose.position.z]) + seen = [rig.camera_pose(t0 + (t1 - t0) * k / 11) for k in range(12)] + positions = [[p.position.x, p.position.y, p.position.z] for p in seen if p is not None] if len(positions) < 2: return 0.0 spread = np.array(positions) @@ -206,12 +179,11 @@ def _camera_span(rig: Rig) -> float: def _heading_series(rig: Rig, spans: list[tuple[float, float]]) -> tuple[np.ndarray, np.ndarray]: - """Camera heading rate over the given time spans: (rate midpoints, rates). + """Camera heading rate over the given spans: (rate midpoints, rates). - A poses stream is read directly - a rigid mount adds a constant offset, - so the base yaw rate is the camera heading rate. A tf rig has no pose - stream to iterate; the optical axis is sampled through tf on each span's - frame-period grid instead. + A rigid mount adds a constant offset, so a poses stream's base yaw rate is + the camera heading rate. A tf rig has no pose stream; the optical axis is + sampled through tf on each span's frame-period grid instead. """ times: list[float] = [] headings: list[float] = [] @@ -243,26 +215,22 @@ def _heading_series(rig: Rig, spans: list[tuple[float, float]]) -> tuple[np.ndar def estimate_color_delay(rig: Rig) -> float: - """Constant lag of color timestamps behind the rig's pose source, in seconds. - - Optical-flow yaw rate between consecutive frames is compared against the - camera heading rate at a grid of candidate lags; the lag maximizing - their absolute correlation is the delay (the sign of the relation is - mount-dependent). The estimate validates itself on the recording: the - full sample and its two interleaved halves must estimate mutually equal - lags to within one frame period - each flow sample integrates motion - over a frame period, so that is the measurement's own resolution - and - every peak must be interior to the searched span. A recording without - usable rotation fails that and keeps 0.0. + """Constant lag of color timestamps behind the pose source, in seconds. + + Optical-flow yaw rate is correlated against camera heading rate over a grid + of candidate lags; the strongest is the delay. The estimate validates + itself: the full sample and its two interleaved halves must agree to within + one frame period - the measurement's own resolution - and every peak must + be interior to the searched span. Without usable rotation it keeps 0. """ import cv2 try: t0, t1 = rig.color.get_time_range() except LookupError: - return 0.0 # live store, nothing recorded yet + return 0.0 - fx = rig.camera_info.K[0] + fx = rig.cameras[rig.optical_frame].K[0] flow_ts: list[float] = [] flow_rate: list[float] = [] flow_dt: list[float] = [] @@ -331,99 +299,106 @@ def lag_of(sel: np.ndarray) -> float | None: return estimates[0] -class RegisterScans(Transformer["PointCloud2", "PointCloud2"]): - """Map sensor-frame scans into a world frame through tf, one transform per scan. +def _images(store: Any, names: list[str]) -> tuple[dict[str, str], list[tuple[str, str]]]: + """Color stream per optical frame, and every metric-depth stream as (frame, name). - A plain stream transformer, so the registered cloud is a derived stream: - ``scans.transform(RegisterScans(tf, world))`` yields world-frame clouds, - ``.save(...)`` persists them, and a live pipeline can tail-register the - same way the embed pipeline does. Scans already in the world frame pass - through untouched; scans with no transform at their time are dropped. + Several color streams in one frame (a recording that also stored a derived + feed) resolve to the highest-rate one. """ + color: dict[str, str] = {} + depth: list[tuple[str, str]] = [] + for name in names: + stream = store.stream(name) + if stream.count() == 0: + continue + image = stream.first().data + frame, data = image.frame_id, image.data + if data.dtype == np.uint16 or data.dtype.kind == "f": + depth.append((frame, name)) + elif ( + data.ndim == 3 + and data.shape[2] == 3 + and not np.array_equal(data[..., 0], data[..., 1]) + ): + held = color.get(frame) + if held is None or _rate(stream) > _rate(store.stream(held)): + color[frame] = name + else: + logger.info(f"rig: image stream {name!r} is neither color nor metric depth") + return color, depth - def __init__(self, tf: TFLookup, world_frame: str, tolerance: float = TF_TOLERANCE) -> None: - self.tf = tf - self.world_frame = world_frame - self.tolerance = tolerance - - def __call__( - self, upstream: Iterator[Observation[PointCloud2]] - ) -> Iterator[Observation[PointCloud2]]: - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - for obs in upstream: - if obs.data.frame_id == self.world_frame: - yield obs - continue - transform = self.tf.get(self.world_frame, obs.data.frame_id, obs.ts, self.tolerance) - if transform is None: - continue - matrix = transform.to_matrix() - points = obs.data.as_numpy()[0] @ matrix[:3, :3].T + matrix[:3, 3] - yield obs.derive( - data=PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=obs.ts) - ) +def _cameras(store: Any, roles: dict[str, Any], names: list[str], types: dict[str, type]): + """Intrinsics per optical frame: an inline manifest dict, or every + CameraInfo stream keyed by the frame it calibrates. - -ROOM_INVENTORY_POLICY = InventoryPolicy( - keyframe_stride_s=1.25, - min_mask_area_px=900, - min_depth_points=30, - max_object_extent_m=2.0, - min_height_above_plane_m=0.08, - band_above_plane_m=(-0.05, 1.5), - min_camera_range_m=0.5, - envelope_pad_m=0.12, - search_radius_m=0.8, - size_gap_max_m=0.8, - support_explained=0.25, - name_attach_iou=0.20, - same_frame_merge_gap_m=0.10, - split_extent_m=1.2, - split_height_m=0.9, - split_eps_m=0.10, -) + Several infos in one frame (a colour/depth pair) resolve by name order, so + the colour one wins; infos in different frames are different cameras. + """ + from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo as CameraInfoMsg + + role = roles.get("camera_info") + if isinstance(role, dict): + info = CameraInfoMsg( + height=role["height"], + width=role["width"], + distortion_model=role.get("distortion_model", ""), + D=role.get("D"), + K=role["K"], + R=role.get("R"), + P=role.get("P"), + frame_id=role["frame_id"], + ) + return {info.frame_id: info} + found = ( + [role] + if isinstance(role, str) + else sorted(n for n in names if types[n] is CameraInfoMsg and store.stream(n).count()) + ) + cameras = {} + for name in found: # sorted, so a color info wins over its depth twin + info = store.stream(name).first().data + cameras.setdefault(info.frame_id, info) + return cameras @dataclass class Rig: - """Everything the stack needs to go from a 2D mask to world geometry. - - Exactly one of ``tf`` / (``base_to_optical`` + ``poses``) provides the - world-to-optical transform, and exactly one of ``depth`` / ``cloud`` - provides 3D geometry. + """2D mask to world geometry. + + Exactly one of ``tf`` / (``mounts`` + ``poses``) gives the world-to-optical + transform, and one of ``depth`` / ``cloud`` gives geometry. ``cameras`` and + ``mounts`` key on optical frame id, resolved per image by + :meth:`camera_frame`. ``color`` is one stream, so a rig with several + cameras needs their frames registered here and their feeds interleaved + into it by whoever builds the rig. """ - camera_info: CameraInfo + cameras: dict[str, CameraInfo] color: Any # color_image stream world_frame: str - optical_frame: str tf: TFLookup | None = None - base_to_optical: Transform | None = None + mounts: dict[str, Transform] = field(default_factory=dict) poses: Any = None # stream carrying world base poses (e.g. odom) depth: Any = None # aligned depth stream, lifted via from_depth cloud: Any = None # pointcloud stream, lifted via from_2d after registration tf_tolerance: float = TF_TOLERANCE cloud_accum_s: float = CLOUD_ACCUM_S - speed_max: float = SPEED_MAX color_delay: float = 0.0 # s - color timestamps lag the pose stream by this - scene_gate: bool = True embed_hz: float = EMBED_HZ - mobile: bool = False # camera rides a mobile base: room-scale policies - # (ts, merged cloud, the rolling map's pitch; None for scan sources) - _cloud_memo: tuple[float, PointCloud2, float | None] | None = field( - default=None, repr=False, init=False + mobile: bool = False # camera rides a mobile base: room-scale defaults + plane_cache: dict[tuple[int, int], SupportPlane] = field(default_factory=dict, repr=False) + # ts -> (merged cloud, the rolling map's pitch; None for scan sources) + _clouds: OrderedDict[float, tuple[PointCloud2, float | None]] = field( + default_factory=OrderedDict, repr=False, init=False ) - # (ts, plane, per-column floor table of the frame's merged cloud) - _shell_memo: tuple[float, Any, tuple[np.ndarray, np.ndarray, float, np.ndarray]] | None = field( - default=None, repr=False, init=False + # (ts, plane) -> per-column floor table of that frame's merged cloud + _shells: OrderedDict[tuple[float, int], tuple[np.ndarray, np.ndarray, float, np.ndarray]] = ( + field(default_factory=OrderedDict, repr=False, init=False) ) - _scan_cache: OrderedDict[float, np.ndarray | None] = field( + _scans: OrderedDict[float, np.ndarray | None] = field( default_factory=OrderedDict, repr=False, init=False ) - _plane_cache: dict[tuple[int, int], SupportPlane] = field( - default_factory=dict, repr=False, init=False - ) @classmethod def from_store( @@ -434,24 +409,13 @@ def from_store( ) -> Rig: """Recognize the store's shape without depending on stream names. - Resolution order per role: ``overrides`` (explicit stream names from - a caller or CLI), then ``manifest`` (passed in, or the ``.rig.json`` - sidecar next to the recording), then discovery by stream data type - and content: TFMessage-typed stream as tf, Image streams classified - by their frames (uint16/float is depth, distinct-channel uint8 is - color; a lossy-coded gray stream is neither), a PointCloud2 stream as - geometry when there is no metric depth, and - when tf is absent - the - highest-rate pose-stamped stream as pose source. On tf rigs the world - frame is the tf tree's root; ambiguity or missing calibration raises - with the candidates rather than guessing. - - The manifest carries roles as stream names plus, for recordings whose - calibration was never recorded, an inline ``camera_info`` dict and a - ``base_to_optical`` mount. + Per role: ``overrides``, then ``manifest`` (passed in, or the + ``.rig.json`` sidecar), then discovery. Cameras key on the optical + frame they calibrate and their colour and depth streams are the Image + streams stamped in that frame, so resolution is frame-driven rather than a + cascade of name tie-breaks. Ambiguity the frames cannot settle, and + missing calibration, raise with the candidates. """ - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo as CameraInfoMsg from dimos.msgs.sensor_msgs.Image import Image as ImageMsg from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as PointCloudMsg from dimos.msgs.tf2_msgs.TFMessage import TFMessage @@ -461,176 +425,124 @@ def from_store( if manifest is None: path = getattr(store.config, "path", None) - if path: - sidecar = Path(f"{path}.rig.json") - if sidecar.exists(): - manifest = json.loads(sidecar.read_text()) - logger.info(f"rig: manifest {sidecar}") - roles: dict[str, Any] = dict(manifest or {}) - roles.update(overrides or {}) - - claimed = {value for value in roles.values() if isinstance(value, str)} - color_name: str | None = roles.get("color") - depth_name: str | None = roles.get("depth") - cloud_name: str | None = roles.get("cloud") - poses_name: str | None = roles.get("poses") + sidecar = Path(f"{path}.rig.json") if path else None + if sidecar is not None and sidecar.exists(): + manifest = json.loads(sidecar.read_text()) + logger.info(f"rig: manifest {sidecar}") + roles: dict[str, Any] = {**(manifest or {}), **(overrides or {})} + claimed = {v for v in roles.values() if isinstance(v, str)} tf_names = [n for n in names if types[n] is TFMessage] tf_name = tf_names[0] if len(tf_names) == 1 else ("tf" if "tf" in tf_names else None) tf = StreamTF.from_store(store, tf_name) if tf_name is not None else None - # image streams classify by content: metric depth or genuine color + cameras = _cameras(store, roles, names, types) image_names = [n for n in names if types[n] is ImageMsg and n not in claimed] - color_candidates: list[str] = [] - depth_candidates: list[str] = [] - empty_images: list[str] = [] - image_frames: dict[str, str] = {} - for name in image_names: - stream = store.stream(name) - if stream.count() == 0: - empty_images.append(name) - continue - image = stream.first().data - image_frames[name] = image.frame_id - frame = image.data - if frame.dtype == np.uint16 or frame.dtype.kind == "f": - depth_candidates.append(name) - elif ( - frame.ndim == 3 - and frame.shape[2] == 3 - and not np.array_equal(frame[..., 0], frame[..., 1]) - ): - color_candidates.append(name) - else: - logger.info(f"rig: image stream {name!r} is neither color nor metric depth") - if color_name is None and depth_name is not None: - depth_frame = store.stream(depth_name).first().data.frame_id - matching = [name for name in color_candidates if image_frames[name] == depth_frame] - if len(matching) == 1: - color_name = matching[0] - logger.info(f"rig: paired depth {depth_name!r} with color {color_name!r}") + by_frame_color, depth_streams = _images(store, image_names) + + color_name = roles.get("color") if color_name is None: - if len(color_candidates) > 1: - color_name = max(color_candidates, key=lambda n: _stream_rate(store.stream(n))) - logger.info(f"rig: several color streams, using highest-rate {color_name!r}") - elif color_candidates: - color_name = color_candidates[0] - elif len(empty_images) == 1: - color_name = empty_images[0] # a live store's not-yet-filled feed + # a camera's own frame first, then any colour feed, then a live store's + # not-yet-filled one + ranked = [by_frame_color[f] for f in cameras if f in by_frame_color] + ranked += [n for n in by_frame_color.values() if n not in ranked] + ranked += [n for n in image_names if store.stream(n).count() == 0] + color_name = next(iter(ranked), None) if color_name is None: raise ValueError(f"no color image stream among {names}; pass a manifest or --color") claimed.add(color_name) - if depth_name is None: - if len(depth_candidates) > 1: - raise ValueError(f"several depth streams {depth_candidates}; pass --depth") - depth_name = depth_candidates[0] if depth_candidates else None + color = store.stream(color_name) + color_frame = color.first().data.frame_id if color.count() else None + depth_name = roles.get("depth") + if depth_name is None: + # aligned depth shares the colour frame; otherwise the only depth stream + # stands, and a choice between several is the caller's + aligned = sorted(n for f, n in depth_streams if f == color_frame) + elsewhere = [n for f, n in depth_streams if f != color_frame] + if aligned: + depth_name = aligned[0] + elif len(elsewhere) > 1: + raise ValueError(f"several depth streams {elsewhere}; pass --depth") + elif elsewhere: + depth_name = elsewhere[0] + + cloud_name = roles.get("cloud") if cloud_name is None and depth_name is None: cloud_names = [n for n in names if types[n] is PointCloudMsg and n not in claimed] if len(cloud_names) > 1: raise ValueError(f"several pointcloud streams {cloud_names}; pass --cloud") cloud_name = cloud_names[0] if cloud_names else None - color = store.stream(color_name) depth = store.stream(depth_name) if depth_name is not None else None cloud = store.stream(cloud_name) if cloud_name is not None and depth is None else None if cloud_name is not None: claimed.add(cloud_name) - # intrinsics: inline manifest dict, named stream, or discovery by - # type with the color camera's frame deciding among several - camera_info = None - camera_info_role = roles.get("camera_info") - if isinstance(camera_info_role, dict): - camera_info = CameraInfoMsg( - height=camera_info_role["height"], - width=camera_info_role["width"], - distortion_model=camera_info_role.get("distortion_model", ""), - D=camera_info_role.get("D"), - K=camera_info_role["K"], - R=camera_info_role.get("R"), - P=camera_info_role.get("P"), - frame_id=camera_info_role["frame_id"], - ) + if not cameras: + # embed-only stores may carry no calibration; geometry raises on use + cameras = {OPTICAL_FRAME: cast("CameraInfo", None)} + if color_frame in cameras: + optical = color_frame + elif len(cameras) == 1: + optical = next(iter(cameras)) # images stamped in a frame calibration never names else: - ci_name = camera_info_role if isinstance(camera_info_role, str) else None - if ci_name is None: - candidates = [ - n for n in names if types[n] is CameraInfoMsg and store.stream(n).count() - ] - try: - color_frame = color.first().data.frame_id - except LookupError: - color_frame = None - matching = [ - n for n in candidates if store.stream(n).first().data.frame_id == color_frame - ] - if matching: - ci_name = sorted(matching)[0] - elif len(candidates) == 1: - ci_name = candidates[0] - elif len(candidates) > 1: - raise ValueError(f"several camera_info streams {candidates}; pass a manifest") - if ci_name is not None: - camera_info = store.stream(ci_name).first().data - - # embed-only stores (no geometry) may carry no calibration at all; - # every geometry API raises on use, embedding never touches it - optical_frame = camera_info.frame_id if camera_info is not None else gates.OPTICAL_FRAME - - world_frame = gates.WORLD_FRAME + raise ValueError( + f"color stream {color_name!r} is stamped {color_frame!r}, which is none of the " + f"cameras {sorted(cameras)}; name the right camera_info in the manifest" + ) + cameras = {optical: cameras[optical], **cameras} # the colour camera answers first + + world_frame = WORLD_FRAME if tf is not None: - world_frame = _tf_root(store, cast("str", tf_name), tf, optical_frame) or world_frame + world_frame = _tf_root(store, cast("str", tf_name), tf, optical) or world_frame elif cloud is not None: try: world_frame = cloud.first().data.frame_id except LookupError: pass # live store, nothing recorded yet - base_to_optical = None - mount = roles.get("base_to_optical") - if isinstance(mount, dict): - base_to_optical = Transform( + mounts: dict[str, Transform] = {} + if isinstance(mount := roles.get("base_to_optical"), dict): + mounts[optical] = Transform( translation=Vector3(*mount["translation"]), rotation=Quaternion(*mount["rotation"]), - frame_id="base_link", - child_frame_id=camera_info.frame_id if camera_info else "camera_optical", + frame_id=mount.get("frame_id", "base_link"), + child_frame_id=optical, ) poses = None if tf is None: + poses_name = roles.get("poses") if poses_name is None: posed = [ n for n in names if n not in claimed - and n != color_name and store.stream(n).count() and store.stream(n).first().pose_tuple is not None ] - if posed: - poses_name = max(posed, key=lambda n: _stream_rate(store.stream(n))) + poses_name = max(posed, key=lambda n: _rate(store.stream(n))) if posed else None poses = store.stream(poses_name) if poses_name is not None else None if depth is not None or cloud is not None: - if camera_info is None: + if cameras[optical] is None: raise ValueError( - "store has 3D geometry but no camera calibration; add a CameraInfo " - "stream role or an inline camera_info to the .rig.json manifest" + "store has 3D geometry but no camera calibration; add a CameraInfo stream " + "role or an inline camera_info to the .rig.json manifest" ) - if tf is None and (poses is None or base_to_optical is None): + if tf is None and (poses is None or not mounts): raise ValueError( "store has no tf; a pose-stamped rig needs a poses stream and a " "base_to_optical mount in the .rig.json manifest" ) rig = cls( - camera_info=cast("CameraInfo", camera_info), + cameras=cameras, color=color, world_frame=world_frame, - optical_frame=optical_frame, tf=tf, - base_to_optical=base_to_optical, + mounts=mounts, poses=poses, depth=depth, cloud=cloud, @@ -639,37 +551,59 @@ def from_store( span = _camera_span(rig) rig.mobile = span > MOBILE_SPAN_M if rig.mobile: - rig.speed_max = WALK_SPEED_MAX - rig.scene_gate = False rig.embed_hz = WALK_EMBED_HZ - if camera_info is not None: + if cameras[optical] is not None: rig.color_delay = estimate_color_delay(rig) logger.info( - f"rig: color={color_name!r} depth={depth_name!r} cloud={cloud_name!r} " - f"tf={tf_name!r} world={world_frame!r} span={span:.1f}m mobile={rig.mobile} " + f"rig: color={color_name!r} depth={depth_name!r} cloud={cloud_name!r} tf={tf_name!r} " + f"world={world_frame!r} span={span:.1f}m mobile={rig.mobile} " f"color_delay={rig.color_delay * 1000:.0f}ms" ) return rig + @property + def optical_frame(self) -> str: + """The default camera: the one the rig's ``color`` stream feeds. + + Discovery registers it first; a hand-built rig orders its own dict. + """ + return next(iter(self.cameras)) + + def camera_frame(self, frame: str | None) -> str: + """Which camera resolves a lookup: the image's own, if the rig has it. + + A recording can stamp its images in a frame its calibration never + names, so on a one-camera rig any frame falls to that camera. With + several, nothing says which one took the image, and folding would + silently project through the wrong model, so it is an error. + """ + if frame is None or frame in self.cameras: + return frame or self.optical_frame + if len(self.cameras) > 1: + raise KeyError( + f"image frame {frame!r} is none of the rig's cameras {sorted(self.cameras)}" + ) + return self.optical_frame + # pose - def world_to_optical(self, ts: float) -> Transform | None: + def world_to_optical(self, ts: float, frame: str | None = None) -> Transform | None: + frame = self.camera_frame(frame) ts -= self.color_delay # color stamps lag; the capture instant is earlier if self.tf is not None: - return self.tf.get(self.optical_frame, self.world_frame, ts, self.tf_tolerance) + return self.tf.get(frame, self.world_frame, ts, self.tf_tolerance) pose = self.pose_at(ts) if pose is None: return None - mount = cast("Transform", self.base_to_optical) - return -(Transform.from_pose("base_link", pose) + mount) + mount = self.mounts[frame] + return -(Transform.from_pose(mount.frame_id, pose) + mount) def pose_at(self, ts: float) -> PoseStamped | None: - """World base pose at ts, interpolated between the bracketing samples. + """World base pose at ts, interpolated between bracketing samples. - A walking robot covers centimeters and degrees per pose period, so - snapping to a recorded sample misplaces the projection; interpolation - follows the motion. Outside the bracketed span the nearest sample in - the tolerance window stands. + A walking robot covers centimeters per pose period, so snapping to a + sample misplaces the projection. Outside the bracketed span the + nearest sample in the tolerance window stands. """ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -697,58 +631,33 @@ def pose_at(self, ts: float) -> PoseStamped | None: orientation=(float(q[0]), float(q[1]), float(q[2]), float(q[3])), ) - def camera_pose(self, ts: float) -> PoseStamped | None: - """World pose of the camera optical frame at ts.""" - transform = self.world_to_optical(ts) + def camera_pose(self, ts: float, frame: str | None = None) -> PoseStamped | None: + """World pose of a camera's optical frame at ts.""" + transform = self.world_to_optical(ts, frame) return (-transform).to_pose() if transform is not None else None def index_pose(self, obs: Observation[Any]) -> Pose | PoseStamped | None: - """The pose an embedded index observation carries. - - tf rigs stamp the derived optical pose; pose-stamped rigs keep the - recorded base pose, which is what ``sees`` expects to find on an - observation when it resolves the transform through the mount. - """ + """tf rigs stamp the optical pose; pose-stamped rigs keep the base one.""" if self.tf is not None: - return self.camera_pose(obs.ts) + return self.camera_pose(obs.ts, obs.data.frame_id) return obs.pose - def camera_speed(self, ts: float, dt: float = 0.06) -> float | None: - """Linear camera speed (m/s) around ts, from pose differencing.""" - a = self.camera_pose(ts - dt) - b = self.camera_pose(ts + dt) - if a is None or b is None: - return None - return float((b.position - a.position).magnitude() / (2 * dt)) - - def camera_still(self, ts: float, envelope: float = STILL_ENVELOPE) -> bool: - """Camera below the rig's speed gate over the whole capture envelope.""" - for offset in (-envelope, 0.0, envelope): - speed = self.camera_speed(ts + offset) - if speed is None or speed > self.speed_max: - return False - return True - - def still_intervals(self, t0: float, t1: float) -> list[tuple[float, float]]: - """Maximal camera-still intervals inside [t0, t1], sampled at 0.25 s.""" - step = 0.25 - times = np.arange(t0, t1 + step, step) - intervals: list[tuple[float, float]] = [] - run_start: float | None = None - for t in times: - speed = self.camera_speed(float(t)) - still = speed is not None and speed <= self.speed_max - if still and run_start is None: - run_start = float(t) - elif not still and run_start is not None: - intervals.append((run_start, float(t) - step)) - run_start = None - if run_start is not None: - intervals.append((run_start, float(times[-1]))) - return [(a, b) for a, b in intervals if b >= a] - # geometry + def frame_at(self, obs: Observation[Any]) -> Image: + """The colour frame an index observation refers to. + + A live index carries the image itself; a replay index carries only what + frame selection needs, so its pixels come back from the colour stream, + which must still hold the frame the index was built from. + """ + from dimos.msgs.sensor_msgs.Image import Image + + if isinstance(obs.data, Image): + return obs.data + image: Image = self.color.at(obs.ts, FRAME_TOLERANCE).first().data + return image + def depth_at(self, ts: float) -> Image | None: """Temporal join: aligned depth frame for a color timestamp.""" try: @@ -758,16 +667,11 @@ def depth_at(self, ts: float) -> Image | None: return depth def registered_scan(self, scan: Observation[PointCloud2]) -> np.ndarray | None: - """A scan's points in the world frame, registered via tf when needed. - - Decode and registration run once per scan per run: accumulation - windows of neighboring frames overlap almost entirely, and every - query of a multi-label run shares every window. - """ + """A scan's world-frame points, registered via tf, decoded once per run.""" key = scan.ts - if key in self._scan_cache: - self._scan_cache.move_to_end(key) - return self._scan_cache[key] + if key in self._scans: + self._scans.move_to_end(key) + return self._scans[key] points: np.ndarray | None = scan.data.as_numpy()[0] frame = scan.data.frame_id if frame != self.world_frame: @@ -781,28 +685,28 @@ def registered_scan(self, scan: Observation[PointCloud2]) -> np.ndarray | None: else: matrix = transform.to_matrix() points = points @ matrix[:3, :3].T + matrix[:3, 3] - self._scan_cache[key] = points - if len(self._scan_cache) > _SCAN_CACHE_MAX: - self._scan_cache.popitem(last=False) + self._scans[key] = points + if len(self._scans) > _SCAN_CACHE_MAX: + self._scans.popitem(last=False) return points def cloud_at(self, ts: float) -> PointCloud2 | None: """World-frame geometry at ts: the window's scans merged. - The merge rule follows the source's shape, measured per window from - the points themselves. A stream whose next snapshot re-reports the - majority of the nearest one's exact points is a rolling occupancy - map: each snapshot is already temporally integrated over the area - it covers and the map clears what moved away, so snapshots merge - nearest-ts first, a farther snapshot contributing only cells - outside the XY coverage of every nearer one - plain accumulation - would resurrect every moved object's trail. Scans that never repeat - are fresh samples of a scene static in world frame, so they - accumulate whole; coordinate quantization alone proves nothing, a - mm-integer wire format grids a scan without making it a map. + A stream whose next snapshot re-reports the majority of the nearest + one's exact points is a rolling occupancy map, already integrated over + the area it covers and cleared of what moved away: those merge + nearest-ts first, a farther snapshot contributing only cells outside + every nearer one's XY coverage, since plain accumulation would + resurrect each moved object's trail. Scans that never repeat are fresh + samples of a static scene and accumulate whole; quantization alone + proves nothing, a mm-integer wire format grids a scan without making + it a map. """ - if self._cloud_memo is not None and self._cloud_memo[0] == ts: - return self._cloud_memo[1] + held = self._clouds.get(ts) + if held is not None: + self._clouds.move_to_end(ts) + return held[0] from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 scans = self.cloud.after(ts - self.cloud_accum_s).before(ts + self.cloud_accum_s) @@ -845,27 +749,35 @@ def cloud_at(self, ts: float) -> PointCloud2 | None: ] = True points = np.vstack(kept) merged = PointCloud2.from_numpy(points, frame_id=self.world_frame, timestamp=ts) - self._cloud_memo = (ts, merged, quantum) + self._clouds[ts] = (merged, quantum) + if len(self._clouds) > _CLOUD_CACHE_MAX: + self._clouds.popitem(last=False) return merged + def _quantum(self, ts: float) -> float | None: + """The lattice pitch of the merged cloud at ts; None for scan sources.""" + self.cloud_at(ts) + return self._clouds[ts][1] + def _shell_table( self, ts: float, plane: SupportPlane ) -> tuple[np.ndarray, np.ndarray, float, np.ndarray] | None: - """Per-column floor of the frame's merged cloud; None for continuous sources. + """Per-column floor of the merged cloud; None for continuous sources. - A rolling map registers each snapshot with its own odometry error, - so the support surface sits at a different absolute level per - region. The floor of a point's own XY column is the local reference - a single global plane cannot be. + A rolling map registers each snapshot with its own odometry error, so + the support sits at a different absolute level per region; a point's + own XY column is the local reference one global plane cannot be. """ cloud = self.cloud_at(ts) - if cloud is None or self._cloud_memo[2] is None: # type: ignore[index] + quantum = self._clouds[ts][1] if cloud is not None else None + if quantum is None: return None - memo = self._shell_memo - if memo is not None and memo[0] == ts and memo[1] is plane: - return memo[2] - quantum = cast("float", self._cloud_memo[2]) # type: ignore[index] - points = cloud.as_numpy()[0] + key = (ts, id(plane)) + held = self._shells.get(key) + if held is not None: + self._shells.move_to_end(key) + return held + points = cloud.as_numpy()[0] # type: ignore[union-attr] anchor = points[0, :2] keys = _column_keys(points, quantum, anchor) heights = plane.height_above(points) @@ -873,17 +785,17 @@ def _shell_table( keys_sorted = keys[order] starts = np.nonzero(np.concatenate(([True], np.diff(keys_sorted) != 0)))[0] table = (keys_sorted[starts], np.minimum.reduceat(heights[order], starts), quantum, anchor) - self._shell_memo = (ts, plane, table) + self._shells[key] = table + if len(self._shells) > _CLOUD_CACHE_MAX: + self._shells.popitem(last=False) return table def support_heights(self, ts: float, plane: SupportPlane, points: np.ndarray) -> np.ndarray: - """Signed heights of points above the frame's support shell. + """Signed heights above the frame's support shell. - Depth rigs and continuous-scan sources measure against the fitted - plane directly. A rolling-map source measures against the local - column floor of the frame's own merged cloud (see ``_shell_table``); - ``points`` must come from that cloud, which is what every projected - lift produces. + Depth and continuous-scan sources measure against the fitted plane; a + rolling map measures against its own column floor, so ``points`` must + come from that frame's merged cloud, which every projected lift is. """ heights = plane.height_above(points) if self.cloud is None: @@ -897,16 +809,13 @@ def support_heights(self, ts: float, plane: SupportPlane, points: np.ndarray) -> return local def _support_strip(self, ts: float, plane: SupportPlane, shell: float, gap: float = 0.3) -> Any: - """Filter dropping support-shell points outside the object's own stance. - - A misaligned mask row collects the support surface along the whole - view ray; shell points survive only within the camera-range span of - the detection's above-shell structure - under the object, not along - the approach. The stance is the dominant range cluster of the - above-shell points (the same ``gap`` split as ``range_cluster``), so - background caught on the mask rim cannot widen it. Points below the - shell (mirror returns through a glossy surface) never survive. A - detection with no structure above the shell passes untouched. + """Drop support-shell points outside the object's own stance. + + A misaligned mask row collects the support along the whole view ray; + shell points survive only within the camera-range span of the + detection's above-shell structure - under the object, not along the + approach. Points below the shell never survive; a detection with no + structure above it passes untouched. """ def filter_func(det: Any, pc: Any, ci: Any, tf: Any) -> Any: @@ -936,32 +845,32 @@ def filter_func(det: Any, pc: Any, ci: Any, tf: Any) -> Any: def lift( self, detections: ImageDetections2D, plane: SupportPlane | None = None ) -> ImageDetections3DPC | None: - """2D detections to world-frame 3D clouds, or None without geometry/pose. - - With a support ``plane``, a projected lift strips support-shell - points outside each detection's stance before the generic filters. - The shell is the support surface's own occupied band: a plane - crossing a lattice straddles at most two adjacent levels, so a - rolling map's shell ends halfway to the third; a continuous source's - shell is the plane fit's inlier distance. + """2D detections to world-frame clouds, or None without geometry/pose. + + The shell is the support's own occupied band: a plane crossing a + lattice straddles at most two levels, so a rolling map's shell ends + halfway to the third, and a continuous source's is the fit's inlier + distance. """ - transform = self.world_to_optical(detections.ts) + frame = self.camera_frame(detections.image.frame_id) + transform = self.world_to_optical(detections.ts, frame) if transform is None: return None + camera_info = self.cameras[frame] if self.depth is not None: depth = self.depth_at(detections.ts) if depth is None: return None - return ImageDetections3DPC.from_depth(detections, depth, self.camera_info, transform) + return ImageDetections3DPC.from_depth(detections, depth, camera_info, transform) cloud = self.cloud_at(detections.ts) if cloud is None: return None filters = _CLOUD_LIFT_FILTERS if plane is not None: - quantum = cast("float | None", self._cloud_memo[2]) # type: ignore[index] + quantum = self._clouds[detections.ts][1] shell = 1.5 * quantum if quantum is not None else PLANE_DISTANCE_CLOUD filters = [self._support_strip(detections.ts, plane, shell), *_CLOUD_LIFT_FILTERS] - return ImageDetections3DPC.from_2d(detections, cloud, self.camera_info, transform, filters) + return ImageDetections3DPC.from_2d(detections, cloud, camera_info, transform, filters) def backdrop(self, ts: float, depth_trunc: float = 1.5) -> PointCloud2 | None: """World-frame scene cloud around ts, for plane fits and rendering.""" @@ -970,97 +879,22 @@ def backdrop(self, ts: float, depth_trunc: float = 1.5) -> PointCloud2 | None: from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 depth = self.depth_at(ts) - transform = self.world_to_optical(ts) - if depth is None or transform is None: + if depth is None: return None try: color = self.color.at(ts, 0.1).first().data except LookupError: return None + frame = self.camera_frame(color.frame_id) + transform = self.world_to_optical(ts, frame) + if transform is None: + return None return PointCloud2.from_rgbd( - color, depth, self.camera_info, depth_scale=0.001, depth_trunc=depth_trunc + color, depth, self.cameras[frame], depth_scale=0.001, depth_trunc=depth_trunc ).transform(-transform) - # predicates - - def sees( - self, - point: Any, - *, - extent: Any | None = None, - min_fraction: float = 1.0, - max_range: float | None = None, - ) -> Callable[[Observation[Any]], bool]: - """Predicate: does an observation's camera see the world point. - - Occlusion checking through measured depth only exists on depth rigs; - projected-cloud rigs rely on the geometric visibility test alone. - """ - return project_sees( - point, - self.camera_info, - tf=self.tf, - base_to_optical=self.base_to_optical, - world_frame=self.world_frame, - optical_frame=self.optical_frame, - time_tolerance=self.tf_tolerance, - extent=extent, - min_fraction=min_fraction, - max_range=max_range, - depth=(lambda obs: self.depth_at(obs.ts)) if self.depth is not None else None, - ) - - def keyframes( - self, - t0: float, - t1: float, - stride: float, - motion_threshold: float = gates.MOTION_THRESHOLD, - ) -> list[Observation[Image]]: - """Camera-still (and, on scene-gated rigs, scene-still) frames on a grid. - - For each grid point the nearest passing frame within half a stride is - selected, so a grid point landing mid-sweep snaps to the neighboring - pause instead of being lost. - """ - intervals = self.still_intervals(t0, t1) if self.scene_gate else [] - gray: dict[float, np.ndarray | None] = {} - selected: list[Observation[Image]] = [] - seen: set[float] = set() - offsets = [0.0] - probe = 0.35 - while probe <= stride / 2: - offsets.extend([probe, -probe]) - probe += 0.35 - - t = t0 + 0.5 - while t < t1: - for offset in offsets: - ts = t + offset - if ts < t0 or ts > t1: - continue - if not self.camera_still(ts): - continue - if self.scene_gate and not gates.scene_still( - self.color, ts, intervals, gray, motion_threshold - ): - continue - try: - obs = self.color.at(ts, 0.1).first() - except LookupError: - continue - if obs.ts in seen: - break - if self.world_to_optical(obs.ts) is None: - continue - seen.add(obs.ts) - selected.append(obs) - break - t += stride - return selected - def default_localize_policy(self) -> LocalizePolicy: - return ROOM_LOCALIZE_POLICY if self.mobile else LocalizePolicy() - - def default_inventory_policy(self) -> InventoryPolicy: - return ROOM_INVENTORY_POLICY if self.mobile else InventoryPolicy() + base = ROOM_LOCALIZE_POLICY if self.mobile else LocalizePolicy() + if self.cloud is None: + return base + return replace(base, min_points=CLOUD_MIN_POINTS) diff --git a/dimos/perception/memory/support_plane.py b/dimos/perception/memory/support_plane.py index 52a4b0f921..3ab4ad1ad5 100644 --- a/dimos/perception/memory/support_plane.py +++ b/dimos/perception/memory/support_plane.py @@ -12,15 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Support-surface fit and the scope predicate derived from it. - -The plane is RANSAC-fit from the window's own frames - nothing scene-specific -is passed in and no caller supplies coordinates. On a wrist-camera rig over a -workspace the dominant horizontal plane is the tabletop; on a mobile rig it -is the floor. The plane's inlier footprint is the workspace; the scope -predicate accepts a support when its cloud sits in a band above the plane and -its footprint intersects the plane footprint. -""" +"""RANSAC fit of the dominant horizontal support surface from window keyframes.""" from __future__ import annotations @@ -41,58 +33,33 @@ PLANE_DISTANCE = 0.01 # m - RANSAC inlier distance for depth-camera clouds PLANE_DISTANCE_CLOUD = 0.03 # m - registered lidar is noisier MIN_HORIZONTAL_DOT = 0.90 # |normal . z| for a plane to count as horizontal -PLANE_SEED = 0 # RANSAC is seeded so one window always fits the same plane -FOOTPRINT_DILATE_M = 0.03 +PLANE_SEED = 0 # so one window always fits the same plane @dataclass class SupportPlane: - """A horizontal support surface: plane coefficients plus inlier footprint.""" + """Plane coefficients (a, b, c, d), normal up, plus the inlier footprint.""" - # Plane (a, b, c, d): a*x + b*y + c*z + d = 0, normal pointing up (+z). coefficients: tuple[float, float, float, float] footprint_hull: np.ndarray # (K, 2) convex hull of inlier (x, y), world inlier_count: int - @property - def normal(self) -> np.ndarray: - return np.array(self.coefficients[:3]) - def height_above(self, points: np.ndarray) -> np.ndarray: - """Signed height of (N, 3) world points above the plane.""" a, b, c, d = self.coefficients heights: np.ndarray = points @ np.array([a, b, c]) + d return heights - def footprint_contains(self, points_xy: np.ndarray) -> np.ndarray: - """Boolean mask: which (N, 2) world XY points fall inside the (dilated) hull.""" - from matplotlib.path import Path as MplPath - - hull = self.footprint_hull - center = hull.mean(axis=0) - offsets = hull - center - norms = np.linalg.norm(offsets, axis=1, keepdims=True) - dilated = hull + offsets / np.maximum(norms, 1e-9) * FOOTPRINT_DILATE_M - return MplPath(dilated).contains_points(points_xy) - def fit_support_plane(rig: Rig, keyframes: list[Observation[Image]]) -> SupportPlane | None: - """Fit the dominant near-horizontal plane from a handful of window keyframes. - - Non-horizontal dominant planes (a wall, a screen) are peeled off and the - fit repeats on the remainder. Among horizontal candidates the one with - the most inliers wins. - """ + """Dominant near-horizontal plane; non-horizontal candidates are peeled off.""" if not keyframes: return None - picks = keyframes[:: max(1, len(keyframes) // 5)][:5] clouds = [] - for obs in picks: + for obs in keyframes[:: max(1, len(keyframes) // 5)][:5]: cloud = rig.backdrop(obs.ts) - if cloud is None: - continue - clouds.append(cloud.voxel_downsample(0.01)) + if cloud is not None: + clouds.append(cloud.voxel_downsample(0.01)) if not clouds: return None @@ -105,13 +72,11 @@ def fit_support_plane(rig: Rig, keyframes: list[Observation[Image]]) -> SupportP import open3d as o3d - # Unseeded, the plane fit lands on a different set of inliers each run o3d.utility.random.seed(PLANE_SEED) - distance = PLANE_DISTANCE if rig.depth is not None else PLANE_DISTANCE_CLOUD remaining = o3d.geometry.PointCloud() remaining.points = o3d.utility.Vector3dVector(points) - best: tuple[np.ndarray, np.ndarray] | None = None # (coefficients, inlier points) + best: tuple[np.ndarray, np.ndarray] | None = None for _ in range(4): if len(remaining.points) < 500: break @@ -119,10 +84,8 @@ def fit_support_plane(rig: Rig, keyframes: list[Observation[Image]]) -> SupportP distance_threshold=distance, ransac_n=3, num_iterations=1000, probability=1.0 ) inliers = np.asarray(remaining.points)[inlier_idx] - normal = np.array(model[:3]) - if abs(normal[2]) >= MIN_HORIZONTAL_DOT: - if best is None or len(inliers) > len(best[1]): - best = (np.array(model), inliers) + if abs(model[2]) >= MIN_HORIZONTAL_DOT and (best is None or len(inliers) > len(best[1])): + best = (np.array(model), inliers) remaining = remaining.select_by_index(inlier_idx, invert=True) if best is None: @@ -130,14 +93,13 @@ def fit_support_plane(rig: Rig, keyframes: list[Observation[Image]]) -> SupportP return None model, inliers = best - if model[2] < 0: # normal points up + if model[2] < 0: model = -model from scipy.spatial import ConvexHull xy = inliers[:, :2] - hull = ConvexHull(xy) return SupportPlane( coefficients=(float(model[0]), float(model[1]), float(model[2]), float(model[3])), - footprint_hull=xy[hull.vertices], + footprint_hull=xy[ConvexHull(xy).vertices], inlier_count=len(inliers), ) diff --git a/dimos/perception/memory/tool_localize.py b/dimos/perception/memory/tool_localize.py index 093cab0725..1591dd23df 100644 --- a/dimos/perception/memory/tool_localize.py +++ b/dimos/perception/memory/tool_localize.py @@ -121,7 +121,7 @@ def image_colors(det: Any) -> np.ndarray: points = det.pointcloud.points_f32() matrix = det.transform.to_matrix() cam = points @ matrix[:3, :3].T + matrix[:3, 3] - pixels = Detection3DPC.project_pixels(cam, rig.camera_info) + pixels = Detection3DPC.project_pixels(cam, rig.cameras[rig.optical_frame]) cols = np.round(pixels[:, 0]).astype(int) rows = np.round(pixels[:, 1]).astype(int) rgb = det.image.to_rgb().data @@ -145,7 +145,7 @@ def image_colors(det: Any) -> np.ndarray: elif rig.depth is not None: backdrop_ts = next((t.backdrop_ts for _, t in traces if t.backdrop_ts is not None), None) if backdrop_ts is None: - backdrop_ts = next((t.matched[0][0] for _, t in traces if t.matched), None) + backdrop_ts = next((t.first_match_ts for _, t in traces if t.first_match_ts), None) if backdrop_ts is not None: backdrop = rig.backdrop(backdrop_ts) if backdrop is not None: @@ -165,7 +165,7 @@ def image_colors(det: Any) -> np.ndarray: # live camera feed + frustum tracking the capture pose along the # timeline; on mobile rigs a translucent box marks the robot - rr.log("camera", rig.camera_info.to_rerun(), static=True) + rr.log("camera", rig.cameras[rig.optical_frame].to_rerun(), static=True) if rig.mobile: rr.log( "robot", @@ -207,7 +207,7 @@ def image_colors(det: Any) -> np.ndarray: continue frame = f"{root}/frames/{i}" rr.log(frame, pose.to_rerun()) - rr.log(frame, rig.camera_info.to_rerun()) + rr.log(frame, rig.cameras[rig.optical_frame].to_rerun()) rr.log(f"{frame}/image", annotated.to_rerun()) # the answers: per verified instance, every sighting's cloud colored diff --git a/dimos/perception/memory/types.py b/dimos/perception/memory/types.py index 0d7e30b8bb..eca151e881 100644 --- a/dimos/perception/memory/types.py +++ b/dimos/perception/memory/types.py @@ -12,20 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Object registration types: supports, instances, localizations, policies. - -The identity key throughout is the *support* - the object's occupied volume -in world coordinates. Labels and appearance are metadata attached to a -support, never a key. -""" +"""The localize answer types and the thresholds a caller tunes per call.""" from __future__ import annotations -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Literal - -import numpy as np +from dataclasses import dataclass +from typing import TYPE_CHECKING if TYPE_CHECKING: from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @@ -44,45 +36,9 @@ class Support: frame_id: str -@dataclass -class SupportObservation: - """One accepted per-frame observation of a support: a masked depth lift.""" - - ts: float - cloud: PointCloud2 - centroid: np.ndarray # (3,) world - aabb_min: np.ndarray # (3,) world - aabb_max: np.ndarray # (3,) world - n_points: int - mask_area_px: int - camera_position: np.ndarray # (3,) world - score: float = 1.0 - bbox: tuple[float, float, float, float] | None = None - - -@dataclass -class Instance: - """A deduplicated object instance computed for one inventory call.""" - - instance_id: str - grounded: bool - primary_label: str | None - labels: tuple[tuple[str, float], ...] - state: Literal["active", "occluded", "stale", "retired"] - identity_confidence: float - support: Support | None - latest_position_xyz: tuple[float, float, float] | None - latest_seen_ts: float - members: list[SupportObservation] = field(default_factory=list) - - @dataclass class Localization: - """One verified instance of a queried object. - - ``point_cloud`` is the union of every viewpoint that saw the instance; - position and timestamps follow the latest sighting. - """ + """One verified instance. Cloud is every viewpoint's union; pose is the latest.""" instance_id: str semantic_score: float @@ -106,171 +62,29 @@ class Localization: @dataclass(frozen=True) class LocalizePolicy: - """Score and geometry thresholds for candidate formation, lift and acceptance. - - The funnel generalizes; these numbers do not. Each was fit to the score - and height distributions of one measured scene, so a different rig, - object scale or detector vocabulary needs its own instance rather than - the defaults. - - `candidate_floor`: OWLv2 score at which a box is formed for a query. - Boxes below this never lift. - - `accept_score`: A support group is returned when its highest member score meets this. - The RGB-only path uses the same floor. - - `min_views`: Unique camera positions, rounded to 1 cm, required before a group is confirmed. - A support seen from one pose only is dropped. - - `cluster_radius_m`: Two lifted detections are the same support when their cloud centers sit within this many meters. - - `peak_prominence`: Minimum SigLIP similarity rise for a semantic peak. - A local maximum below this is not a peak. - - `peak_distance_s`: Minimum seconds between two semantic peaks. - - `peak_width_s`: Minimum peak width in seconds at half prominence. - ``None`` disables the width gate. - - `verify_radius_m`: Index frames whose pose is within this many meters of a peak are gathered for OWLv2. - - `verify_window_s`: Keep the sharpest gathered frame per this many seconds. - - `settled_window_fraction`: Collapse index frames closer than this fraction of ``1/embed_hz`` to the sharper one, so one window does not count as two viewpoints. + """Per-call thresholds. The funnel generalizes; these numbers do not. - `tail_k`: After the last peak, take this many extra frames by query similarity so the window tail can be a latest sighting. - - `max_object_extent_m`: Drop a lift whose longest AABB edge exceeds this. - - `surface_patch_max_rise_m`: Drop a lift whose 95th-percentile height above the support is below this, when the drop test also holds. - A cloud that hugs the support is a patch of the surface, not an object. - - `surface_patch_min_drop_m`: Drop a lift whose 5th-percentile height is above this, when the rise test also holds. - - `min_depth_points`: Minimum points on a depth lift. - Projected-cloud rigs ignore this. - - `min_camera_range_m`: Drop a lift whose median point-to-camera range is below this. - - `fuse_voxel_m`: Voxel size of the union cloud at identity merge. - ``0`` concatenates. - - `plane_cell_m`: XY cell size for the support-plane cache. - - `plane_keyframes`: Candidate frames sampled to fit the support plane. - - `refusal_margin`: If this instance's score minus the best coexisting rival is below this, ``reason`` is set. - The instance is still returned. + Each was fit to one measured scene, so a different rig, object scale or + detector vocabulary passes its own instance instead of the defaults. """ - candidate_floor: float = 0.25 - accept_score: float = 0.40 - min_views: int = 2 - cluster_radius_m: float = 0.08 + candidate_floor: float = 0.25 # OWLv2 score at which a box is formed + accept_score: float = 0.40 # group's best member score to be returned + min_views: int = 2 # distinct camera positions (1 cm) to confirm a group + cluster_radius_m: float = 0.08 # two lifts are one object within this peak_prominence: float = 0.02 peak_distance_s: float = 1.0 - peak_width_s: float | None = 0.5 - verify_radius_m: float = 1.6 - verify_window_s: float = 0.5 - settled_window_fraction: float = 0.5 - tail_k: int = 1 + peak_width_s: float | None = 0.5 # None disables the width gate + verify_radius_m: float = 1.6 # gather index frames this close to a peak + verify_window_s: float = 0.5 # keep the sharpest gathered frame per window + settled_window_fraction: float = 0.5 # of 1/embed_hz; collapses split windows + tail_k: int = 1 # extra frames past the last peak, so the tail can be latest max_object_extent_m: float = 0.60 - surface_patch_max_rise_m: float = 0.003 - surface_patch_min_drop_m: float = -0.02 - min_depth_points: int = 60 + surface_patch_max_rise_m: float = 0.003 # with the drop test, a lift hugging + surface_patch_min_drop_m: float = -0.02 # the support is surface, not object + min_points: int = 60 # points on a lift min_camera_range_m: float = 0.28 - fuse_voxel_m: float = 0.01 - plane_cell_m: float = 2.0 + fuse_voxel_m: float = 0.01 # union cloud voxel at merge; 0 concatenates + plane_cell_m: float = 2.0 # XY cell of the support-plane cache plane_keyframes: int = 5 - refusal_margin: float = 0.15 - - -@dataclass(frozen=True) -class InventoryPolicy: - """Thresholds for discovery, validity, scope, association and naming. - - Every geometric quantity is metric (meters, seconds, pixels, IoU) - a - claim that can be checked against the recording. The two naming numbers - are detector scores and decide whether a name is reported, never whether - an instance exists. The candidate names are data and travel as a call - argument, not as policy. - """ - - keyframe_stride_s: float = 2.5 # proposal keyframe grid - min_mask_area_px: int = 400 - max_mask_area_fraction: float = 0.25 - min_depth_points: int = 60 - max_object_extent_m: float = 0.45 - min_height_above_plane_m: float = 0.003 - band_above_plane_m: tuple[float, float] = (-0.02, 0.30) - min_camera_range_m: float = 0.28 - - envelope_pad_m: float = 0.015 - search_radius_m: float = 0.15 - overlap_accept: float = 0.20 - # Same-object views may differ in bounding size by partiality alone; a - # gap beyond this is two different bodies. - size_gap_max_m: float = 0.25 - # The majority of a candidate's points must lie within the error envelope - # of the track's accumulated support. Partial and newly revealed views of - # one object satisfy this; a different object placed at a vacated rest - # position does not, which is what AABB overlap cannot express. - support_explained: float = 0.5 - # A lifted cloud plainly spanning more than one object: wider than any - # single object at this rig's scale, or taller than one body. Repaired by - # stripping support-surface points and splitting by 3D connectivity. - split_extent_m: float = 0.30 - split_height_m: float = 0.10 - split_eps_m: float = 0.03 - # Same-frame observations whose clouds touch within this gap are one - # body - rigid objects cannot interpenetrate, and distinct objects on a - # workspace sit apart by more than sensor noise. This is what fuses - # whole-and-part duplicate proposals while identical twins, centimeters - # apart, stay two. - same_frame_merge_gap_m: float = 0.02 - # A support observed in a single keyframe is unconfirmed - nothing saw it - # from a second pose or moment, so it never becomes an instance. - min_member_observations: int = 2 - - # Naming abstention: a name is reported only when it clears the accept - # floor and beats the runner-up canonical group by the margin, at the - # frame and again over a track. Otherwise the instance stays unknown-N. - name_accept_score: float = 0.18 - name_refusal_margin: float = 0.06 - # An attachment must be the detector drawing a box around this member. - # Whole-object masks want a strict overlap; fragment masks of a large - # object overlap their object's box only partially. - name_attach_iou: float = 0.45 - - include_object_parts: bool = False - include_surfaces: bool = False - include_containers: bool = True - preserve_unknown_instances: bool = True - - in_scope: Callable[[np.ndarray], bool] | None = None - - -def aabb_overlap( - a_min: np.ndarray, - a_max: np.ndarray, - b_min: np.ndarray, - b_max: np.ndarray, - pad: float = 0.0, -) -> float: - """Intersection volume normalized by the smaller (padded) box volume. - - ``pad`` grows each box by the error envelope on every side, so the value - is an overlap of envelopes, not of raw partial-view boxes. Degenerate - axes are floored at 1 cm so thin objects (a pen, a sticky pad) do not - produce zero volumes. - """ - a_lo, a_hi = a_min - pad, a_max + pad - b_lo, b_hi = b_min - pad, b_max + pad - inter = np.minimum(a_hi, b_hi) - np.maximum(a_lo, b_lo) - if (inter <= 0).any(): - return 0.0 - floor = 0.01 - vol_a = float(np.prod(np.maximum(a_hi - a_lo, floor))) - vol_b = float(np.prod(np.maximum(b_hi - b_lo, floor))) - vol_i = float(np.prod(np.maximum(inter, 0.0))) - return vol_i / max(min(vol_a, vol_b), 1e-9) + refusal_margin: float = 0.15 # below this against a coexisting rival, flagged From 00f0c66a716940c987320731b09ac40c13090b96 Mon Sep 17 00:00:00 2001 From: bogwi Date: Fri, 4 Sep 2026 02:05:25 +0900 Subject: [PATCH 28/28] fix mypy --- dimos/perception/memory/localize.py | 16 ++++++++++++---- dimos/perception/memory/rig.py | 15 ++++++++------- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/dimos/perception/memory/localize.py b/dimos/perception/memory/localize.py index 49bf544a26..a7d64901a8 100644 --- a/dimos/perception/memory/localize.py +++ b/dimos/perception/memory/localize.py @@ -127,7 +127,9 @@ def add(self, det: Detection3DPC, radius: float, voxel: float) -> int: ) return hit - def _fuse(self, i: int, ts: float, points: np.ndarray, centroid: np.ndarray, voxel: float): + def _fuse( + self, i: int, ts: float, points: np.ndarray, centroid: np.ndarray, voxel: float + ) -> None: """Voxel-average a sighting into the group's union. Each fused point carries the raw points behind it, so the union is a @@ -162,6 +164,10 @@ def rows(self, i: int) -> np.ndarray: return np.asarray(self.members[i]).reshape(-1, M_WIDTH) +def _similarity(obs: Any) -> float: + return float(obs.similarity) + + def _settled(index: Stream[Any, Any], spacing: float) -> set[int]: """Ids left once sub-spacing duplicates are dropped. @@ -190,7 +196,9 @@ def _lift( ) -> list[Detection3DPC]: """Gate a frame's lifted detections.""" pose = rig.camera_pose(detections.ts, detections.image.frame_id) - lifted = rig.lift(detections, plane) if pose is not None else None + if pose is None: + return [] + lifted = rig.lift(detections, plane) if lifted is None: return [] camera = np.array([pose.position.x, pose.position.y, pose.position.z]) @@ -297,12 +305,12 @@ def localize( for q in queries: query_embedding = siglip.embed_text(q) - sightings = list( + sightings: list[Any] = list( source.search(query_embedding) .order_by("ts") .transform( peaks( - key=lambda obs: float(obs.similarity), + key=_similarity, prominence=policy.peak_prominence, distance=policy.peak_distance_s, width=policy.peak_width_s, diff --git a/dimos/perception/memory/rig.py b/dimos/perception/memory/rig.py index bef87de093..6635b82d11 100644 --- a/dimos/perception/memory/rig.py +++ b/dimos/perception/memory/rig.py @@ -316,9 +316,7 @@ def _images(store: Any, names: list[str]) -> tuple[dict[str, str], list[tuple[st if data.dtype == np.uint16 or data.dtype.kind == "f": depth.append((frame, name)) elif ( - data.ndim == 3 - and data.shape[2] == 3 - and not np.array_equal(data[..., 0], data[..., 1]) + data.ndim == 3 and data.shape[2] == 3 and not np.array_equal(data[..., 0], data[..., 1]) ): held = color.get(frame) if held is None or _rate(stream) > _rate(store.stream(held)): @@ -328,7 +326,9 @@ def _images(store: Any, names: list[str]) -> tuple[dict[str, str], list[tuple[st return color, depth -def _cameras(store: Any, roles: dict[str, Any], names: list[str], types: dict[str, type]): +def _cameras( + store: Any, roles: dict[str, Any], names: list[str], types: dict[str, type] +) -> dict[str, CameraInfo]: """Intrinsics per optical frame: an inline manifest dict, or every CameraInfo stream keyed by the frame it calibrates. @@ -355,7 +355,7 @@ def _cameras(store: Any, roles: dict[str, Any], names: list[str], types: dict[st if isinstance(role, str) else sorted(n for n in names if types[n] is CameraInfoMsg and store.stream(n).count()) ) - cameras = {} + cameras: dict[str, CameraInfo] = {} for name in found: # sorted, so a color info wins over its depth twin info = store.stream(name).first().data cameras.setdefault(info.frame_id, info) @@ -453,7 +453,7 @@ def from_store( claimed.add(color_name) color = store.stream(color_name) - color_frame = color.first().data.frame_id if color.count() else None + color_frame: str | None = color.first().data.frame_id if color.count() else None depth_name = roles.get("depth") if depth_name is None: # aligned depth shares the colour frame; otherwise the only depth stream @@ -482,7 +482,8 @@ def from_store( if not cameras: # embed-only stores may carry no calibration; geometry raises on use cameras = {OPTICAL_FRAME: cast("CameraInfo", None)} - if color_frame in cameras: + optical: str + if color_frame is not None and color_frame in cameras: optical = color_frame elif len(cameras) == 1: optical = next(iter(cameras)) # images stamped in a frame calibration never names