diff --git a/dimos/hardware/manipulators/xarm/adapter.py b/dimos/hardware/manipulators/xarm/adapter.py index ec3b09b565..2e421e0d9b 100644 --- a/dimos/hardware/manipulators/xarm/adapter.py +++ b/dimos/hardware/manipulators/xarm/adapter.py @@ -60,11 +60,20 @@ class XArmAdapter(ManipulatorAdapter): No inheritance required - just matching method signatures. """ - def __init__(self, address: str, dof: int = 6, **_: object) -> None: + def __init__( + self, + address: str, + dof: int = 6, + initial_positions: list[float] | None = None, + **_: object, + ) -> None: if not address: raise ValueError("address (IP) is required for XArmAdapter") self._ip = address self._dof = dof + self._initial_positions = ( + list(initial_positions) if initial_positions is not None else None + ) self._arm: XArmAPI | None = None self._control_mode: ControlMode = ControlMode.POSITION self._gripper_enabled: bool = False @@ -267,6 +276,8 @@ def _move_to_initial_pose(self) -> bool: return code == 0 def _initial_joints_degrees(self) -> list[float] | None: + if self._initial_positions is not None: + return [math.degrees(position) for position in self._initial_positions] if self._dof == 6: return _XARM6_INITIAL_JOINTS_DEG if self._dof == 7: diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 0dd9edbd9e..fe2e19a753 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -31,3 +31,6 @@ from dimos.robot.manipulators.xarm.blueprints.simulation import ( xarm_perception_sim as xarm_perception_sim, ) +from dimos.robot.manipulators.xarm.blueprints.worldbelief import ( + xarm6_worldbelief as xarm6_worldbelief, +) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index e16c1527e7..0c8278e384 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -359,7 +359,9 @@ def start(self) -> None: for name, port in self.inputs.items(): stream_name = self.config.stream_remapping.get(name, name) - stream: Stream[Any] = self.store.stream(stream_name, port.type) + stream: Stream[Any] = self.store.stream( + stream_name, port.type, **self._stream_kwargs(name) + ) self._port_to_stream(name, port, stream) logger.info("Recording %s -> %s (%s)", name, stream_name, port.type.__name__) @@ -389,10 +391,18 @@ async def on_msg(msg: Any) -> None: ts, getattr(msg, "ts", None), ) - stream.append(msg, ts=ts, pose=pose) + try: + stream.append(msg, ts=ts, pose=pose) + except sqlite3.ProgrammingError as exc: + if "closed database" in str(exc).lower(): + return + raise self.process_observable(input_topic.pure_observable(), on_msg) + def _stream_kwargs(self, name: str) -> dict[str, Any]: + return {} + def _prepare_streams(self) -> None: """On APPEND, drop the streams this recorder is about to (re)write — the remapped In-port streams plus ``tf`` — so a re-run replaces them instead diff --git a/dimos/models/embedding/clip.py b/dimos/models/embedding/clip.py index b0b4c99d76..87e7cefa5d 100644 --- a/dimos/models/embedding/clip.py +++ b/dimos/models/embedding/clip.py @@ -57,7 +57,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: Returns embeddings as torch.Tensor on device for efficient GPU comparisons. """ # Convert to PIL images - pil_images = [PILImage.fromarray(img.to_opencv()) for img in images] + pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] # Process images with torch.inference_mode(): diff --git a/dimos/models/embedding/dino.py b/dimos/models/embedding/dino.py new file mode 100644 index 0000000000..54d500c693 --- /dev/null +++ b/dimos/models/embedding/dino.py @@ -0,0 +1,84 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from functools import cached_property +from typing import Any, overload + +from PIL import Image as PILImage +import torch +import torch.nn.functional as functional +from transformers import AutoImageProcessor, AutoModel + +from dimos.models.base import HuggingFaceModel +from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig +from dimos.msgs.sensor_msgs.Image import Image + + +class DINOModelConfig(HuggingFaceEmbeddingModelConfig): + model_name: str = "facebook/dinov2-base" + dtype: torch.dtype = torch.float32 + + +class DINOModel(EmbeddingModel, HuggingFaceModel): + """DINOv2 image embedding model for visual instance identity.""" + + config: DINOModelConfig + _model_class = AutoModel + + @cached_property + def _model(self) -> Any: + return super()._model.eval() + + @cached_property + def _processor(self) -> AutoImageProcessor: + return AutoImageProcessor.from_pretrained(self.config.model_name) + + @overload + def embed(self, image: Image, /) -> Embedding: ... + @overload + def embed(self, *images: Image) -> list[Embedding]: ... + def embed(self, *images: Image) -> Embedding | list[Embedding]: + pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] + with torch.inference_mode(): + inputs = self._processor(images=pil_images, return_tensors="pt").to(self.config.device) + outputs = self._model(**inputs) + feats = getattr(outputs, "pooler_output", None) + if feats is None: + last_hidden = getattr(outputs, "last_hidden_state", None) + if last_hidden is None: + raise RuntimeError( + "DINO model did not return pooler_output or last_hidden_state" + ) + feats = last_hidden[:, 0] + if self.config.normalize: + feats = functional.normalize(feats, dim=-1) + + embeddings: list[Embedding] = [] + for i, feat in enumerate(feats): + embeddings.append(Embedding(vector=feat, timestamp=images[i].ts)) + return embeddings[0] if len(images) == 1 else embeddings + + @overload + def embed_text(self, text: str, /) -> Embedding: ... + @overload + def embed_text(self, *texts: str) -> list[Embedding]: ... + def embed_text(self, *texts: str) -> Embedding | list[Embedding]: + raise NotImplementedError("DINOModel does not support text embeddings") + + def stop(self) -> None: + if "_processor" in self.__dict__: + del self.__dict__["_processor"] + super().stop() diff --git a/dimos/models/embedding/mobileclip.py b/dimos/models/embedding/mobileclip.py index ff1f2efc5a..744fc3790b 100644 --- a/dimos/models/embedding/mobileclip.py +++ b/dimos/models/embedding/mobileclip.py @@ -66,7 +66,7 @@ def embed(self, *images: Image) -> Embedding | list[Embedding]: Returns embeddings as torch.Tensor on device for efficient GPU comparisons. """ # Convert to PIL images - pil_images = [PILImage.fromarray(img.to_opencv()) for img in images] + pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images] # Preprocess and batch with torch.inference_mode(): diff --git a/dimos/perception/detection/detectors/yoloe.py b/dimos/perception/detection/detectors/yoloe.py index e7cf79590d..8ba0dd06c8 100644 --- a/dimos/perception/detection/detectors/yoloe.py +++ b/dimos/perception/detection/detectors/yoloe.py @@ -25,6 +25,9 @@ from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.utils.data import get_data from dimos.utils.gpu_utils import is_cuda_available +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() class YoloePromptMode(Enum): @@ -43,6 +46,9 @@ def __init__( prompt_mode: YoloePromptMode = YoloePromptMode.LRPC, exclude_class_ids: list[int] | None = None, max_area_ratio: float | None = 0.3, + conf: float = 0.6, + iou: float = 0.6, + max_det: int | None = None, ) -> None: """ Initialize YOLO-E 2D detector. @@ -54,6 +60,9 @@ def __init__( prompt_mode: LRPC for prompt-free detection, PROMPT for text/visual prompting. exclude_class_ids: Class IDs to filter out from results (pass [] to disable). max_area_ratio: Maximum bbox area ratio (0-1) relative to image. + conf: Confidence threshold passed to Ultralytics track(). + iou: NMS IoU threshold passed to Ultralytics track(). + max_det: Optional maximum detections per frame. """ if model_name is None: if prompt_mode == YoloePromptMode.LRPC: @@ -62,9 +71,15 @@ def __init__( model_name = "yoloe-11s-seg.pt" self.model = YOLOE(get_data(model_path) / model_name) + self.detector_id = "yoloe" + self.detector_model_name = model_name + self.tracker_source = "ultralytics.track" self.prompt_mode = prompt_mode self._visual_prompts: dict[str, NDArray[Any]] | None = None self.max_area_ratio = max_area_ratio + self.conf = conf + self.iou = iou + self.max_det = max_det self._lock = threading.Lock() if prompt_mode == YoloePromptMode.PROMPT: @@ -80,6 +95,10 @@ def __init__( self.device = "cuda" else: self.device = "cpu" + logger.info( + f"YOLO-E detector loaded model={model_name} prompt_mode={prompt_mode.value} " + f"device={self.device} conf={self.conf} iou={self.iou} max_det={self.max_det}" + ) def set_prompts( self, @@ -120,11 +139,13 @@ def process_image(self, image: Image) -> "ImageDetections2D[Any]": track_kwargs = { "source": image.to_opencv(), "device": self.device, - "conf": 0.6, - "iou": 0.6, + "conf": self.conf, + "iou": self.iou, "persist": True, "verbose": False, } + if self.max_det is not None: + track_kwargs["max_det"] = self.max_det with self._lock: if self._visual_prompts is not None: diff --git a/dimos/perception/detection/identity_association.py b/dimos/perception/detection/identity_association.py new file mode 100644 index 0000000000..d6a390e8d0 --- /dev/null +++ b/dimos/perception/detection/identity_association.py @@ -0,0 +1,335 @@ +# Copyright 2025-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. + +"""Typed identity-association evidence and policy helpers.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass, field +from typing import Any + + +@dataclass(frozen=True) +class AssociationEvidence: + """Evidence for one observation-to-entity candidate.""" + + observation_name: str + candidate_id: str + reason: str + distance: float | None = None + distance_score: float | None = None + track_match: bool = False + label_match: bool | None = None + support: int = 0 + appearance_similarity: float | None = None + semantic_similarity: float | None = None + visual_similarity: float | None = None + appearance_mismatch: bool = False + semantic_mismatch: bool = False + visual_mismatch: bool = False + score: float = 0.0 + accepted: bool = False + ambiguous: bool = False + margin: float | None = None + vetoes: tuple[str, ...] = field(default_factory=tuple) + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["vetoes"] = list(self.vetoes) + return data + + +@dataclass(frozen=True) +class IdentityAssociationCandidate: + """Candidate entity for an observation-to-identity decision.""" + + candidate_id: str + distance: float + support: int + evidence: AssociationEvidence + departed: bool = False + + +@dataclass(frozen=True) +class AssociationDecision: + """Result of a pure association-policy decision.""" + + candidate_id: str | None + reason: str + accepted: bool + ambiguous: bool = False + margin: float | None = None + evidence: AssociationEvidence | None = None + + +class IdentityAssociationPolicy: + """Choose among already-compatible identity candidates.""" + + def choose_position_candidate( + self, + candidates: Sequence[IdentityAssociationCandidate], + *, + sticky_support: bool, + ambiguity_margin: float = 0.0, + ) -> AssociationDecision: + if not candidates: + return AssociationDecision(candidate_id=None, reason="none", accepted=False) + + if any(candidate.departed for candidate in candidates): + ranked = sorted(candidates, key=lambda candidate: candidate.distance) + else: + ranked = sorted( + candidates, + key=lambda candidate: ( + -candidate.evidence.score, + -candidate.support if sticky_support else 0, + candidate.distance, + candidate.candidate_id, + ), + ) + + best = ranked[0] + margin = None + if len(ranked) > 1: + margin = best.evidence.score - ranked[1].evidence.score + if margin < ambiguity_margin: + return AssociationDecision( + candidate_id=best.candidate_id, + reason="position_ambiguous", + accepted=False, + ambiguous=True, + margin=margin, + evidence=best.evidence, + ) + + return AssociationDecision( + candidate_id=best.candidate_id, + reason="position", + accepted=True, + margin=margin, + evidence=best.evidence, + ) + + def choose_observation_candidate( + self, + candidates: Sequence[IdentityAssociationCandidate], + *, + sticky_support: bool, + ambiguity_margin: float = 0.0, + ) -> AssociationDecision: + """Choose among identity candidates by evidence score.""" + if not candidates: + return AssociationDecision(candidate_id=None, reason="none", accepted=False) + + ranked = sorted( + candidates, + key=lambda candidate: ( + -self._candidate_score(candidate, sticky_support=sticky_support), + -candidate.support if sticky_support else 0, + candidate.distance, + candidate.candidate_id, + candidate.evidence.reason, + ), + ) + best = ranked[0] + margin = None + if len(ranked) > 1: + margin = self._candidate_score(best, sticky_support=sticky_support) - self._candidate_score( + ranked[1], + sticky_support=sticky_support, + ) + if margin < ambiguity_margin: + return AssociationDecision( + candidate_id=best.candidate_id, + reason="identity_ambiguous", + accepted=False, + ambiguous=True, + margin=margin, + evidence=best.evidence, + ) + + return AssociationDecision( + candidate_id=best.candidate_id, + reason=best.evidence.reason, + accepted=True, + margin=margin, + evidence=best.evidence, + ) + + def choose_frame_position_assignments( + self, + candidate_sets: Mapping[Any, Sequence[IdentityAssociationCandidate]], + *, + sticky_support: bool, + ) -> dict[Any, AssociationDecision]: + """Choose one-to-one position assignments for a full frame.""" + nonempty = {key: tuple(cands) for key, cands in candidate_sets.items() if cands} + assignments: dict[Any, AssociationDecision] = {} + for component_keys in self._assignment_components(nonempty): + component = {key: nonempty[key] for key in component_keys} + assignments.update( + self._choose_assignment_component(component, sticky_support=sticky_support) + ) + return assignments + + def _assignment_components( + self, + candidate_sets: Mapping[Any, Sequence[IdentityAssociationCandidate]], + ) -> list[tuple[Any, ...]]: + obs_to_eids = { + key: {candidate.candidate_id for candidate in candidates} + for key, candidates in candidate_sets.items() + } + eid_to_obs: dict[str, set[Any]] = {} + for key, eids in obs_to_eids.items(): + for eid in eids: + eid_to_obs.setdefault(eid, set()).add(key) + + components: list[tuple[Any, ...]] = [] + unseen = set(obs_to_eids) + while unseen: + root = unseen.pop() + stack = [root] + component = {root} + while stack: + key = stack.pop() + for eid in obs_to_eids.get(key, ()): + for other in eid_to_obs.get(eid, ()): + if other in unseen: + unseen.remove(other) + component.add(other) + stack.append(other) + components.append(tuple(component)) + return components + + def _candidate_score( + self, + candidate: IdentityAssociationCandidate, + *, + sticky_support: bool, + ) -> float: + score = float(candidate.evidence.score) + if sticky_support: + score += 1e-6 * float(candidate.support) + score -= 1e-9 * float(candidate.distance) + return score + + def _choose_assignment_component( + self, + candidate_sets: Mapping[Any, Sequence[IdentityAssociationCandidate]], + *, + sticky_support: bool, + ) -> dict[Any, AssociationDecision]: + # Keep pathological components bounded. + entity_ids = {candidate.candidate_id for candidates in candidate_sets.values() for candidate in candidates} + if len(candidate_sets) > 10 or len(entity_ids) > 12: + return self._choose_assignment_component_greedy( + candidate_sets, + sticky_support=sticky_support, + ) + + ordered_keys = tuple(sorted(candidate_sets, key=lambda key: (len(candidate_sets[key]), str(key)))) + ordered_candidates = { + key: tuple( + sorted( + candidate_sets[key], + key=lambda candidate: ( + -self._candidate_score(candidate, sticky_support=sticky_support), + candidate.distance, + candidate.candidate_id, + ), + ) + ) + for key in ordered_keys + } + best_metric = (-1, float("-inf"), float("-inf")) + best_assignment: dict[Any, IdentityAssociationCandidate] = {} + + def search( + index: int, + used: set[str], + score: float, + count: int, + distance_sum: float, + current: dict[Any, IdentityAssociationCandidate], + ) -> None: + nonlocal best_metric, best_assignment + if index >= len(ordered_keys): + metric = (count, score, -distance_sum) + if metric > best_metric: + best_metric = metric + best_assignment = dict(current) + return + + key = ordered_keys[index] + search(index + 1, used, score, count, distance_sum, current) + for candidate in ordered_candidates[key]: + if candidate.candidate_id in used: + continue + candidate_score = self._candidate_score(candidate, sticky_support=sticky_support) + if candidate_score <= 0.0: + continue + used.add(candidate.candidate_id) + current[key] = candidate + search( + index + 1, + used, + score + candidate_score, + count + 1, + distance_sum + float(candidate.distance), + current, + ) + current.pop(key, None) + used.remove(candidate.candidate_id) + + search(0, set(), 0.0, 0, 0.0, {}) + return { + key: AssociationDecision( + candidate_id=candidate.candidate_id, + reason="position", + accepted=True, + evidence=candidate.evidence, + ) + for key, candidate in best_assignment.items() + } + + def _choose_assignment_component_greedy( + self, + candidate_sets: Mapping[Any, Sequence[IdentityAssociationCandidate]], + *, + sticky_support: bool, + ) -> dict[Any, AssociationDecision]: + ranked: list[tuple[float, float, str, Any, IdentityAssociationCandidate]] = [] + for key, candidates in candidate_sets.items(): + for candidate in candidates: + score = self._candidate_score(candidate, sticky_support=sticky_support) + if score > 0.0: + ranked.append((score, -candidate.distance, candidate.candidate_id, key, candidate)) + ranked.sort(reverse=True) + used_keys: set[Any] = set() + used_entities: set[str] = set() + assignments: dict[Any, AssociationDecision] = {} + for _, _, _, key, candidate in ranked: + if key in used_keys or candidate.candidate_id in used_entities: + continue + used_keys.add(key) + used_entities.add(candidate.candidate_id) + assignments[key] = AssociationDecision( + candidate_id=candidate.candidate_id, + reason="position", + accepted=True, + evidence=candidate.evidence, + ) + return assignments diff --git a/dimos/perception/detection/identity_features.py b/dimos/perception/detection/identity_features.py new file mode 100644 index 0000000000..bb29cb63f3 --- /dev/null +++ b/dimos/perception/detection/identity_features.py @@ -0,0 +1,80 @@ +# Copyright 2025-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. + +"""Pure embedding helpers for object association.""" + +from __future__ import annotations + +from collections.abc import MutableSequence, Sequence +from typing import Any + +import numpy as np + + +def normalize_embedding(value: Any) -> np.ndarray | None: + """Return a unit float32 vector from an embedding-like value, or None if unusable.""" + if value is None: + return None + vector = np.asarray(value, dtype=np.float32).reshape(-1) + if vector.size == 0: + return None + norm = float(np.linalg.norm(vector)) + return vector / norm if norm > 0.0 else None + + +def gallery_cos(embedding: Any, gallery: Sequence[Any]) -> float | None: + """Best compatible cosine of an embedding against a stored view gallery.""" + emb = normalize_embedding(embedding) + if emb is None or not gallery: + return None + values: list[float] = [] + for item in gallery: + view = normalize_embedding(item) + if view is not None and view.size == emb.size: + values.append(float(np.dot(emb, view))) + return max(values) if values else None + + +def add_diverse_embedding_view( + gallery: MutableSequence[Any], + embedding: Any, + *, + novelty: float, + max_size: int, +) -> None: + """Add an embedding view only when it increases gallery diversity.""" + emb = normalize_embedding(embedding) + if emb is None or max_size <= 0: + return + if not gallery: + gallery.append(emb) + return + cosine = gallery_cos(emb, gallery) + if cosine is not None and cosine >= novelty: + return + if len(gallery) < max_size: + gallery.append(emb) + return + gallery[int(np.argmax(_embedding_redundancy_scores(gallery)))] = emb + + +def _embedding_redundancy_scores(gallery: Sequence[Any]) -> list[float]: + scores: list[float] = [] + for i, item in enumerate(gallery): + score = 0.0 + for j, other in enumerate(gallery): + if i != j: + score += gallery_cos(item, [other]) or 0.0 + scores.append(score) + return scores diff --git a/dimos/perception/detection/type/detection3d/object.py b/dimos/perception/detection/type/detection3d/object.py index f832eedaf4..4fdb679f42 100644 --- a/dimos/perception/detection/type/detection3d/object.py +++ b/dimos/perception/detection/type/detection3d/object.py @@ -21,6 +21,7 @@ import cv2 from dimos_lcm.geometry_msgs import Pose +from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose import numpy as np import open3d as o3d # type: ignore[import-untyped] @@ -58,6 +59,20 @@ class Object(Detection3D): camera_transform: Transform | None = None mask: np.ndarray[Any, np.dtype[np.uint8]] | None = None detections_count: int = 1 + embedding: np.ndarray[Any, np.dtype[np.float32]] | None = None + visual_embedding: np.ndarray[Any, np.dtype[np.float32]] | None = None + visual_embedding_model: str | None = None + visual_embedding_device: str | None = None + visual_embedding_dim: int | None = None + observation_partial: bool = False + observation_partial_from_camera_motion: bool = False + detector_id: str | None = None + detector_model: str | None = None + tracker_source: str | None = None + embedding_model: str | None = None + embedding_device: str | None = None + embedding_dim: int | None = None + observation_source: str | None = None def update_object(self, other: Object) -> None: """Update this object with data from another detection. @@ -92,12 +107,50 @@ def update_object(self, other: Object) -> None: self.ts = other.ts self.frame_id = other.frame_id self.image = other.image + if other.embedding is not None: + self.embedding = other.embedding + self.embedding_model = other.embedding_model + self.embedding_device = other.embedding_device + self.embedding_dim = other.embedding_dim + if other.visual_embedding is not None: + self.visual_embedding = other.visual_embedding + self.visual_embedding_model = other.visual_embedding_model + self.visual_embedding_device = other.visual_embedding_device + self.visual_embedding_dim = other.visual_embedding_dim + self.observation_partial = other.observation_partial + self.observation_partial_from_camera_motion = other.observation_partial_from_camera_motion + self.detector_id = other.detector_id + self.detector_model = other.detector_model + self.tracker_source = other.tracker_source + self.observation_source = other.observation_source self.detections_count += 1 def get_oriented_bounding_box(self) -> Any: """Get oriented bounding box of the pointcloud.""" return self.pointcloud.oriented_bounding_box + def _detection3d_bbox_components(self) -> tuple[Vector3, Quaternion, Vector3]: + """Return the fitted object geometry selected by ObjectSceneRegistration. + + The object point cloud is raw evidence. The stored center/size/pose are the + fitted planner-facing geometry after OSR has applied its box policy, such + as AABB/clamping. Publishing must not recompute a second OBB from raw + points, because a few table/background outliers can stretch the visible + Detection3D box and diverge from the geometry used by WorldBelief. + """ + center = self.center + size = self.size + orientation = getattr(self.pose, "orientation", Quaternion(0.0, 0.0, 0.0, 1.0)) + return ( + Vector3(center.x, center.y, center.z), + Quaternion(orientation.x, orientation.y, orientation.z, orientation.w), + Vector3( + max(float(size.x), 1e-3), + max(float(size.y), 1e-3), + max(float(size.z), 1e-3), + ), + ) + def scene_entity_label(self) -> str: """Get label for scene visualization.""" if self.detections_count > 1: @@ -105,18 +158,33 @@ def scene_entity_label(self) -> str: return f"{self.track_id}/{self.name} ({self.confidence:.0%})" def to_detection3d_msg(self) -> ROSDetection3D: - """Convert to ROS Detection3D message.""" - obb = self.get_oriented_bounding_box() - orientation = Quaternion.from_rotation_matrix(obb.R) + """Convert to ROS Detection3D message. + + ``Detection3D.id`` is the stable object_id. The detector's frame-local + track id remains available on the Object stream, but downstream obstacle + consumers need this message id to survive tracker resets. + """ + center, orientation, size = self._detection3d_bbox_components() msg = ROSDetection3D() msg.header = Header(self.ts, self.frame_id) - msg.id = str(self.track_id) - msg.bbox.center = Pose( - position=Vector3(obb.center[0], obb.center[1], obb.center[2]), - orientation=orientation, + msg.id = self.object_id + msg.results = [ + ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis( + class_id=str(self.class_id), + score=self.confidence, + ) + ) + ] + msg.results_length = len(msg.results) + msg.bbox = BoundingBox3D( + center=Pose( + position=center, + orientation=orientation, + ), + size=size, ) - msg.bbox.size = Vector3(obb.extent[0], obb.extent[1], obb.extent[2]) return msg @@ -159,6 +227,10 @@ def from_2d_to_list( max_distance: float = 0.0, use_aabb: bool = False, max_obstacle_width: float = 0.0, + detector_id: str | None = None, + detector_model: str | None = None, + tracker_source: str | None = None, + observation_source: str | None = None, ) -> list[Object]: """Create 3D Objects from 2D detections and RGBD images. @@ -184,6 +256,10 @@ def from_2d_to_list( Produces upright obstacles with identity orientation. max_obstacle_width: Clamp X/Y size to this value (meters). Useful for manipulation where obstacles must fit within the gripper. 0 disables. + detector_id: Optional detector implementation identifier for provenance. + detector_model: Optional model/checkpoint identifier for provenance. + tracker_source: Optional tracker implementation/source identifier. + observation_source: Optional producer module/source identifier. Returns: List of Object instances with pointclouds @@ -310,6 +386,10 @@ def from_2d_to_list( pose=pose, camera_transform=camera_transform, mask=store_mask, + detector_id=detector_id, + detector_model=detector_model, + tracker_source=tracker_source, + observation_source=observation_source, ) ) @@ -319,7 +399,7 @@ def from_2d_to_list( def aggregate_pointclouds(objects: list[Object]) -> PointCloud2: """Aggregate all object pointclouds into a single colored pointcloud. - Each object's points are colored based on its track_id. + Each object's points are colored based on its object_id. Args: objects: List of Object instances with pointclouds @@ -373,21 +453,22 @@ def aggregate_pointclouds(objects: list[Object]) -> PointCloud2: return pc -def to_detection3d_array(objects: list[Object]) -> Detection3DArray: - """Convert a list of Objects to a ROS Detection3DArray message. - - Args: - objects: List of Object instances - - Returns: - Detection3DArray ROS message - """ - array = Detection3DArray() - - if objects: - array.header = Header(objects[0].ts, objects[0].frame_id) - - for obj in objects: - array.detections.append(obj.to_detection3d_msg()) - - return array +def to_detection3d_array( + objects: list[Object], + *, + frame_id: str | None = None, + ts: float | None = None, +) -> Detection3DArray: + """Convert a list of Objects to a ROS Detection3DArray message.""" + detections = [obj.to_detection3d_msg() for obj in objects] + resolved_frame_id = frame_id + if resolved_frame_id is None: + resolved_frame_id = objects[0].frame_id if objects else "" + resolved_ts = ts + if resolved_ts is None: + resolved_ts = objects[0].ts if objects else 0.0 + return Detection3DArray( + detections_length=len(detections), + header=Header(resolved_ts, resolved_frame_id), + detections=detections, + ) diff --git a/dimos/perception/detection/world_belief.py b/dimos/perception/detection/world_belief.py new file mode 100644 index 0000000000..4a172544ea --- /dev/null +++ b/dimos/perception/detection/world_belief.py @@ -0,0 +1,1373 @@ +# Copyright 2025-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. + +"""Support-gated object identity association. + +``WorldBelief`` is API-compatible with ``ObjectDB`` but separates maintained +identity from the currently present set. An identity stays in the maintained +table until eviction; it becomes present only with recent support or explicit +promotion. + +Association uses explicit evidence records for tracker continuity, density-aware +geometry, embedding compatibility, and optional re-acquisition for recently +departed identities. Tracker ids are evidence, not authority. Per-frame +co-occurrence prevents two simultaneous tracks from collapsing onto one id. + +Time is driven by detection timestamps or explicit empty-frame timestamps. The +wall clock is only a bootstrap fallback, which keeps replayed streams and full +occlusion frames from accidentally evicting every maintained identity. +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from dataclasses import replace as dataclass_replace +import threading +import time +from typing import TYPE_CHECKING, Any + +import numpy as np +import open3d as o3d # type: ignore[import-untyped] + +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.perception.detection.identity_association import ( + AssociationDecision, + AssociationEvidence, + IdentityAssociationCandidate, + IdentityAssociationPolicy, +) +from dimos.perception.detection.identity_features import ( + add_diverse_embedding_view, + gallery_cos, + normalize_embedding, +) +from dimos.perception.detection.world_belief_history import WorldBeliefHistoryStore +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.perception.detection.type.detection3d.object import Object + +logger = setup_logger() + + +class WorldBelief: + """Support-gated object identity table with the ObjectDB public API.""" + + def __init__( + self, + distance_threshold: float = 0.2, + track_id_ttl_s: float = 5.0, + *, + min_support: int = 4, + recent_window: float = 1.5, + eviction_ttl_s: float = 60.0, + anchor_window: int = 9, + sticky_support: bool = True, + label_gate: bool = True, + cooccurrence_gate: bool = True, + enable_history: bool = False, + history_path: str | None = None, + history_stream: str = "worldbelief_obs", + reid_reacquire: bool = True, + reacq_window: float = 8.0, + reacq_radius: float | None = 1.0, + reacq_cos: float = 0.6, + reacq_margin: float = 0.10, + gallery_size: int = 8, + gallery_novelty: float = 0.88, + density_gate: bool = True, + density_factor: float = 0.45, + position_jitter_margin: float = 0.03, + position_ambiguity_margin: float = 0.03, + suppress_partial_observation_creates: bool = True, + candidate_gate_new_objects: bool = True, + new_candidate_min_age_s: float = 5.0, + new_candidate_ttl_s: float = 10.0, + new_candidate_max_count: int = 32, + reacq_override: bool = True, + ) -> None: + self._distance_threshold = distance_threshold + self._track_id_ttl_s = track_id_ttl_s + self._min_support = int(min_support) + self._recent_window = recent_window + self._eviction_ttl_s = eviction_ttl_s + self._anchor_window = int(anchor_window) + self._sticky_support = sticky_support + self._label_gate = label_gate + self._cooccurrence_gate = cooccurrence_gate + self._reid_reacquire = reid_reacquire + self._reacq_window = reacq_window + self._reacq_radius = reacq_radius + self._reacq_cos = reacq_cos + self._reacq_margin = reacq_margin + self._gallery_size = int(gallery_size) + self._gallery_novelty = gallery_novelty + self._density_gate = density_gate + self._density_factor = density_factor + self._position_jitter_margin = position_jitter_margin + self._position_ambiguity_margin = position_ambiguity_margin + self._identity_policy = IdentityAssociationPolicy() + self._suppress_partial_observation_creates = suppress_partial_observation_creates + self._candidate_gate_new_objects = candidate_gate_new_objects + self._new_candidate_min_age_s = new_candidate_min_age_s + self._new_candidate_ttl_s = new_candidate_ttl_s + self._new_candidate_max_count = int(new_candidate_max_count) + self._reacq_override = reacq_override + + self._entities: dict[str, Object] = {} + self._meta: dict[str, dict[str, Any]] = {} + self._track_id_map: dict[tuple[str, int], str] = {} + self._promoted: set[str] = set() # objects force-kept present (e.g. select_object RPC) + self._new_candidates: dict[str, dict[str, Any]] = {} + self._last_add_stats: dict[str, int] = {} + self._last_association_evidence: list[AssociationEvidence] = [] + self._now: float = 0.0 + self._lock = threading.RLock() + + self._history = WorldBeliefHistoryStore( + enabled=enable_history, + path=history_path, + stream_name=history_stream, + gallery_size=self._gallery_size, + gallery_novelty=self._gallery_novelty, + ) + + def _frame_time(self, objects: list[Object], frame_ts: float | None) -> float: + ts_vals = [float(o.ts) for o in objects if getattr(o, "ts", None) is not None] + if ts_vals: + return max(ts_vals) + if frame_ts is not None: + return max(self._now, float(frame_ts)) + return self._now or time.time() + + def _recent_cut(self) -> float: + return self._now - self._recent_window + + def _recent_support(self, eid: str, cut: float | None = None) -> int: + if cut is None: + cut = self._recent_cut() + return sum(1 for t in self._meta[eid]["window"] if t >= cut) + + def _is_present(self, eid: str, cut: float | None = None) -> bool: + return self._recent_support(eid, cut) >= self._min_support or eid in self._promoted + + def _is_departed(self, eid: str, cut: float | None = None) -> bool: + if cut is None: + cut = self._recent_cut() + return self._meta[eid]["last_seen"] < cut + + def _set_empty_add_stats(self) -> None: + self._last_add_stats = { + "input": 0, + "created": 0, + "updated": 0, + "matched_track": 0, + "matched_distance": 0, + "present": len(self.get_objects()), + "maintained": len(self._entities), + } + self._last_association_evidence = [] + + + def add_objects( + self, + objects: list[Object], + evict_exempt: set[str] | None = None, + frame_ts: float | None = None, + ) -> list[Object]: + """Ingest one frame of detections and return the current present objects.""" + stats = { + "input": len(objects), + "created": 0, + "updated": 0, + "matched_track": 0, + "matched_distance": 0, + } + with self._lock: + now = self._frame_time(objects, frame_ts) + self._now = now + scene_established_at_frame_start = len(self._entities) >= max(3, self._min_support) + frame_claims: dict[ + str, tuple[str, int] + ] = {} # eid -> observation claim key for THIS frame + frame_candidate_claims: dict[str, tuple[str, int]] = {} + self._last_association_evidence = [] + position_assignments = ( + self._frame_position_assignments(objects) + if len(objects) > 1 and self._entities + else {} + ) + processing_order = sorted( + range(len(objects)), + key=lambda idx: (idx not in position_assignments, idx), + ) + results_by_index: dict[int, Any] = {} + for idx in processing_order: + obj = objects[idx] + eid, reason = self._associate( + obj, + frame_claims, + position_assignment=position_assignments.get(idx), + ) + if eid is None: + # During cold start, even partial detections are useful inventory. Once a + # scene is established, suppress only moving-camera partials here; stable + # edge-visible objects can still mature through the candidate gate below. + suppress_partial = ( + self._suppress_partial_observation_creates + and getattr(obj, "observation_partial_from_camera_motion", False) + and scene_established_at_frame_start + ) + if suppress_partial: + stats["suppressed_partial"] = stats.get("suppressed_partial", 0) + 1 + continue + if self._candidate_gate_new_objects and scene_established_at_frame_start: + promoted = self._update_new_candidate(obj, now, frame_candidate_claims) + if promoted is None: + stats["pending_candidate"] = stats.get("pending_candidate", 0) + 1 + continue + new_obj = promoted + else: + new_obj = self._insert(obj, now) + results_by_index[idx] = new_obj + stats["created"] += 1 + # Prevent another same-frame detection from merging onto this new entity. + frame_claims[new_obj.object_id] = self._frame_claim_key(obj) + continue + if reason == "reid": + # Re-acquisition means the old position anchor should not resist the new location. + self._meta[eid]["positions"] = [] + self._update(eid, obj, now) + results_by_index[idx] = self._entities[eid] + stats["updated"] += 1 + stats["matched_track" if reason == "track" else "matched_distance"] += 1 + frame_claims[eid] = self._frame_claim_key(obj) + self._evict_stale(now, evict_exempt) + stats["present"] = len(self.get_objects()) + stats["maintained"] = len(self._entities) + self._last_add_stats = stats + if stats["created"] > 0: + logger.info(f"WorldBelief: {stats}") + return [results_by_index[idx] for idx in range(len(objects)) if idx in results_by_index] + + def advance_time(self, frame_ts: float, evict_exempt: set[str] | None = None) -> None: + """Advance the observation clock for a frame that produced no 3D objects. + + This lets the support-confirmed present-set decay during full occlusion or detector blackout + without falling back to wall-clock time on replayed streams. + """ + with self._lock: + if frame_ts <= 0: + return + now = max(self._now, float(frame_ts)) + self._now = now + self._evict_stale(now, evict_exempt) + self._set_empty_add_stats() + + def get_last_add_stats(self) -> dict[str, int]: + with self._lock: + return dict(self._last_add_stats) + + def get_last_association_evidence(self) -> list[dict[str, Any]]: + """Return evidence records from the most recent add/advance call.""" + with self._lock: + return [e.to_dict() for e in self._last_association_evidence] + + def get_identity_gallery_stats(self) -> dict[str, int]: + """Return public diagnostics for restored/stored identity feature galleries.""" + with self._lock: + semantic_views = sum(len(m.get("semantic_gallery") or ()) for m in self._meta.values()) + visual_views = sum(len(m.get("visual_gallery") or ()) for m in self._meta.values()) + return { + "entity_count": len(self._meta), + "semantic_gallery_entity_count": sum( + 1 for m in self._meta.values() if m.get("semantic_gallery") + ), + "semantic_gallery_view_count": semantic_views, + "visual_gallery_entity_count": sum( + 1 for m in self._meta.values() if m.get("visual_gallery") + ), + "visual_gallery_view_count": visual_views, + } + + def get_table_snapshot(self) -> list[dict[str, Any]]: + """Return compact metadata for the maintained object table.""" + with self._lock: + rows: list[dict[str, Any]] = [] + track_ids_by_entity: dict[str, list[int]] = defaultdict(list) + for (_namespace, track_id), eid in self._track_id_map.items(): + track_ids_by_entity[eid].append(int(track_id)) + + for eid, obj in sorted(self._entities.items()): + m = self._meta.get(eid, {}) + center = getattr(obj, "center", None) + size = getattr(obj, "size", None) + present = self._is_present(eid) + lifecycle = "present" if present else "maintained" + + rows.append( + { + "object_id": eid, + "name": getattr(obj, "name", ""), + "present": present, + "lifecycle": lifecycle, + "center": None + if center is None + else [float(center.x), float(center.y), float(center.z)], + "size": None + if size is None + else [float(size.x), float(size.y), float(size.z)], + "entered_t": float(m.get("entered_t", 0.0) or 0.0), + "last_seen": float(m.get("last_seen", 0.0) or 0.0), + "support": int(m.get("support", 0) or 0), + "recent_support": int(self._recent_support(eid)) + if eid in self._meta + else 0, + "track_ids": sorted(set(track_ids_by_entity.get(eid, []))), + "semantic_gallery_size": len(m.get("semantic_gallery") or ()), + "visual_gallery_size": len(m.get("visual_gallery") or ()), + "partial": bool(getattr(obj, "observation_partial", False)), + } + ) + return rows + + def get_objects(self) -> list[Object]: + """Present-set = support-confirmed within the recent window (or force-promoted).""" + with self._lock: + cut = self._recent_cut() + return [self._entities[eid] for eid in self._meta if self._is_present(eid, cut)] + + def get_all_objects(self) -> list[Object]: + with self._lock: + return list(self._entities.values()) + + def promote(self, object_id: str) -> bool: + with self._lock: + if object_id in self._entities: + self._promoted.add(object_id) + return True + return False + + def find_by_name(self, name: str) -> list[Object]: + with self._lock: + return [o for o in self.get_objects() if o.name == name] + + def find_by_object_id(self, object_id: str) -> Object | None: + with self._lock: + return self._entities.get(object_id) + + def find_nearest(self, position: Vector3, name: str | None = None) -> Object | None: + with self._lock: + cands = [ + o + for o in self.get_objects() + if o.center is not None and (name is None or o.name == name) + ] + if not cands: + return None + return min(cands, key=lambda o: position.distance(o.center)) + + def clear(self) -> None: + with self._lock: + for obj in self._entities.values(): + obj.pointcloud = PointCloud2( + pointcloud=o3d.geometry.PointCloud(), + frame_id=obj.pointcloud.frame_id, + ts=obj.pointcloud.ts, + ) + self._entities.clear() + self._meta.clear() + self._track_id_map.clear() + self._promoted.clear() + self._new_candidates.clear() + self._last_association_evidence = [] + logger.info("WorldBelief cleared") + + def agent_encode(self) -> list[dict[str, Any]]: + with self._lock: + return [o.agent_encode() for o in self.get_objects()] + + + @staticmethod + def _median_center(positions) -> Vector3: + m = np.median(np.asarray(positions, dtype=float), axis=0) + return Vector3(m) + + def _jitter_radius(self) -> float: + """Radius for recovering an object after small position jitter.""" + return max(2.0 * self._distance_threshold, 0.16) + + @staticmethod + def _semantic_emb_of(obj: Object): + return normalize_embedding(getattr(obj, "embedding", None)) + + @staticmethod + def _visual_emb_of(obj: Object): + return normalize_embedding(getattr(obj, "visual_embedding", None)) + + @staticmethod + def _role_gallery(m: dict[str, Any], role: str) -> list[Any]: + if role == "visual": + return m.get("visual_gallery") or ( + [m["visual_emb"]] if m.get("visual_emb") is not None else [] + ) + return m.get("semantic_gallery") or ( + [m["semantic_emb"]] if m.get("semantic_emb") is not None else [] + ) + + def _embedding_compatible_for_role(self, obj: Object, eid: str, role: str) -> bool: + """True when object and entity embeddings are comparable for one evidence role.""" + m = self._meta.get(eid) + if m is None: + return False + if role == "visual": + if self._visual_emb_of(obj) is None: + return False + obj_model = getattr(obj, "visual_embedding_model", None) + meta_model = m.get("visual_embedding_model") + obj_dim = getattr(obj, "visual_embedding_dim", None) + meta_dim = m.get("visual_embedding_dim") + else: + if self._semantic_emb_of(obj) is None: + return False + obj_model = getattr(obj, "embedding_model", None) + meta_model = m.get("semantic_embedding_model") + obj_dim = getattr(obj, "embedding_dim", None) + meta_dim = m.get("semantic_embedding_dim") + if obj_model is not None and meta_model is not None and obj_model != meta_model: + return False + if obj_dim is not None and meta_dim is not None and int(obj_dim) != int(meta_dim): + return False + return bool(self._role_gallery(m, role)) + + def _candidate_embedding_role(self, obj: Object, eid: str) -> str | None: + """Best comparable embedding role for this object/entity pair.""" + if self._embedding_compatible_for_role(obj, eid, "visual"): + return "visual" + if self._embedding_compatible_for_role(obj, eid, "semantic"): + return "semantic" + return None + + @staticmethod + def _tracker_key(obj: Object) -> tuple[str, int] | None: + track_id = getattr(obj, "track_id", -1) + if track_id is None or int(track_id) < 0: + return None + source = ( + getattr(obj, "tracker_source", None) or getattr(obj, "detector_id", None) or "default" + ) + return str(source), int(track_id) + + def _frame_claim_key(self, obj: Object) -> tuple[str, int]: + """Per-frame co-occurrence token, even when a detector has no tracker ids.""" + return self._tracker_key(obj) or ("observation", id(obj)) + + def _same_frame_duplicate(self, obj: Object, other: Object) -> bool: + """Allow near-identical same-frame duplicate detections to collapse onto one entity.""" + if obj.center is None or other.center is None: + return False + if self._label_gate and obj.name != other.name: + return False + if self._observation_identity_mismatch(obj, other): + return False + duplicate_radius = max(0.005, min(0.015, 0.15 * self._distance_threshold)) + return obj.center.distance(other.center) <= duplicate_radius + + def _appearance_similarity_for_role(self, obj: Object, eid: str, role: str) -> float | None: + emb = self._visual_emb_of(obj) if role == "visual" else self._semantic_emb_of(obj) + if emb is None or not self._embedding_compatible_for_role(obj, eid, role): + return None + m = self._meta.get(eid) + if m is None: + return None + return gallery_cos(emb, self._role_gallery(m, role)) + + def _appearance_mismatch(self, obj: Object, eid: str) -> bool: + """Return True when available identity evidence strongly contradicts a candidate. + + Visual/DINO evidence is preferred for physical identity. Semantic CLIP/SigLIP evidence is a fallback + only when no visual descriptor is available, so semantic similarity does not dominate lookalike cans. + """ + visual_similarity = self._appearance_similarity_for_role(obj, eid, "visual") + if visual_similarity is not None: + return visual_similarity < self._reacq_cos + semantic_similarity = self._appearance_similarity_for_role(obj, eid, "semantic") + if semantic_similarity is not None and self._visual_emb_of(obj) is None: + return semantic_similarity < self._reacq_cos + return False + + def _association_evidence( + self, + obj: Object, + eid: str, + *, + reason: str, + distance: float | None = None, + track_match: bool = False, + hard_appearance: bool = True, + accepted: bool = False, + ambiguous: bool = False, + margin: float | None = None, + ) -> AssociationEvidence: + distance_score = None + if distance is not None: + gate = max(self._distance_threshold, 1e-6) + distance_score = max(0.0, 1.0 - min(float(distance) / gate, 1.0)) + + meta = self._meta.get(eid, {}) + labels = meta.get("labels") + label_match = labels.most_common(1)[0][0] == obj.name if labels else None + support = int(meta.get("support", 0)) + semantic_similarity = self._appearance_similarity_for_role(obj, eid, "semantic") + visual_similarity = self._appearance_similarity_for_role(obj, eid, "visual") + appearance_similarity = ( + visual_similarity if visual_similarity is not None else semantic_similarity + ) + visual_mismatch = visual_similarity is not None and visual_similarity < self._reacq_cos + semantic_mismatch = ( + semantic_similarity is not None + and semantic_similarity < self._reacq_cos + and visual_similarity is None + ) + appearance_mismatch = visual_mismatch or semantic_mismatch + vetoes: list[str] = [] + if hard_appearance and visual_mismatch: + vetoes.append("visual") + elif hard_appearance and semantic_mismatch: + vetoes.append("appearance") + + support_score = min(support / max(self._min_support, 1), 1.0) + score = 0.0 + if distance_score is not None: + score += 0.45 * distance_score + if track_match: + score += 0.35 + if label_match is True: + score += 0.10 + if visual_similarity is not None: + score += 0.25 * max(0.0, visual_similarity) + elif semantic_similarity is not None: + score += 0.12 * max(0.0, semantic_similarity) + score += 0.05 * support_score + score -= float(len(vetoes)) + + return AssociationEvidence( + observation_name=obj.name, + candidate_id=eid, + reason=reason, + distance=distance, + distance_score=distance_score, + track_match=track_match, + label_match=label_match, + support=support, + appearance_similarity=appearance_similarity, + semantic_similarity=semantic_similarity, + visual_similarity=visual_similarity, + appearance_mismatch=appearance_mismatch, + semantic_mismatch=semantic_mismatch, + visual_mismatch=visual_mismatch, + score=score, + accepted=accepted, + ambiguous=ambiguous, + margin=margin, + vetoes=tuple(vetoes), + ) + + def _record_association_evidence(self, evidence: AssociationEvidence) -> None: + self._last_association_evidence.append(evidence) + + def _record_accepted_decision(self, decision: AssociationDecision) -> None: + if decision.evidence is None: + return + self._record_association_evidence( + dataclass_replace( + decision.evidence, + accepted=True, + ambiguous=False, + margin=decision.margin, + ) + ) + + def _record_accepted_association( + self, + obj: Object, + eid: str, + *, + reason: str, + distance: float | None = None, + track_match: bool = False, + hard_appearance: bool = True, + ) -> None: + self._record_association_evidence( + self._association_evidence( + obj, + eid, + reason=reason, + distance=distance, + track_match=track_match, + hard_appearance=hard_appearance, + accepted=True, + ) + ) + + def _record_ambiguous_association( + self, + obj: Object, + eid: str, + *, + reason: str, + distance: float | None = None, + track_match: bool = False, + hard_appearance: bool = True, + margin: float | None = None, + ) -> None: + self._record_association_evidence( + self._association_evidence( + obj, + eid, + reason=reason, + distance=distance, + track_match=track_match, + hard_appearance=hard_appearance, + accepted=False, + ambiguous=True, + margin=margin, + ) + ) + + def _observation_identity_mismatch(self, obj: Object, other: Object) -> bool: + """Negative identity evidence between two observations before an entity exists in _meta.""" + a_visual = self._visual_emb_of(obj) + b_visual = self._visual_emb_of(other) + if a_visual is not None and b_visual is not None and a_visual.size == b_visual.size: + a_model = getattr(obj, "visual_embedding_model", None) + b_model = getattr(other, "visual_embedding_model", None) + if (a_model is None or b_model is None or a_model == b_model) and float( + np.dot(a_visual, b_visual) + ) < self._reacq_cos: + return True + return False + a_emb = self._semantic_emb_of(obj) + b_emb = self._semantic_emb_of(other) + if a_emb is not None and b_emb is not None and a_emb.size == b_emb.size: + a_model = getattr(obj, "embedding_model", None) + b_model = getattr(other, "embedding_model", None) + if (a_model is None or b_model is None or a_model == b_model) and float( + np.dot(a_emb, b_emb) + ) < self._reacq_cos: + return True + return False + + def _reacquire( + self, + obj: Object, + frame_claims: dict[str, tuple[str, int]], + ) -> str | None: + """Re-acquire a recently departed identity by appearance.""" + best, best_cos, second_cos, _best_role = self._best_departed(obj, frame_claims) + if best is not None and best_cos >= self._reacq_cos: + if (best_cos - second_cos) >= self._reacq_margin: + return best + be = self._entities.get(best) + bd = ( + obj.center.distance(be.center) + if be is not None and be.center is not None and obj.center is not None + else None + ) + self._record_ambiguous_association( + obj, + best, + reason="reid_ambiguous", + distance=bd, + margin=best_cos - second_cos, + ) + return None + + def _best_departed( + self, + obj: Object, + frame_claims: dict[str, tuple[str, int]], + ): + """Best DEPARTED identity for this detection by gallery max-cos, with covisibility ELIMINATION + (present entities excluded) + recency/radius bounds. Returns (eid|None, best_cos, second_cos). + Shared by re-acquisition (no-position-match path) and the cross-collision override.""" + cut = self._recent_cut() + best, best_cos, second_cos, best_role = None, -1.0, -1.0, None + for eid, e in self._entities.items(): + if frame_claims.get(eid) is not None: # already taken this frame + continue + m = self._meta[eid] + if self._is_present(eid, cut): + continue + if (self._now - m["last_seen"]) > self._reacq_window: # departed too long ago + continue + if ( + self._reacq_radius is not None + and e.center is not None + and obj.center is not None + and obj.center.distance(e.center) > self._reacq_radius + ): # generous sanity bound only + continue + role = self._candidate_embedding_role(obj, eid) + if role is None: + continue + emb = self._visual_emb_of(obj) if role == "visual" else self._semantic_emb_of(obj) + if emb is None: + continue + gal = self._role_gallery(m, role) + cos = gallery_cos(emb, gal) + if cos is None: + continue + if cos > best_cos: + best, best_cos, second_cos, best_role = eid, cos, best_cos, role + elif cos > second_cos: + second_cos = cos + return best, best_cos, second_cos, best_role + + def _position_candidates( + self, + obj: Object, + frame_claims: dict[str, tuple[str, int]], + *, + record_evidence: bool = True, + ) -> list[tuple[float, str, Object, AssociationEvidence]]: + assert obj.center is not None + ents = [(eid, e) for eid, e in self._entities.items() if e.center is not None] + if self._density_gate and len(ents) > 1: + cands = self._density_candidates(obj, ents) + else: + cands = [] + for eid, e in ents: + d = obj.center.distance(e.center) + if d <= self._distance_threshold: + cands.append((d, eid, e)) + + claim_key = self._frame_claim_key(obj) + if self._cooccurrence_gate and cands: + cands = [ + c + for c in cands + if (claimed := frame_claims.get(c[1])) is None + or claimed == claim_key + or self._same_frame_duplicate(obj, c[2]) + ] + + if self._label_gate and cands: + same = [c for c in cands if self._meta[c[1]]["labels"].most_common(1)[0][0] == obj.name] + if same: + cands = same + if cands: + filtered: list[tuple[float, str, Object]] = [] + for d, eid, entity in cands: + evidence = self._association_evidence( + obj, + eid, + reason="position", + distance=d, + hard_appearance=True, + ) + if record_evidence: + self._record_association_evidence(evidence) + if evidence.vetoes: + continue + filtered.append((d, eid, entity, evidence)) + cands = filtered + return cands + + def _frame_position_assignments( + self, + objects: list[Object], + ) -> dict[int, AssociationDecision]: + candidate_sets: dict[int, list[IdentityAssociationCandidate]] = {} + empty_claims: dict[str, tuple[str, int]] = {} + for idx, obj in enumerate(objects): + if obj.center is None: + continue + cands = self._position_candidates(obj, empty_claims, record_evidence=False) + if not cands: + continue + candidate_sets[idx] = [ + IdentityAssociationCandidate( + candidate_id=eid, + distance=d, + support=int(self._meta[eid]["support"]), + evidence=evidence, + departed=self._is_departed(eid), + ) + for d, eid, _, evidence in cands + ] + if not candidate_sets: + return {} + return self._identity_policy.choose_frame_position_assignments( + candidate_sets, + sticky_support=self._sticky_support, + ) + + def _density_candidates( + self, + obj: Object, + ents: list[tuple[str, Object]], + ) -> list[tuple[float, str, Object]]: + assert obj.center is not None + recent_ents = [(eid, e) for eid, e in ents if not self._is_departed(eid)] + nearest: dict[str, float] = {} + for eid, e in ents: + neighbours = ( + recent_ents + if self._is_departed(eid) + else [(oeid, o) for oeid, o in recent_ents if oeid != eid] + ) + if neighbours: + nearest[eid] = min(e.center.distance(o.center) for _, o in neighbours) + else: + nearest[eid] = self._distance_threshold / max(self._density_factor, 1e-6) + + cands: list[tuple[float, str, Object]] = [] + for eid, e in ents: + d = obj.center.distance(e.center) + gate = ( + self._distance_threshold + if self._is_departed(eid) + else min(self._distance_threshold, self._density_factor * nearest[eid]) + ) + if d <= gate: + cands.append((d, eid, e)) + return cands + + def _choose_position_candidate( + self, + obj: Object, + cands: list[tuple[float, str, Object, AssociationEvidence]], + ): + policy_candidates = [ + IdentityAssociationCandidate( + candidate_id=eid, + distance=d, + support=int(self._meta[eid]["support"]), + evidence=evidence, + departed=self._is_departed(eid), + ) + for d, eid, _, evidence in cands + ] + decision = self._identity_policy.choose_position_candidate( + policy_candidates, + sticky_support=self._sticky_support, + ambiguity_margin=self._position_ambiguity_margin, + ) + if decision.candidate_id is None: + raise RuntimeError( + "position association policy returned no candidate for non-empty input" + ) + return decision + + def _prune_new_candidates(self, now: float) -> None: + if not self._new_candidates: + return + ttl = max(0.0, self._new_candidate_ttl_s) + if ttl > 0.0: + for cid, cand in list(self._new_candidates.items()): + if (now - float(cand.get("last_seen", now))) > ttl: + self._new_candidates.pop(cid, None) + max_count = max(0, self._new_candidate_max_count) + if max_count and len(self._new_candidates) > max_count: + ordered = sorted( + self._new_candidates.items(), + key=lambda item: float(item[1].get("last_seen", 0.0)), + ) + for cid, _ in ordered[: len(self._new_candidates) - max_count]: + self._new_candidates.pop(cid, None) + + def _update_new_candidate( + self, + obj: Object, + now: float, + frame_candidate_claims: dict[str, tuple[str, int]], + ) -> Object | None: + if obj.center is None: + return None + self._prune_new_candidates(now) + radius = self._jitter_radius() + key = None + claim_key = self._frame_claim_key(obj) + for cid, cand in self._new_candidates.items(): + center = cand.get("center") + cand_obj = cand.get("obj") + if ( + cand.get("name") == obj.name + and center is not None + and obj.center.distance(center) <= radius + ): + if cand_obj is not None and self._observation_identity_mismatch(obj, cand_obj): + continue + claimed = frame_candidate_claims.get(cid) + if ( + self._cooccurrence_gate + and claimed is not None + and claimed != claim_key + and cand_obj is not None + and not self._same_frame_duplicate(obj, cand_obj) + ): + continue + key = cid + break + if key is None: + key = obj.object_id + self._new_candidates[key] = { + "object_id": key, + "name": obj.name, + "center": obj.center, + "support": 1, + "first_seen": now, + "last_seen": now, + "last_support_ts": now, + "obj": obj, + } + frame_candidate_claims[key] = claim_key + return None + cand = self._new_candidates[key] + same_frame_support = float(cand.get("last_support_ts", float("nan"))) == now + if not same_frame_support: + cand["support"] += 1 + cand["last_support_ts"] = now + cand["center"] = obj.center + cand["last_seen"] = now + cand["obj"] = obj + frame_candidate_claims[key] = claim_key + self._prune_new_candidates(now) + if ( + cand["support"] < self._min_support + or (now - cand["first_seen"]) < self._new_candidate_min_age_s + ): + return None + promoted = cand["obj"] + promoted.object_id = cand["object_id"] + self._new_candidates.pop(key, None) + return self._insert(promoted, now) + + def _recent_position_jitter_candidate( + self, + obj: Object, + frame_claims: dict[str, tuple[str, int]], + ) -> str | None: + """Recover a recent, unclaimed same-label object after small camera jitter.""" + if obj.center is None: + return None + if len(self._entities) < 2 and self._tracker_key(obj) is not None: + return None + radius = self._jitter_radius() + recency_s = max(self._track_id_ttl_s, self._recent_window) + cands: list[tuple[float, str]] = [] + for eid, e in self._entities.items(): + if e.center is None: + continue + claimed = frame_claims.get(eid) + claim_key = self._frame_claim_key(obj) + if ( + self._cooccurrence_gate + and claimed is not None + and claimed != claim_key + and not self._same_frame_duplicate(obj, e) + ): + continue + m = self._meta[eid] + if (self._now - m["last_seen"]) > recency_s: + continue + if self._label_gate and m["labels"].most_common(1)[0][0] != obj.name: + continue + if self._appearance_mismatch(obj, eid): + continue + d = obj.center.distance(e.center) + if d <= radius: + cands.append((d, eid)) + if not cands: + return None + cands.sort(key=lambda x: x[0]) + if len(cands) > 1 and (cands[1][0] - cands[0][0]) < self._position_jitter_margin: + self._record_ambiguous_association( + obj, + cands[0][1], + reason="jitter_ambiguous", + distance=cands[0][0], + margin=cands[1][0] - cands[0][0], + hard_appearance=False, + ) + return None + return cands[0][1] + + def _associate( + self, + obj: Object, + frame_claims: dict[str, tuple[str, int]], + *, + position_assignment: AssociationDecision | None = None, + ) -> tuple[str | None, str | None]: + """Sticky, support-prioritised association. Returns (object_id|None, reason).""" + # priority 1: detector/tracker track_id (generous gate rejects stale-id reuse far away) + track_key = self._tracker_key(obj) + if track_key is not None: + eid = self._track_id_map.get(track_key) + if eid is not None and eid in self._entities: + # Do not let track-id continuity collapse two same-frame detections onto one entity. + claimed = frame_claims.get(eid) + claim_key = self._frame_claim_key(obj) + e = self._entities[eid] + blocked = ( + self._cooccurrence_gate + and claimed is not None + and claimed != claim_key + and not self._same_frame_duplicate(obj, e) + ) + meta = self._meta[eid] + label_blocked = self._label_gate and meta["labels"].most_common(1)[0][0] != obj.name + # Track ids are continuity evidence; visual mismatch can still veto them. + visual_similarity = self._appearance_similarity_for_role(obj, eid, "visual") + visual_blocked = ( + visual_similarity is not None and visual_similarity < self._reacq_cos + ) + semantic_blocked = self._is_departed(eid) and self._appearance_mismatch(obj, eid) + identity_blocked = visual_blocked or semantic_blocked + if identity_blocked: + block_dist = ( + obj.center.distance(e.center) + if e.center is not None and obj.center is not None + else None + ) + self._record_association_evidence( + self._association_evidence( + obj, + eid, + reason="track", + distance=block_dist, + track_match=True, + hard_appearance=True, + accepted=False, + ) + ) + if ( + not blocked + and not label_blocked + and not identity_blocked + and e.center is not None + and obj.center is not None + and (track_dist := obj.center.distance(e.center)) + <= 2 * self._distance_threshold + and (self._now - meta["last_seen"]) <= self._track_id_ttl_s + ): + # Let track and position evidence compete before accepting continuity. + track_evidence = self._association_evidence( + obj, + eid, + reason="track", + distance=track_dist, + track_match=True, + hard_appearance=self._is_departed(eid), + ) + self._record_association_evidence(track_evidence) + identity_candidates: list[IdentityAssociationCandidate] = [ + IdentityAssociationCandidate( + candidate_id=eid, + distance=track_dist, + support=int(meta["support"]), + evidence=track_evidence, + departed=self._is_departed(eid), + ) + ] + cands = self._position_candidates(obj, frame_claims) + identity_candidates.extend( + IdentityAssociationCandidate( + candidate_id=pos_eid, + distance=pos_dist, + support=int(self._meta[pos_eid]["support"]), + evidence=pos_evidence, + departed=self._is_departed(pos_eid), + ) + for pos_dist, pos_eid, _, pos_evidence in cands + if pos_eid != eid + ) + decision = self._identity_policy.choose_observation_candidate( + identity_candidates, + sticky_support=self._sticky_support, + ) + if decision.accepted and decision.candidate_id is not None: + self._record_accepted_decision(decision) + return decision.candidate_id, ( + "track" if decision.reason == "track" else "distance" + ) + if not blocked and not identity_blocked: + # Stale or teleported track id: fall through to geometry. + del self._track_id_map[track_key] + + if obj.center is None: + return None, None + + if ( + position_assignment is not None + and position_assignment.accepted + and position_assignment.candidate_id is not None + and position_assignment.candidate_id in self._entities + ): + assigned_eid = position_assignment.candidate_id + assigned_entity = self._entities[assigned_eid] + claimed = frame_claims.get(assigned_eid) + claim_key = self._frame_claim_key(obj) + blocked = ( + self._cooccurrence_gate + and claimed is not None + and claimed != claim_key + and not self._same_frame_duplicate(obj, assigned_entity) + ) + if not blocked and not self._appearance_mismatch(obj, assigned_eid): + assigned_dist = ( + obj.center.distance(assigned_entity.center) + if assigned_entity.center is not None + else None + ) + self._record_accepted_association( + obj, assigned_eid, reason="distance", distance=assigned_dist + ) + return assigned_eid, "distance" + + cands = self._position_candidates(obj, frame_claims) + if not cands: + jitter = self._recent_position_jitter_candidate(obj, frame_claims) + if jitter is not None: + je = self._entities[jitter] + jd = obj.center.distance(je.center) if je.center is not None else None + self._record_accepted_association(obj, jitter, reason="jitter", distance=jd) + return jitter, "distance" + if self._reid_reacquire: + reid = self._reacquire(obj, frame_claims) + if reid is not None: + re = self._entities[reid] + rd = obj.center.distance(re.center) if re.center is not None else None + self._record_accepted_association(obj, reid, reason="reid", distance=rd) + return reid, "reid" + return None, None + + pos_decision = self._choose_position_candidate(obj, cands) + pos_eid = pos_decision.candidate_id + pos_dist = min(d for d, cid, _, _ in cands if cid == pos_eid) + if pos_decision.ambiguous: + self._record_ambiguous_association( + obj, + pos_eid, + reason=pos_decision.reason, + distance=pos_dist, + margin=pos_decision.margin, + ) + return None, None + + if self._reacq_override and self._reid_reacquire: + bd, bc, sc, role = self._best_departed(obj, frame_claims) + emb = ( + None + if role is None + else (self._visual_emb_of(obj) if role == "visual" else self._semantic_emb_of(obj)) + ) + if ( + emb is not None + and role is not None + and bd is not None + and bd != pos_eid + and bc >= self._reacq_cos + and (bc - sc) >= self._reacq_margin + ): + pg = self._role_gallery(self._meta[pos_eid], role) + pos_cos = gallery_cos(emb, pg) + if pos_cos is None or bc >= pos_cos + self._reacq_margin: + be = self._entities[bd] + bd_dist = obj.center.distance(be.center) if be.center is not None else None + self._record_accepted_association(obj, bd, reason="reid", distance=bd_dist) + return bd, "reid" + self._record_accepted_association(obj, pos_eid, reason="distance", distance=pos_dist) + return pos_eid, "distance" + + def _insert(self, obj: Object, now: float) -> Object: + if not obj.ts: + obj.ts = now + self._entities[obj.object_id] = obj + c = obj.center + semantic0 = self._semantic_emb_of(obj) + visual0 = self._visual_emb_of(obj) + self._meta[obj.object_id] = { + "entered_t": now, + "last_seen": now, + "support": 1, + "window": [now], + "labels": Counter([obj.name]), + "semantic_emb": semantic0, + "semantic_gallery": [semantic0] if semantic0 is not None else [], + "semantic_embedding_model": getattr(obj, "embedding_model", None), + "semantic_embedding_device": getattr(obj, "embedding_device", None), + "semantic_embedding_dim": getattr( + obj, "embedding_dim", int(semantic0.size) if semantic0 is not None else None + ), + "visual_emb": visual0, + "visual_gallery": [visual0] if visual0 is not None else [], + "visual_embedding_model": getattr(obj, "visual_embedding_model", None), + "visual_embedding_device": getattr(obj, "visual_embedding_device", None), + "visual_embedding_dim": getattr( + obj, "visual_embedding_dim", int(visual0.size) if visual0 is not None else None + ), + "positions": [[c.x, c.y, c.z]] if c is not None else [], + } + track_key = self._tracker_key(obj) + if track_key is not None: + self._track_id_map[track_key] = obj.object_id + self._history.append_object(obj, now, now) + return obj + + def _update(self, eid: str, obj: Object, now: float) -> None: + existing = self._entities[eid] + m = self._meta[eid] + partial = getattr(obj, "observation_partial", False) + stable_center = ( + self._median_center(m["positions"]) if partial and m["positions"] else existing.center + ) + stable_pose = existing.pose if partial else None + stable_size = existing.size if partial else None + stable_pointcloud = existing.pointcloud if partial else None + + existing.update_object(obj) # latest frame metadata + detection count + existing.ts = obj.ts or now + # Partial observations refresh identity without moving stable geometry. + if partial: + if stable_center is not None: + existing.center = stable_center + if stable_pose is not None: + existing.pose = stable_pose + if stable_center is not None: + existing.pose.position = stable_center + if stable_size is not None: + existing.size = stable_size + if stable_pointcloud is not None: + existing.pointcloud = stable_pointcloud + elif existing.center is not None: + m["positions"].append([existing.center.x, existing.center.y, existing.center.z]) + if len(m["positions"]) > self._anchor_window: + m["positions"] = m["positions"][-self._anchor_window :] + existing.center = self._median_center(m["positions"]) + m["support"] += 1 + m["last_seen"] = now + if not partial: + m["labels"][obj.name] += 1 + # Keep the reported label stable across detector wording jitter. + existing.name = m["labels"].most_common(1)[0][0] + semantic_emb = self._semantic_emb_of(obj) + visual_emb = self._visual_emb_of(obj) + if not partial: + if semantic_emb is not None: + obj_model = getattr(obj, "embedding_model", None) + meta_model = m.get("semantic_embedding_model") + if obj_model is not None and meta_model is not None and obj_model != meta_model: + m["semantic_gallery"] = [] + m["semantic_emb"] = semantic_emb + m["semantic_embedding_model"] = getattr(obj, "embedding_model", None) + m["semantic_embedding_device"] = getattr(obj, "embedding_device", None) + m["semantic_embedding_dim"] = getattr(obj, "embedding_dim", int(semantic_emb.size)) + add_diverse_embedding_view( + m.setdefault("semantic_gallery", []), + semantic_emb, + novelty=self._gallery_novelty, + max_size=self._gallery_size, + ) + if visual_emb is not None: + obj_model = getattr(obj, "visual_embedding_model", None) + meta_model = m.get("visual_embedding_model") + if obj_model is not None and meta_model is not None and obj_model != meta_model: + m["visual_gallery"] = [] + m["visual_emb"] = visual_emb + m["visual_embedding_model"] = getattr(obj, "visual_embedding_model", None) + m["visual_embedding_device"] = getattr(obj, "visual_embedding_device", None) + m["visual_embedding_dim"] = getattr( + obj, "visual_embedding_dim", int(visual_emb.size) + ) + add_diverse_embedding_view( + m.setdefault("visual_gallery", []), + visual_emb, + novelty=self._gallery_novelty, + max_size=self._gallery_size, + ) + m["window"].append(now) + cut = now - self._recent_window + if len(m["window"]) > 16 and m["window"][0] < cut: + m["window"] = [t for t in m["window"] if t >= cut] + track_key = self._tracker_key(obj) + if track_key is not None: + self._track_id_map[track_key] = eid + self._history.append_object( + existing, + m["entered_t"], + now, + identity_obj=obj, + include_identity=not partial, + ) + + def _evict_stale(self, now: float, evict_exempt: set[str] | None = None) -> None: + exempt = evict_exempt or () + dead = [ + eid + for eid, m in self._meta.items() + if (now - m["last_seen"]) > self._eviction_ttl_s and eid not in exempt + ] + for eid in dead: + self._entities.pop(eid, None) + self._meta.pop(eid, None) + self._promoted.discard(eid) + for tid, mapped in list(self._track_id_map.items()): + if mapped == eid: + del self._track_id_map[tid] + + + def when_entered(self, object_id: str) -> float | None: + """First-seen sim time of an object from live state or durable evidence history.""" + with self._lock: + m = self._meta.get(object_id) + if m is not None: + return m["entered_t"] + entered = self._history.when_entered(object_id) + return entered + + def rehydrate(self) -> None: + """Seed the online maintained table from durable Memory2 identity evidence.""" + if not self._history.enabled: + return + with self._lock: + state = self._history.rehydrate() + self._entities.clear() + self._entities.update(state.entities) + self._meta.clear() + self._meta.update(state.meta) + self._track_id_map.clear() + self._promoted.clear() + self._new_candidates.clear() + self._now = state.now + logger.info( + f"WorldBelief rehydrated {len(self._entities)} entities from {self._history.path}" + ) + + def close(self) -> None: + """Release the delegated history store (call on shutdown).""" + self._history.close() + + def __len__(self) -> int: + with self._lock: + return len(self.get_objects()) + + def __repr__(self) -> str: + with self._lock: + return ( + f"WorldBelief(present={len(self.get_objects())}, maintained={len(self._entities)})" + ) diff --git a/dimos/perception/detection/world_belief_history.py b/dimos/perception/detection/world_belief_history.py new file mode 100644 index 0000000000..40f9c8f944 --- /dev/null +++ b/dimos/perception/detection/world_belief_history.py @@ -0,0 +1,375 @@ +# Copyright 2025-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. + +"""Memory2-backed durable history for WorldBelief.""" + +from __future__ import annotations + +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +import re +from typing import TYPE_CHECKING, Any + +import numpy as np +import open3d as o3d # type: ignore[import-untyped] + +from dimos.models.embedding.base import Embedding +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.perception.detection.identity_features import ( + add_diverse_embedding_view, + normalize_embedding, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.perception.detection.type.detection3d.object import Object + +logger = setup_logger() + + +@dataclass(frozen=True) +class RehydratedWorldBeliefState: + """Compact state restored from durable object evidence history.""" + + entities: dict[str, Object] + meta: dict[str, dict[str, Any]] + now: float + + +class WorldBeliefHistoryStore: + """Memory2 history adapter for WorldBelief durable identity evidence.""" + + def __init__( + self, + *, + enabled: bool, + path: str | None, + stream_name: str, + gallery_size: int, + gallery_novelty: float, + ) -> None: + self.enabled = enabled + self.path = path + self.stream_name = stream_name + self.gallery_size = int(gallery_size) + self.gallery_novelty = float(gallery_novelty) + self._store: Any = None + self._stream: Any = None + self._vector_streams: dict[str, Any] = {} + + def ensure_open(self) -> None: + """Lazily open the Memory2 SqliteStore used for durable identity history.""" + if not self.enabled or self._stream is not None: + return + from dimos.memory2.store.sqlite import SqliteStore + + if self.path is None: + return + Path(self.path).parent.mkdir(parents=True, exist_ok=True) + self._store = SqliteStore(path=self.path) + self._store.start() + self._stream = self._store.stream(self.stream_name, dict, codec="pickle") + + def close(self) -> None: + if self._store is not None: + try: + self._store.stop() + except Exception: + pass + self._store = None + self._stream = None + self._vector_streams = {} + + @staticmethod + def _safe_stream_part(value: Any) -> str: + text = str(value or "unknown").lower() + text = re.sub(r"[^a-z0-9]+", "_", text).strip("_") + return text or "unknown" + + def _embedding_stream_name(self, role: str, model: Any, dim: int) -> str: + role_part = self._safe_stream_part(role) + model_part = self._safe_stream_part(model) + return f"{self.stream_name}_{role_part}_{model_part}_{int(dim)}" + + def _embedding_stream(self, role: str, model: Any, dim: int) -> Any: + name = self._embedding_stream_name(role, model, dim) + stream = self._vector_streams.get(name) + if stream is None: + assert self._store is not None + stream = self._store.stream(name, dict, codec="pickle") + self._vector_streams[name] = stream + return stream + + def _append_vector_evidence( + self, + *, + role: str, + obj: Object, + entered_t: float, + now: float, + emb: np.ndarray[Any, np.dtype[np.float32]], + model: Any, + device: Any, + ) -> None: + stream = self._embedding_stream(role, model, int(emb.size)) + c = obj.center + stream.append( + { + "object_id": obj.object_id, + "name": obj.name, + "role": role, + "pos": [c.x, c.y, c.z] if c is not None else None, + "entered_t": entered_t, + "embedding_model": model, + "embedding_device": device, + "embedding_dim": int(emb.size), + }, + ts=now, + tags={"object_id": obj.object_id, "name": obj.name, "role": role}, + embedding=Embedding(vector=emb.astype(np.float32)), + ) + + def append_object( + self, + obj: Object, + entered_t: float, + now: float, + *, + identity_obj: Object | None = None, + include_identity: bool = True, + ) -> None: + """Append one durable object evidence event to Memory2.""" + if not self.enabled: + return + try: + self.ensure_open() + if self._stream is None: + return + c = obj.center + payload: dict[str, Any] = { + "object_id": obj.object_id, + "name": obj.name, + "pos": [c.x, c.y, c.z] if c is not None else None, + "entered_t": entered_t, + } + if include_identity: + source = identity_obj or obj + semantic_emb = normalize_embedding(getattr(source, "embedding", None)) + if semantic_emb is not None: + payload["semantic_embedding"] = semantic_emb.astype(np.float32).tolist() + payload["semantic_embedding_model"] = getattr(source, "embedding_model", None) + payload["semantic_embedding_device"] = getattr(source, "embedding_device", None) + payload["semantic_embedding_dim"] = int(semantic_emb.size) + visual_emb = normalize_embedding(getattr(source, "visual_embedding", None)) + if visual_emb is not None: + payload["visual_embedding"] = visual_emb.astype(np.float32).tolist() + payload["visual_embedding_model"] = getattr( + source, "visual_embedding_model", None + ) + payload["visual_embedding_device"] = getattr( + source, "visual_embedding_device", None + ) + payload["visual_embedding_dim"] = int(visual_emb.size) + self._stream.append( + payload, + ts=now, + tags={"object_id": obj.object_id, "name": obj.name}, + ) + if include_identity: + source = identity_obj or obj + if semantic_emb is not None: + self._append_vector_evidence( + role="semantic", + obj=obj, + entered_t=entered_t, + now=now, + emb=semantic_emb, + model=getattr(source, "embedding_model", None), + device=getattr(source, "embedding_device", None), + ) + if visual_emb is not None: + self._append_vector_evidence( + role="visual", + obj=obj, + entered_t=entered_t, + now=now, + emb=visual_emb, + model=getattr(source, "visual_embedding_model", None), + device=getattr(source, "visual_embedding_device", None), + ) + except Exception as exc: + logger.warning( + f"WorldBelief history append failed; continuing without persistence: {exc}" + ) + + def when_entered(self, object_id: str) -> float | None: + if not self.enabled: + return None + self.ensure_open() + if self._stream is None: + return None + best = None + for obs in self._stream.tags(object_id=object_id): + d = getattr(obs, "data", None) or obs._data + et = d.get("entered_t", obs.ts) if isinstance(d, dict) else obs.ts + best = et if best is None else min(best, et) + return best + + def _aggregate_events(self) -> dict[str, dict[str, Any]]: + self.ensure_open() + if self._stream is None: + return {} + agg: dict[str, dict[str, Any]] = {} + for obs in self._stream: + d = getattr(obs, "data", None) or obs._data + if not isinstance(d, dict): + continue + oid = d.get("object_id") + if oid is None or d.get("pos") is None: + continue + a = agg.setdefault( + oid, + { + "entered": d.get("entered_t", obs.ts), + "last": obs.ts, + "pos": d["pos"], + "labels": Counter(), + "n": 0, + "semantic_emb": None, + "semantic_emb_ts": -1.0, + "semantic_gallery": [], + "semantic_embedding_model": None, + "semantic_embedding_device": None, + "semantic_embedding_dim": None, + "visual_emb": None, + "visual_emb_ts": -1.0, + "visual_gallery": [], + "visual_embedding_model": None, + "visual_embedding_device": None, + "visual_embedding_dim": None, + }, + ) + a["entered"] = min(a["entered"], d.get("entered_t", obs.ts)) + if obs.ts >= a["last"]: + a["last"], a["pos"] = obs.ts, d["pos"] + a["labels"][d.get("name", "object")] += 1 + a["n"] += 1 + + raw_semantic = d.get("semantic_embedding") + if raw_semantic is None and d.get("visual_embedding") is None: + raw_semantic = d.get("emb") + if raw_semantic is None: + raw_semantic = d.get("embedding") + semantic_emb = normalize_embedding(raw_semantic) + if semantic_emb is not None: + add_diverse_embedding_view( + a.setdefault("semantic_gallery", []), + semantic_emb, + novelty=self.gallery_novelty, + max_size=self.gallery_size, + ) + if obs.ts >= a["semantic_emb_ts"]: + a["semantic_emb"] = semantic_emb + a["semantic_emb_ts"] = obs.ts + a["semantic_embedding_model"] = d.get( + "semantic_embedding_model", + d.get("embedding_model"), + ) + a["semantic_embedding_device"] = d.get( + "semantic_embedding_device", + d.get("embedding_device"), + ) + a["semantic_embedding_dim"] = d.get( + "semantic_embedding_dim", + d.get("embedding_dim", int(semantic_emb.size)), + ) + + visual_emb = normalize_embedding(d.get("visual_embedding")) + if visual_emb is not None: + add_diverse_embedding_view( + a.setdefault("visual_gallery", []), + visual_emb, + novelty=self.gallery_novelty, + max_size=self.gallery_size, + ) + if obs.ts >= a["visual_emb_ts"]: + a["visual_emb"] = visual_emb + a["visual_emb_ts"] = obs.ts + a["visual_embedding_model"] = d.get("visual_embedding_model") + a["visual_embedding_device"] = d.get("visual_embedding_device") + a["visual_embedding_dim"] = d.get("visual_embedding_dim", int(visual_emb.size)) + + return agg + + def rehydrate(self) -> RehydratedWorldBeliefState: + """Return compact maintained-entity state restored from durable evidence.""" + from dimos.perception.detection.type.detection3d.object import Object + + if not self.enabled: + return RehydratedWorldBeliefState(entities={}, meta={}, now=0.0) + agg = self._aggregate_events() + img = Image(np.zeros((2, 2, 3), np.uint8)) + entities: dict[str, Object] = {} + meta: dict[str, dict[str, Any]] = {} + for oid, a in agg.items(): + p = a["pos"] + v = Vector3(p) + pc = PointCloud2(pointcloud=o3d.geometry.PointCloud(), frame_id="world", ts=a["last"]) + name = a["labels"].most_common(1)[0][0] + entities[oid] = Object( + object_id=oid, + center=v, + size=Vector3(0.05, 0.05, 0.1), + pose=PoseStamped(position=v), + pointcloud=pc, + image=img, + bbox=(0, 0, 1, 1), + track_id=-1, + class_id=0, + confidence=0.5, + name=name, + ts=a["last"], + embedding=a["semantic_emb"], + embedding_model=a["semantic_embedding_model"], + embedding_device=a["semantic_embedding_device"], + embedding_dim=a["semantic_embedding_dim"], + visual_embedding=a["visual_emb"], + visual_embedding_model=a["visual_embedding_model"], + visual_embedding_device=a["visual_embedding_device"], + visual_embedding_dim=a["visual_embedding_dim"], + ) + meta[oid] = { + "entered_t": a["entered"], + "last_seen": a["last"], + "support": a["n"], + "window": [], + "labels": a["labels"], + "semantic_emb": a["semantic_emb"], + "semantic_gallery": a["semantic_gallery"], + "semantic_embedding_model": a["semantic_embedding_model"], + "semantic_embedding_device": a["semantic_embedding_device"], + "semantic_embedding_dim": a["semantic_embedding_dim"], + "visual_emb": a["visual_emb"], + "visual_gallery": a["visual_gallery"], + "visual_embedding_model": a["visual_embedding_model"], + "visual_embedding_device": a["visual_embedding_device"], + "visual_embedding_dim": a["visual_embedding_dim"], + "positions": [p], + } + now = max((a["last"] for a in agg.values()), default=0.0) + return RehydratedWorldBeliefState(entities=entities, meta=meta, now=now) diff --git a/dimos/perception/object_scene_registration.py b/dimos/perception/object_scene_registration.py index 9689e5b7b9..ac36dd9d16 100644 --- a/dimos/perception/object_scene_registration.py +++ b/dimos/perception/object_scene_registration.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import dataclass +import hashlib +import threading import time from typing import Any @@ -24,21 +27,22 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, ImageFormat from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.msgs.std_msgs.Header import Header from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray -from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode +from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.objectDB import ObjectDB from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection3d.object import ( Object, - Object as DetObject, aggregate_pointclouds, to_detection3d_array, ) +from dimos.perception.detection.world_belief import WorldBelief from dimos.types.timestamped import align_timestamped from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure @@ -46,8 +50,233 @@ logger = setup_logger() +@dataclass(frozen=True) +class _FrustumConfig: + enabled: bool = True + margin_px: float = 24.0 + near_m: float = 0.02 + + +_CAMERA_TRANSLATION_MOTION_THRESHOLD_M = 0.005 +_CAMERA_ROTATION_MOTION_THRESHOLD_RAD = 0.01 +_CAMERA_MOTION_SETTLE_FRAMES = 6 + +_OBJECTDB_TTL_S = 5.0 + +_WB_REACQ = { + "reacq_window": 30.0, + "reacq_radius": 1.0, + "reacq_cos": 0.6, + "reacq_margin": 0.10, +} +_WB_CANDIDATE = { + "suppress_partial_observation_creates": True, + "candidate_gate_new_objects": True, + "new_candidate_min_age_s": 1.0, + "new_candidate_ttl_s": 10.0, + "new_candidate_max_count": 32, +} + + +def _render_tracking_overlay( + color_image: Image, + frame_objects: list[Object], + present_objects: list[Object], +) -> Image | None: + img = color_image.to_opencv() + if img is None or getattr(img, "ndim", 0) != 3: + return None + out = img.copy() + height, width = out.shape[:2] + present_ids = {o.object_id for o in present_objects} + seen: set[str] = set() + visible_ids: set[str] = set() + + for obj in frame_objects: + oid = getattr(obj, "object_id", "") or "?" + if oid in seen: + continue + seen.add(oid) + bbox = getattr(obj, "bbox", None) + if bbox is None: + continue + try: + x1, y1, x2, y2 = (round(float(v)) for v in bbox) + except Exception: + continue + x1, x2 = sorted((max(0, min(width - 1, x1)), max(0, min(width - 1, x2)))) + y1, y2 = sorted((max(0, min(height - 1, y1)), max(0, min(height - 1, y2)))) + if x2 - x1 < 2 or y2 - y1 < 2: + continue + + visible_ids.add(oid) + digest = hashlib.md5(oid.encode("utf-8")).digest() + color = (int(60 + digest[0] % 180), int(60 + digest[1] % 180), int(60 + digest[2] % 180)) + state = "present" if oid in present_ids else "warming" + label = f"{oid[:6]} {getattr(obj, 'name', '') or 'object'} {state}" + cv2.rectangle(out, (x1, y1), (x2, y2), color, 2 if oid in present_ids else 1) + (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1) + top = max(0, y1 - th - 7) + cv2.rectangle(out, (x1, top), (min(width - 1, x1 + tw + 5), y1), color, -1) + cv2.putText( + out, + label, + (x1 + 2, max(th + 1, y1 - 4)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (0, 0, 0), + 1, + lineType=cv2.LINE_AA, + ) + + header = f"WorldBelief: {len(visible_ids)} visible, {len(present_objects)} present" + memory_only = [o for o in present_objects if o.object_id not in visible_ids] + if memory_only: + header += f", {len(memory_only)} memory" + text = "memory: " + ", ".join( + f"{o.object_id[:6]} {getattr(o, 'name', '') or 'object'}" for o in memory_only[:4] + ) + if len(memory_only) > 4: + text += f" +{len(memory_only) - 4}" + cv2.putText( + out, + text, + (8, height - 12), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (220, 220, 220), + 1, + lineType=cv2.LINE_AA, + ) + cv2.putText( + out, + header, + (8, 24), + cv2.FONT_HERSHEY_SIMPLEX, + 0.6, + (255, 255, 255), + 2, + lineType=cv2.LINE_AA, + ) + return Image.from_opencv(out, frame_id=color_image.frame_id, ts=color_image.ts) + + +def _camera_intrinsics_for_image( + camera_info: CameraInfo, + color_image: Image, +) -> tuple[float, float, float, float, int, int] | None: + K = camera_info.get_K_matrix() + fx, fy = float(K[0, 0]), float(K[1, 1]) + cx, cy = float(K[0, 2]), float(K[1, 2]) + if fx <= 0.0 or fy <= 0.0: + return None + + width = int(color_image.width) + height = int(color_image.height) + if width <= 0 or height <= 0: + return None + + info_width = int(getattr(camera_info, "width", 0) or width) + info_height = int(getattr(camera_info, "height", 0) or height) + sx = width / info_width if info_width > 0 else 1.0 + sy = height / info_height if info_height > 0 else 1.0 + return fx * sx, fy * sy, cx * sx, cy * sy, width, height + + +def _object_frustum_points(obj: Object) -> NDArray[np.float64] | None: + center = getattr(obj, "center", None) + size = getattr(obj, "size", None) + if center is None or size is None: + return None + + hx = max(float(size.x), 1e-3) * 0.5 + hy = max(float(size.y), 1e-3) * 0.5 + hz = max(float(size.z), 1e-3) * 0.5 + points = [[float(center.x), float(center.y), float(center.z), 1.0]] + for dx in (-hx, hx): + for dy in (-hy, hy): + for dz in (-hz, hz): + points.append( + [ + float(center.x) + dx, + float(center.y) + dy, + float(center.z) + dz, + 1.0, + ] + ) + return np.asarray(points, dtype=np.float64) + + +def _is_object_inside_camera_frustum( + obj: Object, + *, + camera_info: CameraInfo, + color_image: Image, + target_to_camera: Transform, + config: _FrustumConfig, +) -> bool | None: + """Return whether a target-frame object should be visible in the current camera image. + + None means the object lacks enough geometry/calibration for a safe decision; callers should not + exempt it from eviction based on an unknown. + """ + intrinsics = _camera_intrinsics_for_image(camera_info, color_image) + if intrinsics is None: + return None + fx, fy, cx, cy, width, height = intrinsics + + points = _object_frustum_points(obj) + if points is None: + return None + + cam_points = (target_to_camera.to_matrix() @ points.T).T[:, :3] + for x, y, z in cam_points: + z = float(z) + if z <= config.near_m: + continue + u = fx * (float(x) / z) + cx + v = fy * (float(y) / z) + cy + margin = config.margin_px + if -margin <= u <= width + margin and -margin <= v <= height + margin: + return True + return False + + +def _frustum_exempt_object_ids( + objects: list[Object], + *, + camera_info: CameraInfo, + color_image: Image, + camera_transform: Transform | None, + config: _FrustumConfig, +) -> set[str]: + if not config.enabled or not objects: + return set() + target_to_camera = ( + camera_transform.inverse() + if camera_transform is not None + else Transform( + frame_id=color_image.frame_id or "", + child_frame_id=color_image.frame_id or "", + ts=color_image.ts, + ) + ) + exempt: set[str] = set() + for obj in objects: + inside = _is_object_inside_camera_frustum( + obj, + camera_info=camera_info, + color_image=color_image, + target_to_camera=target_to_camera, + config=config, + ) + if inside is False: + exempt.add(obj.object_id) + return exempt + + class ObjectSceneRegistrationModule(Module): - """Module for detecting objects in camera images using YOLO-E with 2D and 3D detection.""" + """Module for registering detector observations as stable 3D scene objects.""" color_image: In[Image] depth_image: In[Image] @@ -55,52 +284,221 @@ class ObjectSceneRegistrationModule(Module): detections_2d: Out[Detection2DArray] detections_3d: Out[Detection3DArray] - objects: Out[list[DetObject]] + objects: Out[list[Object]] pointcloud: Out[PointCloud2] + annotated_image: Out[Image] + worldbelief_audit: Out[dict] - _detector: Yoloe2DDetector | None = None + _detector: Detector | None = None _camera_info: CameraInfo | None = None - _object_db: ObjectDB + _object_db: ObjectDB | WorldBelief _latest_depth_image: Image | None = None _latest_camera_transform: Any = None def __init__( self, target_frame: str = "map", - prompt_mode: YoloePromptMode = YoloePromptMode.LRPC, - # ObjectDB tuning + prompt_mode: Any = "lrpc", + belief_engine: str = "objectdb", + text_prompts: list[str] | None = None, + detector_model_name: str | None = None, + detector_device: str | None = None, + detector_conf: float = 0.6, + detector_iou: float = 0.6, + detector_max_det: int | None = None, + detector_max_area_ratio: float | None = 0.3, + detector_tracking_enabled: bool = True, + embedding_device: str | None = None, + embedding_model_id: str = "openai/clip-vit-base-patch32", + compute_visual_embeddings: bool = False, + visual_embedding_model_id: str = "facebook/dinov2-base", + visual_embedding_device: str | None = None, distance_threshold: float = 0.2, + # min_detections_for_permanent -> ObjectDB: cumulative detections (ever) to + # promote a pending object to permanent. + # min_support -> WorldBelief: detections within recent_window to trust an + # entity as "present". min_detections_for_permanent: int = 6, - # Object 3D reconstruction tuning + min_support: int = 4, + recent_window: float = 8, + eviction_ttl_s: float = 60.0, max_distance: float = 0.0, use_aabb: bool = False, max_obstacle_width: float = 0.0, + history_path: str | None = None, + compute_embeddings: bool = False, + # On/off toggles for the WorldBelief-only frustum + camera-motion lifecycle handling + frustum_eviction_exemption: bool = False, + camera_motion_marks_partial: bool = False, **kwargs: Any, ) -> None: super().__init__(**kwargs) self._target_frame = target_frame self._prompt_mode = prompt_mode - self._object_db = ObjectDB( - distance_threshold=distance_threshold, - min_detections_for_permanent=min_detections_for_permanent, - ) + self._text_prompts = tuple(text_prompts) if text_prompts else None + self._detector_model_name = detector_model_name + self._detector_device = detector_device + self._detector_conf = detector_conf + self._detector_iou = detector_iou + self._detector_max_det = detector_max_det + self._detector_max_area_ratio = detector_max_area_ratio + self._detector_tracking_enabled = detector_tracking_enabled + self._embedding_device = embedding_device + self._embedding_model_id = embedding_model_id + self._compute_visual_embeddings = compute_visual_embeddings + self._visual_embedding_model_id = visual_embedding_model_id + self._visual_embedding_device = visual_embedding_device or embedding_device + self._history_path = history_path + self._compute_embeddings = compute_embeddings + self._embedder: Any = None + self._visual_embedder: Any = None + self._active = False + self._processing_lock = threading.RLock() + self._last_camera_observation_transform: Transform | None = None + self._camera_motion_grace_frames_remaining = 0 + self._camera_motion_marks_partial = camera_motion_marks_partial + if belief_engine == "objectdb": + self._object_db = ObjectDB( + distance_threshold=distance_threshold, + min_detections_for_permanent=min_detections_for_permanent, + pending_ttl_s=_OBJECTDB_TTL_S, + track_id_ttl_s=_OBJECTDB_TTL_S, + ) + elif belief_engine == "world_belief": + self._object_db = WorldBelief( + distance_threshold=distance_threshold, + min_support=min_support, + recent_window=recent_window, + eviction_ttl_s=eviction_ttl_s, + history_path=history_path, + enable_history=history_path is not None, + **_WB_REACQ, + **_WB_CANDIDATE, + ) + else: + raise ValueError( + f"unknown belief engine: {belief_engine!r} (use 'world_belief' or 'objectdb')" + ) self._max_distance = max_distance self._use_aabb = use_aabb self._max_obstacle_width = max_obstacle_width + self._frustum_config = _FrustumConfig(enabled=frustum_eviction_exemption) + + def _prompt_mode_value(self) -> str: + return str(getattr(self._prompt_mode, "value", self._prompt_mode or "lrpc")).lower() + + def _coerce_yoloe_prompt_mode(self) -> Any: + from dimos.perception.detection.detectors.yoloe import YoloePromptMode + + value = self._prompt_mode_value() + try: + return YoloePromptMode(value) + except ValueError as exc: + allowed = ", ".join(mode.value for mode in YoloePromptMode) + raise ValueError(f"Unknown YOLO-E prompt mode {value!r}. Available: {allowed}") from exc + + def _apply_detector_prompts(self, detector: Detector) -> None: + if self._prompt_mode_value() != "prompt" or not self._text_prompts: + return + set_prompts = getattr(detector, "set_prompts", None) + if not callable(set_prompts): + logger.warning( + "Configured prompt text ignored because detector does not support set_prompts" + ) + return + set_prompts(text=list(self._text_prompts)) + + def _detector_provenance(self) -> dict[str, str | None]: + detector = self._detector + if detector is None: + return {} + detector_id = getattr(detector, "detector_id", type(detector).__name__) + detector_model = getattr(detector, "detector_model_name", None) + tracker_source = getattr(detector, "tracker_source", None) + if not self._detector_tracking_enabled: + tracker_source = None + return { + "detector_id": str(detector_id) if detector_id is not None else None, + "detector_model": str(detector_model) if detector_model is not None else None, + "tracker_source": str(tracker_source) if tracker_source is not None else None, + "observation_source": "object_scene_registration", + } @rpc def start(self) -> None: super().start() - - if self._prompt_mode == YoloePromptMode.LRPC: - model_name = "yoloe-11l-seg-pf.pt" - else: - model_name = "yoloe-11l-seg.pt" + self._active = True + + # Best-effort cross-session restore from the configured WorldBelief history DB. + if self._history_path is not None and hasattr(self._object_db, "rehydrate"): + try: + self._object_db.rehydrate() + except Exception as e: + logger.warning( + f"WorldBelief cross-session rehydrate failed (continuing fresh): {e}" + ) + + from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode + + prompt_mode = self._coerce_yoloe_prompt_mode() + detector_model_name = self._detector_model_name + if detector_model_name is None: + detector_model_name = ( + "yoloe-11l-seg-pf.pt" + if prompt_mode == YoloePromptMode.LRPC + else "yoloe-11l-seg.pt" + ) self._detector = Yoloe2DDetector( - model_name=model_name, - prompt_mode=self._prompt_mode, + model_name=detector_model_name, + device=self._detector_device, + prompt_mode=prompt_mode, + max_area_ratio=self._detector_max_area_ratio, + conf=self._detector_conf, + iou=self._detector_iou, + max_det=self._detector_max_det, ) + self._apply_detector_prompts(self._detector) + + if self._compute_embeddings and self._embedder is None: + try: + from dimos.models.embedding.clip import CLIPModel + + kwargs = {"model_name": self._embedding_model_id} + if self._embedding_device is not None: + kwargs["device"] = self._embedding_device + self._embedder = CLIPModel(**kwargs) + start = getattr(self._embedder, "start", None) + if callable(start): + start() + logger.info( + "CLIP embedder loaded model=%s device=%s", + self._embedding_model_id, + self._embedding_device, + ) + except Exception as e: + self._embedder = None + raise RuntimeError("semantic embedder requested but unavailable") from e + + if self._compute_visual_embeddings and self._visual_embedder is None: + try: + from dimos.models.embedding.dino import DINOModel + + kwargs = {"model_name": self._visual_embedding_model_id} + if self._visual_embedding_device is not None: + kwargs["device"] = self._visual_embedding_device + self._visual_embedder = DINOModel(**kwargs) + start = getattr(self._visual_embedder, "start", None) + if callable(start): + start() + logger.info( + "DINO embedder loaded model=%s device=%s", + self._visual_embedding_model_id, + self._visual_embedding_device, + ) + except Exception as e: + self._visual_embedder = None + raise RuntimeError("visual embedder requested but unavailable") from e self.camera_info.subscribe(lambda msg: setattr(self, "_camera_info", msg)) @@ -116,11 +514,27 @@ def start(self) -> None: def stop(self) -> None: """Stop the module and clean up resources.""" - if self._detector: - self._detector.stop() - self._detector = None + self._active = False + + detector = self._detector + self._detector = None + if detector: + stop = getattr(detector, "stop", None) + if callable(stop): + stop() + + for embedder_name in ("_embedder", "_visual_embedder"): + embedder = getattr(self, embedder_name, None) + setattr(self, embedder_name, None) + stop = getattr(embedder, "stop", None) + if callable(stop): + stop() - self._object_db.clear() + with self._processing_lock: + self._object_db.clear() + close = getattr(self._object_db, "close", None) + if callable(close): + close() logger.info("ObjectSceneRegistrationModule stopped") super().stop() @@ -132,27 +546,33 @@ def set_prompts( bboxes: NDArray[np.float64] | None = None, ) -> None: """Set prompts for detection. Provide either text or bboxes, not both.""" - if self._detector is not None: - self._detector.set_prompts(text=text, bboxes=bboxes) + if self._detector is None: + return + set_prompts = getattr(self._detector, "set_prompts", None) + if not callable(set_prompts): + logger.warning("Current detector does not support prompts") + return + set_prompts(text=text, bboxes=bboxes) @rpc - def select_object(self, track_id: int) -> dict[str, Any] | None: - """Get object data by track_id and promote to permanent.""" + def select_object(self, object_ref: int | str) -> dict[str, Any] | None: + """Get object data by stable object_id or detector track_id and promote it.""" + ref = str(object_ref) for obj in self._object_db.get_all_objects(): - if obj.track_id == track_id: + if obj.object_id == ref or str(obj.track_id) == ref: self._object_db.promote(obj.object_id) return obj.to_dict() return None @rpc def get_object_track_ids(self) -> list[int]: - """Get track_ids of all permanent objects.""" - return [obj.track_id for obj in self._object_db.get_all_objects()] + """Get track_ids of currently present objects.""" + return [obj.track_id for obj in self._object_db.get_objects()] @rpc def get_detected_objects(self) -> list[dict[str, Any]]: - """Get all detected objects with object_id (UUID) and name.""" - return [obj.agent_encode() for obj in self._object_db.get_all_objects()] + """Get currently present objects with object_id (UUID) and name.""" + return [obj.agent_encode() for obj in self._object_db.get_objects()] @rpc def get_object_pointcloud_by_name(self, name: str) -> PointCloud2 | None: @@ -210,10 +630,10 @@ def get_full_scene_pointcloud( depth_cv = depth_cv.copy() depth_cv[exclude_mask > 0] = 0 - # Build pointcloud from depth - fx, fy = self._camera_info.K[0], self._camera_info.K[4] - cx, cy = self._camera_info.K[2], self._camera_info.K[5] - intrinsic = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy) + K = self._camera_info.get_K_matrix() + intrinsic = o3d.camera.PinholeCameraIntrinsic( + w, h, float(K[0, 0]), float(K[1, 1]), float(K[0, 2]), float(K[1, 2]) + ) depth_o3d = o3d.geometry.Image(depth_cv.astype(np.float32)) pcd = o3d.geometry.PointCloud.create_from_depth_image( @@ -258,7 +678,10 @@ def detect(self, *prompts: str) -> str: if self._detector is None: return "Detector not initialized." - self._detector.set_prompts(text=list(prompts)) + set_prompts = getattr(self._detector, "set_prompts", None) + if not callable(set_prompts): + return "Current detector does not support prompts." + set_prompts(text=list(prompts)) time.sleep(2.0) detected = self.get_detected_objects() @@ -269,16 +692,17 @@ def detect(self, *prompts: str) -> str: return f"Detected {len(detected)} object(s):\n" + "\n".join(obj_list) @skill - def select(self, track_id: int) -> str: - """Select an object by track_id and promote it to permanent. + def select(self, object_ref: int | str) -> str: + """Select an object by stable object_id or detector track_id and promote it. Example: + select("2f4c...") select(5) """ - result = self.select_object(track_id) + result = self.select_object(object_ref) if result is None: - return f"No object found with track_id {track_id}." - return f"Selected object {track_id}: {result['name']}" + return f"No object found with object_id or track_id {object_ref!r}." + return f"Selected object {object_ref}: {result['name']}" def _on_aligned_frames(self, frames) -> None: # type: ignore[no-untyped-def] color_msg, depth_msg = frames @@ -286,7 +710,8 @@ def _on_aligned_frames(self, frames) -> None: # type: ignore[no-untyped-def] def _process_images(self, color_msg: Image, depth_msg: Image) -> None: """Process synchronized color and depth images (runs in background thread).""" - if not self._detector or not self._camera_info: + detector = self._detector + if not self._active or detector is None or self._camera_info is None: return color_image = color_msg @@ -301,7 +726,9 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> None: ) # Run 2D detection - detections_2d: ImageDetections2D[Any] = self._detector.process_image(color_image) + detections_2d: ImageDetections2D[Any] = detector.process_image(color_image) + if not self._active: + return detections_2d_msg = Detection2DArray( detections_length=len(detections_2d.detections), @@ -320,7 +747,7 @@ def _process_3d_detections( depth_image: Image, ) -> None: """Convert 2D detections to 3D and publish.""" - if self._camera_info is None: + if not self._active or self._camera_info is None: return # Cache depth image for full scene pointcloud generation @@ -337,10 +764,24 @@ def _process_3d_detections( ) if camera_transform is None: logger.info("Failed to lookup transform from camera frame to target frame") + evict_exempt = self._all_maintained_object_ids() + self._advance_belief_time(color_image.ts, evict_exempt) + self._publish_present_set( + color_image, + self._object_db.get_objects(), + frame_objects=[], + pointcloud_objects=[], + ) return # Cache camera transform for full scene pointcloud self._latest_camera_transform = camera_transform + camera_moving = ( + self._camera_transform_moved(camera_transform) + if self._camera_motion_marks_partial + else False + ) + evict_exempt = self._frustum_evict_exempt(color_image, camera_transform) objects = Object.from_2d_to_list( detections_2d=detections_2d, @@ -351,22 +792,229 @@ def _process_3d_detections( max_distance=self._max_distance, use_aabb=self._use_aabb, max_obstacle_width=self._max_obstacle_width, + **self._detector_provenance(), ) + if not self._active: + return + width = int(color_image.width) + height = int(color_image.height) + for obj in objects: + bbox = getattr(obj, "bbox", None) + if bbox is None: + continue + try: + x1, y1, x2, y2 = (float(v) for v in bbox) + except Exception: + continue + obj.observation_partial = ( + x1 <= 3.0 or y1 <= 3.0 or x2 >= width - 3.0 or y2 >= height - 3.0 + ) + if camera_moving: + obj.observation_partial = True + obj.observation_partial_from_camera_motion = True + if not objects: + self._advance_belief_time(color_image.ts, evict_exempt) + self._publish_present_set( + color_image, + self._object_db.get_objects(), + frame_objects=[], + pointcloud_objects=[], + ) return - # Add objects to spatial memory database - self._object_db.add_objects(objects) + self._attach_embeddings(objects, color_image) + if not self._active: + return + + frame_objects = self._add_objects_to_memory(objects, evict_exempt, color_image.ts) + + present_objects = self._object_db.get_objects() + self._publish_present_set( + color_image, + present_objects, + frame_objects=frame_objects, + pointcloud_objects=objects, + ) + return + + def _strip_detector_tracking(self, objects: list[Object]) -> None: + if self._detector_tracking_enabled: + return + for obj in objects: + obj.track_id = -1 + obj.tracker_source = None + + def _all_maintained_object_ids(self) -> set[str]: + if not isinstance(self._object_db, WorldBelief): + return set() + return {obj.object_id for obj in self._object_db.get_all_objects()} + + def _camera_transform_moved(self, camera_transform: Transform | None) -> bool: + if camera_transform is None: + self._last_camera_observation_transform = None + self._camera_motion_grace_frames_remaining = 0 + return False + prev = self._last_camera_observation_transform + self._last_camera_observation_transform = camera_transform + if prev is None: + return False + dp = camera_transform.translation.distance(prev.translation) + dr = prev.rotation.normalize().angle_to(camera_transform.rotation.normalize()) + moved = ( + dp > _CAMERA_TRANSLATION_MOTION_THRESHOLD_M + or dr > _CAMERA_ROTATION_MOTION_THRESHOLD_RAD + ) + if moved: + self._camera_motion_grace_frames_remaining = _CAMERA_MOTION_SETTLE_FRAMES + return True + if self._camera_motion_grace_frames_remaining > 0: + self._camera_motion_grace_frames_remaining -= 1 + return True + return False + + def _frustum_evict_exempt( + self, + color_image: Image, + camera_transform: Transform | None, + ) -> set[str]: + if self._camera_info is None or not isinstance(self._object_db, WorldBelief): + return set() + return _frustum_exempt_object_ids( + self._object_db.get_all_objects(), + camera_info=self._camera_info, + color_image=color_image, + camera_transform=camera_transform, + config=self._frustum_config, + ) - # Publish ALL permanent objects so downstream consumers get the full set, - # not just this frame's batch (which may be a subset of what's on the table). - all_permanent = self._object_db.get_objects() + def _advance_belief_time(self, frame_ts: float, evict_exempt: set[str]) -> None: + with self._processing_lock: + if isinstance(self._object_db, WorldBelief): + self._object_db.advance_time(frame_ts, evict_exempt=evict_exempt) + else: + self._object_db.add_objects([]) - detections_3d = to_detection3d_array(all_permanent) + def _add_objects_to_memory( + self, + objects: list[Object], + evict_exempt: set[str], + frame_ts: float, + ) -> list[Object]: + self._strip_detector_tracking(objects) + with self._processing_lock: + if isinstance(self._object_db, WorldBelief): + return self._object_db.add_objects( + objects, + evict_exempt=evict_exempt, + frame_ts=frame_ts, + ) + return self._object_db.add_objects(objects) + + def _publish_worldbelief_audit( + self, + color_image: Image, + present_objects: list[Object], + frame_objects: list[Object] | None = None, + ) -> None: + if not isinstance(self._object_db, WorldBelief): + return + frame_objects = frame_objects if frame_objects is not None else present_objects + try: + stats = self._object_db.get_last_add_stats() + event = { + "type": "worldbelief_frame_audit", + "ts": float(color_image.ts), + "frame_id": color_image.frame_id or "", + "target_frame": self._target_frame, + "detector_tracking_enabled": self._detector_tracking_enabled, + "present_count": len(present_objects), + "maintained_count": len(self._object_db.get_all_objects()), + "present_ids": [obj.object_id for obj in present_objects], + "frame_object_ids": [obj.object_id for obj in frame_objects], + "stats": stats, + "association_evidence": self._object_db.get_last_association_evidence(), + "gallery_stats": self._object_db.get_identity_gallery_stats(), + "table": self._object_db.get_table_snapshot(), + } + self.worldbelief_audit.publish(event) + except Exception as exc: + logger.debug("Failed to publish WorldBelief audit event: %s", exc) + + def _publish_present_set( + self, + color_image: Image, + present_objects: list[Object], + frame_objects: list[Object] | None = None, + pointcloud_objects: list[Object] | None = None, + ) -> None: + detections_3d = to_detection3d_array( + present_objects, frame_id=self._target_frame, ts=color_image.ts + ) self.detections_3d.publish(detections_3d) - self.objects.publish(all_permanent) + self.objects.publish(present_objects) + + overlay = _render_tracking_overlay( + color_image, + frame_objects=frame_objects if frame_objects is not None else present_objects, + present_objects=present_objects, + ) + if overlay is not None: + self.annotated_image.publish(overlay) - objects_for_pc = all_permanent - aggregated_pc = aggregate_pointclouds(objects_for_pc) + pointcloud_objects = pointcloud_objects if pointcloud_objects is not None else [] + if pointcloud_objects: + aggregated_pc = aggregate_pointclouds(pointcloud_objects) + else: + aggregated_pc = PointCloud2( + pointcloud=o3d.geometry.PointCloud(), + frame_id=self._target_frame, + ts=color_image.ts, + ) self.pointcloud.publish(aggregated_pc) - return + self._publish_worldbelief_audit(color_image, present_objects, frame_objects) + + def _attach_embeddings(self, objects: list[Object], color_image: Image) -> None: + def attach(model: Any, target: str) -> None: + crops: list[Image] = [] + idxs: list[int] = [] + for i, obj in enumerate(objects): + if getattr(obj, "bbox", None) is None: + continue + crop = obj.cropped_image(padding=0).to_rgb() + if crop.width < 2 or crop.height < 2: + continue + crop.frame_id = color_image.frame_id + crop.ts = color_image.ts + crops.append(crop) + idxs.append(i) + if not crops: + return + + try: + embeddings = model.embed(*crops) + if not isinstance(embeddings, list): + embeddings = [embeddings] + model_name = getattr(model, "model_name", None) or getattr( + getattr(model, "config", None), "model_name", None + ) + device = getattr(model, "device", None) + for i, embedding in zip(idxs, embeddings, strict=False): + emb = embedding.to_numpy().reshape(-1).astype(np.float32) + if target == "visual": + objects[i].visual_embedding = emb + objects[i].visual_embedding_model = model_name + objects[i].visual_embedding_device = device + objects[i].visual_embedding_dim = int(emb.size) + else: + objects[i].embedding = emb + objects[i].embedding_model = model_name + objects[i].embedding_device = device + objects[i].embedding_dim = int(emb.size) + except Exception as exc: + logger.warning("%s embedding step failed (skipping this frame): %s", target, exc) + + if self._embedder is not None: + attach(self._embedder, "semantic") + if self._visual_embedder is not None: + attach(self._visual_embedder, "visual") diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 61220a0fae..55832d304f 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -135,6 +135,7 @@ "xarm-perception-sim": "dimos.robot.manipulators.xarm.blueprints.simulation:xarm_perception_sim", "xarm-perception-sim-agent": "dimos.robot.manipulators.xarm.blueprints.agentic:xarm_perception_sim_agent", "xarm6-planner-only": "dimos.robot.manipulators.xarm.blueprints.basic:xarm6_planner_only", + "xarm6-worldbelief": "dimos.robot.manipulators.xarm.blueprints.worldbelief:xarm6_worldbelief", "xarm7-planner-coordinator": "dimos.robot.manipulators.xarm.blueprints.basic:xarm7_planner_coordinator", "xarm7-planner-coordinator-agent": "dimos.robot.manipulators.xarm.blueprints.agentic:xarm7_planner_coordinator_agent", } @@ -263,5 +264,6 @@ "wavefront-frontier-explorer": "dimos.navigation.frontier_exploration.wavefront_frontier_goal_selector.WavefrontFrontierExplorer", "web-input": "dimos.agents.web_human_input.WebInput", "websocket-vis-module": "dimos.web.websocket_vis.websocket_vis_module.WebsocketVisModule", + "x-arm6-world-belief-recorder": "dimos.robot.manipulators.xarm.worldbelief_recorder.XArm6WorldBeliefRecorder", "zed-camera": "dimos.hardware.sensors.camera.zed.camera.ZEDCamera", } diff --git a/dimos/robot/manipulators/xarm/blueprints/worldbelief.py b/dimos/robot/manipulators/xarm/blueprints/worldbelief.py new file mode 100644 index 0000000000..b244689837 --- /dev/null +++ b/dimos/robot/manipulators/xarm/blueprints/worldbelief.py @@ -0,0 +1,173 @@ +# 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. + +"""xArm6 WorldBelief hardware perception blueprint.""" + +from __future__ import annotations + +import math +from typing import Any + +import rerun.blueprint as rrb + +from dimos.core.coordination.blueprints import autoconnect +from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera +from dimos.manipulation.pick_and_place_module import PickAndPlaceModule +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.perception.object_scene_registration import ObjectSceneRegistrationModule +from dimos.robot.manipulators.common.blueprints import coordinator, trajectory_task +from dimos.robot.manipulators.xarm.config import ( + make_xarm6_model_config, + xarm6_hardware, +) +from dimos.robot.manipulators.xarm.worldbelief_recorder import ( + XArm6WorldBeliefRecorder, + xarm6_worldbelief_history_path, +) +from dimos.visualization.rerun.bridge import RerunBridgeModule + +XARM6_WORLDBELIEF_CAMERA_TRANSFORM = Transform( + translation=Vector3(x=0.06693724, y=-0.0309563, z=0.00691482), + rotation=Quaternion(0.70513398, 0.00535696, 0.70897578, -0.01052180), +) + +XARM6_WORLDBELIEF_HOME_JOINTS = ( + 0.0, + math.radians(-22.3), + math.radians(-22.3), + 0.0, + math.radians(-20.2), + 0.0, +) + + +def _topic_to_entity(topic: Any) -> str: + topic_name = str(getattr(topic, "name", topic)).split("#", 1)[0] + return { + "/annotated_image": "world/color_camera/annotated_image", + "/color_image": "world/color_camera/color_image", + "/camera_info": "world/color_camera", + "/depth_image": "world/depth_camera/depth_image", + "/depth_camera_info": "world/depth_camera", + "/detections_2d": "world/detections_2d", + "/detections_3d": "world/detections_3d", + "/pointcloud": "world/pointcloud", + }.get(topic_name, f"world/{topic_name.lstrip('/')}") + + +def _color_camera_info_to_rerun(msg: Any) -> list[tuple[str, Any]]: + optical_frame = getattr(msg, "frame_id", None) + return [ + *msg.to_rerun(image_topic="world/color_camera/color_image", optical_frame=optical_frame), + *msg.to_rerun( + image_topic="world/color_camera/annotated_image", optical_frame=optical_frame + ), + ] + + +def _depth_camera_info_to_rerun(msg: Any) -> list[tuple[str, Any]]: + optical_frame = getattr(msg, "frame_id", None) + return msg.to_rerun(image_topic="world/depth_camera/depth_image", optical_frame=optical_frame) + + +def _rerun_blueprint() -> rrb.Blueprint: + return rrb.Blueprint( + rrb.Horizontal( + rrb.Spatial2DView(origin="world/color_camera/annotated_image", name="Camera"), + rrb.Spatial3DView(origin="world", name="3D"), + ) + ) + + +_xarm6_worldbelief_hw = xarm6_hardware( + "arm", + gripper=True, + home_joints=list(XARM6_WORLDBELIEF_HOME_JOINTS), +) + +xarm6_worldbelief = autoconnect( + PickAndPlaceModule.blueprint( + robots=[ + make_xarm6_model_config( + name="arm", + add_gripper=True, + tf_extra_links=["link6"], + home_joints=list(XARM6_WORLDBELIEF_HOME_JOINTS), + ), + ], + planning_timeout=10.0, + visualization={"backend": "meshcat"}, + floor_z=-0.02, + ), + RealSenseCamera.blueprint( + width=640, + height=480, + fps=15, + base_frame_id="link6", + base_transform=XARM6_WORLDBELIEF_CAMERA_TRANSFORM, + ), + ObjectSceneRegistrationModule.blueprint( + target_frame="world", + belief_engine="world_belief", + prompt_mode="prompt", + text_prompts=["object", "container", "cup", "bottle", "can", "box"], + detector_model_name="yoloe-11s-seg.pt", + detector_conf=0.25, + detector_iou=0.8, + detector_max_det=20, + detector_device="cuda", + detector_tracking_enabled=False, + history_path=xarm6_worldbelief_history_path(), + compute_embeddings=True, + embedding_model_id="openai/clip-vit-base-patch32", + embedding_device="cuda", + compute_visual_embeddings=True, + visual_embedding_model_id="facebook/dinov2-base", + visual_embedding_device="cuda", + distance_threshold=0.08, + # WorldBelief engine: trust an object as "present" after 3 detections within + # recent_window. (min_support is WorldBelief's gate; min_detections_for_permanent + # is the ObjectDB equivalent and is unused here.) + min_support=3, + recent_window=8.0, + max_distance=1.0, + use_aabb=True, + max_obstacle_width=0.06, + frustum_eviction_exemption=True, + camera_motion_marks_partial=True, + ), + RerunBridgeModule.blueprint( + blueprint=_rerun_blueprint, + topic_to_entity=_topic_to_entity, + visual_override={ + "world/color_camera": _color_camera_info_to_rerun, + "world/depth_camera": _depth_camera_info_to_rerun, + }, + max_hz={ + "world/color_camera/annotated_image": 10.0, + "world/color_camera/color_image": 10.0, + "world/depth_camera/depth_image": 5.0, + "world/pointcloud": 5.0, + "world/detections_2d": 10.0, + "world/detections_3d": 10.0, + }, + ), + XArm6WorldBeliefRecorder.blueprint(), + coordinator( + hardware=[_xarm6_worldbelief_hw], + tasks=[trajectory_task(_xarm6_worldbelief_hw, name="traj_xarm")], + ), +).global_config(n_workers=7) diff --git a/dimos/robot/manipulators/xarm/worldbelief_recorder.py b/dimos/robot/manipulators/xarm/worldbelief_recorder.py new file mode 100644 index 0000000000..7c22ffcecc --- /dev/null +++ b/dimos/robot/manipulators/xarm/worldbelief_recorder.py @@ -0,0 +1,97 @@ +# Copyright 2025-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. + +"""Memory2 recorder for the xArm6 WorldBelief hardware blueprint.""" + +from __future__ import annotations + +from datetime import datetime +import inspect +from pathlib import Path +from typing import Any + +from dimos.constants import STATE_DIR +from dimos.core.core import rpc +from dimos.core.stream import In +from dimos.memory2.module import Recorder, RecorderConfig +from dimos.msgs.geometry_msgs.Transform import Transform +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.vision_msgs.Detection2DArray import Detection2DArray +from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.utils.logging_config import setup_logger + +_STATE_DIR = STATE_DIR / "worldbelief" / "xarm6" +_HISTORY_PATH = _STATE_DIR / "worldbelief_history.db" +_RECORDING_BASE_PATH = _STATE_DIR / "recordings" / "xarm6_worldbelief.db" + +logger = setup_logger() + + +def xarm6_worldbelief_history_path() -> str: + return str(_HISTORY_PATH) + + +def _timestamped_recording_path(base: str | Path, now: datetime | None = None) -> Path: + base_path = Path(base) + timestamp = (now or datetime.now()).strftime("%Y%m%d_%H%M%S_%f") + candidate = base_path.with_name(f"{base_path.stem}_{timestamp}{base_path.suffix}") + suffix = 1 + while candidate.exists(): + candidate = base_path.with_name(f"{base_path.stem}_{timestamp}_{suffix}{base_path.suffix}") + suffix += 1 + return candidate + + +class XArm6WorldBeliefRecorderConfig(RecorderConfig): + db_path: str | Path = _RECORDING_BASE_PATH + default_frame_id: str = "world" + tf_tolerance: float = 0.5 + record_tf: bool = False + + +class XArm6WorldBeliefRecorder(Recorder): + """Record replay-critical xArm6 WorldBelief streams for offline QA.""" + + color_image: In[Image] + depth_image: In[Image] + camera_info: In[CameraInfo] + depth_camera_info: In[CameraInfo] + annotated_image: In[Image] + detections_2d: In[Detection2DArray] + detections_3d: In[Detection3DArray] + pointcloud: In[PointCloud2] + worldbelief_audit: In[dict] + config: XArm6WorldBeliefRecorderConfig + + @rpc + def start(self) -> None: + if not self.config.g.replay: + db_path = _timestamped_recording_path(self.config.db_path) + db_path.parent.mkdir(parents=True, exist_ok=True) + self.config.db_path = db_path + logger.info("xArm6 WorldBelief recording DB: %s", db_path) + super().start() + + @staticmethod + def _stream_kwargs(name: str) -> dict[str, Any]: + return {"codec": "lcm"} if name == "depth_image" else {} + + async def _resolve_pose(self, name: str, msg: Any, ts: float) -> Any: + frame_id = getattr(msg, "frame_id", None) or self.config.default_frame_id + if frame_id == "world": + return Transform.identity().to_pose() + pose = super()._resolve_pose(name, msg, ts) + return await pose if inspect.isawaitable(pose) else pose