diff --git a/.envrc.nix b/.envrc.nix new file mode 100644 index 0000000000..4a6ade8151 --- /dev/null +++ b/.envrc.nix @@ -0,0 +1,5 @@ +if ! has nix_direnv_version || ! nix_direnv_version 3.0.6; then + source_url "https://raw.githubusercontent.com/nix-community/nix-direnv/3.0.6/direnvrc" "sha256-RYcUJaRMf8oF5LznDrlCXbkOQrywm0HDv1VjYGaJGdM=" +fi +use flake . +dotenv_if_exists diff --git a/.envrc.venv b/.envrc.venv new file mode 100644 index 0000000000..a4b314c6f7 --- /dev/null +++ b/.envrc.venv @@ -0,0 +1,2 @@ +source env/bin/activate +dotenv_if_exists diff --git a/.gitignore b/.gitignore index 12cb51509a..18fd575c85 100644 --- a/.gitignore +++ b/.gitignore @@ -48,3 +48,4 @@ yolo11n.pt # symlink one of .envrc.* if you'd like to use .envrc +.claude diff --git a/data/.lfs/models_mobileclip.tar.gz b/data/.lfs/models_mobileclip.tar.gz new file mode 100644 index 0000000000..874c94de07 --- /dev/null +++ b/data/.lfs/models_mobileclip.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1f8022e365d9e456dcbd3913d36bf8c68a4cd086eb777c92a773c8192cd8235d +size 277814612 diff --git a/data/.lfs/models_yolo.tar.gz b/data/.lfs/models_yolo.tar.gz index aca0915dfd..650d4617ca 100644 --- a/data/.lfs/models_yolo.tar.gz +++ b/data/.lfs/models_yolo.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0ed4a5160d4edfda145b6752b5c49ad22bc2887b66b9b9c38bd8c35fb5ffaf8f -size 9315806 +oid sha256:01796d5884cf29258820cf0e617bf834e9ffb63d8a4c7a54eea802e96fe6a818 +size 72476992 diff --git a/dimos/conftest.py b/dimos/conftest.py index 7e52a6191f..495afa8a24 100644 --- a/dimos/conftest.py +++ b/dimos/conftest.py @@ -14,6 +14,7 @@ import asyncio import threading + import pytest @@ -24,12 +25,35 @@ def event_loop(): loop.close() +_session_threads = set() _seen_threads = set() _seen_threads_lock = threading.RLock() +_before_test_threads = {} # Map test name to set of thread IDs before test _skip_for = ["lcm", "heavy", "ros"] +@pytest.hookimpl() +def pytest_sessionfinish(session): + """Track threads that exist at session start - these are not leaks.""" + + yield + + # Check for session-level thread leaks at teardown + final_threads = [ + t + for t in threading.enumerate() + if t.name != "MainThread" and t.ident not in _session_threads + ] + + if final_threads: + thread_info = [f"{t.name} (daemon={t.daemon})" for t in final_threads] + pytest.fail( + f"\n{len(final_threads)} thread(s) leaked during test session: {thread_info}\n" + "Session-scoped fixtures must clean up all threads in their teardown." + ) + + @pytest.fixture(autouse=True) def monitor_threads(request): # Skip monitoring for tests marked with specified markers @@ -37,24 +61,45 @@ def monitor_threads(request): yield return + # Capture threads before test runs + test_name = request.node.nodeid + with _seen_threads_lock: + _before_test_threads[test_name] = { + t.ident for t in threading.enumerate() if t.ident is not None + } + yield - threads = [t for t in threading.enumerate() if t.name != "MainThread"] + # Only check for threads created BY THIS TEST, not existing ones + with _seen_threads_lock: + before = _before_test_threads.get(test_name, set()) + current = {t.ident for t in threading.enumerate() if t.ident is not None} - if not threads: - return + # New threads are ones that exist now but didn't exist before this test + new_thread_ids = current - before - with _seen_threads_lock: - new_leaks = [t for t in threads if t.ident not in _seen_threads] - for t in threads: - _seen_threads.add(t.ident) + if not new_thread_ids: + return - if not new_leaks: - return + # Get the actual thread objects for new threads + new_threads = [ + t for t in threading.enumerate() if t.ident in new_thread_ids and t.name != "MainThread" + ] + + # Filter out threads we've already seen (from previous tests) + truly_new = [t for t in new_threads if t.ident not in _seen_threads] + + # Mark all new threads as seen + for t in new_threads: + if t.ident is not None: + _seen_threads.add(t.ident) + + if not truly_new: + return - thread_names = [t.name for t in new_leaks] + thread_names = [t.name for t in truly_new] - pytest.fail( - f"Non-closed threads before or during this test. The thread names: {thread_names}. " - "Please look at the first test that fails and fix that." - ) + pytest.fail( + f"Non-closed threads created during this test. Thread names: {thread_names}. " + "Please look at the first test that fails and fix that." + ) diff --git a/dimos/models/embedding/__init__.py b/dimos/models/embedding/__init__.py new file mode 100644 index 0000000000..981e25e5c2 --- /dev/null +++ b/dimos/models/embedding/__init__.py @@ -0,0 +1,30 @@ +from dimos.models.embedding.base import Embedding, EmbeddingModel + +__all__ = [ + "Embedding", + "EmbeddingModel", +] + +# Optional: CLIP support +try: + from dimos.models.embedding.clip import CLIPEmbedding, CLIPModel + + __all__.extend(["CLIPEmbedding", "CLIPModel"]) +except ImportError: + pass + +# Optional: MobileCLIP support +try: + from dimos.models.embedding.mobileclip import MobileCLIPEmbedding, MobileCLIPModel + + __all__.extend(["MobileCLIPEmbedding", "MobileCLIPModel"]) +except ImportError: + pass + +# Optional: TorchReID support +try: + from dimos.models.embedding.treid import TorchReIDEmbedding, TorchReIDModel + + __all__.extend(["TorchReIDEmbedding", "TorchReIDModel"]) +except ImportError: + pass diff --git a/dimos/models/embedding/base.py b/dimos/models/embedding/base.py new file mode 100644 index 0000000000..7f2e1896b9 --- /dev/null +++ b/dimos/models/embedding/base.py @@ -0,0 +1,148 @@ +# Copyright 2025 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 + +import time +from abc import ABC, abstractmethod +from typing import Generic, Optional, TypeVar + +import numpy as np +import torch + +from dimos.msgs.sensor_msgs import Image +from dimos.types.timestamped import Timestamped + + +class Embedding(Timestamped): + """Base class for embeddings with vector data. + + Supports both torch.Tensor (for GPU-accelerated comparisons) and np.ndarray. + Embeddings are kept as torch.Tensor on device by default for efficiency. + """ + + vector: torch.Tensor | np.ndarray + + def __init__(self, vector: torch.Tensor | np.ndarray, timestamp: Optional[float] = None): + self.vector = vector + if timestamp: + self.timestamp = timestamp + else: + self.timestamp = time.time() + + def __matmul__(self, other: "Embedding") -> float: + """Compute cosine similarity via @ operator.""" + if isinstance(self.vector, torch.Tensor): + other_tensor = other.to_torch(self.vector.device) + result = self.vector @ other_tensor + return result.item() + return float(self.vector @ other.to_numpy()) + + def to_numpy(self) -> np.ndarray: + """Convert to numpy array (moves to CPU if needed).""" + if isinstance(self.vector, torch.Tensor): + return self.vector.detach().cpu().numpy() + return self.vector + + def to_torch(self, device: str | torch.device | None = None) -> torch.Tensor: + """Convert to torch tensor on specified device.""" + if isinstance(self.vector, np.ndarray): + tensor = torch.from_numpy(self.vector) + return tensor.to(device) if device else tensor + + if device is not None and self.vector.device != torch.device(device): + return self.vector.to(device) + return self.vector + + def to_cpu(self) -> "Embedding": + """Move embedding to CPU, returning self for chaining.""" + if isinstance(self.vector, torch.Tensor): + self.vector = self.vector.cpu() + return self + + +E = TypeVar("E", bound="Embedding") + + +class EmbeddingModel(ABC, Generic[E]): + """Abstract base class for embedding models supporting vision and language.""" + + device: str + normalize: bool = True + + @abstractmethod + def embed(self, *images: Image) -> E | list[E]: + """ + Embed one or more images. + Returns single Embedding if one image, list if multiple. + """ + pass + + @abstractmethod + def embed_text(self, *texts: str) -> E | list[E]: + """ + Embed one or more text strings. + Returns single Embedding if one text, list if multiple. + """ + pass + + def compare_one_to_many(self, query: E, candidates: list[E]) -> torch.Tensor: + """ + Efficiently compare one query against many candidates on GPU. + + Args: + query: Query embedding + candidates: List of candidate embeddings + + Returns: + torch.Tensor of similarities (N,) + """ + query_tensor = query.to_torch(self.device) + candidate_tensors = torch.stack([c.to_torch(self.device) for c in candidates]) + return query_tensor @ candidate_tensors.T + + def compare_many_to_many(self, queries: list[E], candidates: list[E]) -> torch.Tensor: + """ + Efficiently compare all queries against all candidates on GPU. + + Args: + queries: List of query embeddings + candidates: List of candidate embeddings + + Returns: + torch.Tensor of similarities (M, N) where M=len(queries), N=len(candidates) + """ + query_tensors = torch.stack([q.to_torch(self.device) for q in queries]) + candidate_tensors = torch.stack([c.to_torch(self.device) for c in candidates]) + return query_tensors @ candidate_tensors.T + + def query(self, query_emb: E, candidates: list[E], top_k: int = 5) -> list[tuple[int, float]]: + """ + Find top-k most similar candidates to query (GPU accelerated). + + Args: + query_emb: Query embedding + candidates: List of candidate embeddings + top_k: Number of top results to return + + Returns: + List of (index, similarity) tuples sorted by similarity (descending) + """ + similarities = self.compare_one_to_many(query_emb, candidates) + top_values, top_indices = similarities.topk(k=min(top_k, len(candidates))) + return [(idx.item(), val.item()) for idx, val in zip(top_indices, top_values)] + + def warmup(self) -> None: + """Optional warmup method to pre-load model.""" + pass diff --git a/dimos/models/embedding/clip.py b/dimos/models/embedding/clip.py new file mode 100644 index 0000000000..e751e9ee33 --- /dev/null +++ b/dimos/models/embedding/clip.py @@ -0,0 +1,123 @@ +# Copyright 2025 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. + +import torch +import torch.nn.functional as F +from PIL import Image as PILImage +from transformers import CLIPModel as HFCLIPModel +from transformers import CLIPProcessor + +from dimos.models.embedding.base import Embedding, EmbeddingModel +from dimos.msgs.sensor_msgs import Image + +_CUDA_INITIALIZED = False + + +class CLIPEmbedding(Embedding): ... + + +class CLIPModel(EmbeddingModel[CLIPEmbedding]): + """CLIP embedding model for vision-language re-identification.""" + + def __init__( + self, + model_name: str = "openai/clip-vit-base-patch32", + device: str | None = None, + normalize: bool = False, + ): + """ + Initialize CLIP model. + + Args: + model_name: HuggingFace model name (e.g., "openai/clip-vit-base-patch32") + device: Device to run on (cuda/cpu), auto-detects if None + normalize: Whether to L2 normalize embeddings + """ + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.normalize = normalize + + # Load model and processor + self.model = HFCLIPModel.from_pretrained(model_name).eval().to(self.device) + self.processor = CLIPProcessor.from_pretrained(model_name) + + def embed(self, *images: Image) -> CLIPEmbedding | list[CLIPEmbedding]: + """Embed one or more images. + + 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] + + # Process images + with torch.inference_mode(): + inputs = self.processor(images=pil_images, return_tensors="pt").to(self.device) + image_features = self.model.get_image_features(**inputs) + + if self.normalize: + image_features = F.normalize(image_features, dim=-1) + + # Create embeddings (keep as torch.Tensor on device) + embeddings = [] + for i, feat in enumerate(image_features): + timestamp = images[i].ts + embeddings.append(CLIPEmbedding(vector=feat, timestamp=timestamp)) + + return embeddings[0] if len(images) == 1 else embeddings + + def embed_text(self, *texts: str) -> CLIPEmbedding | list[CLIPEmbedding]: + """Embed one or more text strings. + + Returns embeddings as torch.Tensor on device for efficient GPU comparisons. + """ + with torch.inference_mode(): + inputs = self.processor(text=list(texts), return_tensors="pt", padding=True).to( + self.device + ) + text_features = self.model.get_text_features(**inputs) + + if self.normalize: + text_features = F.normalize(text_features, dim=-1) + + # Create embeddings (keep as torch.Tensor on device) + embeddings = [] + for feat in text_features: + embeddings.append(CLIPEmbedding(vector=feat)) + + return embeddings[0] if len(texts) == 1 else embeddings + + def warmup(self) -> None: + """Warmup the model with a dummy forward pass.""" + # WORKAROUND: HuggingFace CLIP fails with CUBLAS_STATUS_ALLOC_FAILED when it's + # the first model to use CUDA. Initialize CUDA context with a dummy operation. + # This only needs to happen once per process. + global _CUDA_INITIALIZED + if self.device == "cuda" and not _CUDA_INITIALIZED: + try: + # Initialize CUDA with a small matmul operation to setup cuBLAS properly + _ = torch.zeros(1, 1, device="cuda") @ torch.zeros(1, 1, device="cuda") + torch.cuda.synchronize() + _CUDA_INITIALIZED = True + except Exception: + # If initialization fails, continue anyway - the warmup might still work + pass + + dummy_image = torch.randn(1, 3, 224, 224).to(self.device) + dummy_text_inputs = self.processor(text=["warmup"], return_tensors="pt", padding=True).to( + self.device + ) + + with torch.inference_mode(): + # Use pixel_values directly for image warmup + self.model.get_image_features(pixel_values=dummy_image) + self.model.get_text_features(**dummy_text_inputs) diff --git a/dimos/models/embedding/embedding_models_disabled_tests.py b/dimos/models/embedding/embedding_models_disabled_tests.py new file mode 100644 index 0000000000..52e9fd08af --- /dev/null +++ b/dimos/models/embedding/embedding_models_disabled_tests.py @@ -0,0 +1,404 @@ +# Copyright 2025 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. + +import numpy as np +import pytest + +from dimos.msgs.sensor_msgs import Image +from dimos.utils.data import get_data + + +@pytest.fixture(scope="session", params=["clip", "mobileclip", "treid"]) +def embedding_model(request): + """Load embedding model once for all tests. Parametrized for different models.""" + if request.param == "mobileclip": + from dimos.models.embedding.mobileclip import MobileCLIPModel + + model_path = get_data("models_mobileclip") / "mobileclip2_s0.pt" + model = MobileCLIPModel(model_name="MobileCLIP2-S0", model_path=model_path) + elif request.param == "clip": + from dimos.models.embedding.clip import CLIPModel + + model = CLIPModel(model_name="openai/clip-vit-base-patch32") + elif request.param == "treid": + from dimos.models.embedding.treid import TorchReIDModel + + model = TorchReIDModel(model_name="osnet_x1_0") + else: + raise ValueError(f"Unknown model: {request.param}") + + model.warmup() + return model + + +@pytest.fixture(scope="session") +def test_image(): + """Load test image.""" + return Image.from_file(get_data("cafe.jpg")).to_rgb() + + +@pytest.mark.heavy +def test_single_image_embedding(embedding_model, test_image): + """Test embedding a single image.""" + embedding = embedding_model.embed(test_image) + + # Embedding should be torch.Tensor on device + import torch + + assert isinstance(embedding.vector, torch.Tensor), "Embedding should be torch.Tensor" + assert embedding.vector.device.type in ["cuda", "cpu"], "Should be on valid device" + + # Test conversion to numpy + vector_np = embedding.to_numpy() + print(f"\nEmbedding shape: {vector_np.shape}") + print(f"Embedding dtype: {vector_np.dtype}") + print(f"Embedding norm: {np.linalg.norm(vector_np):.4f}") + + assert vector_np.shape[0] > 0, "Embedding should have features" + assert np.isfinite(vector_np).all(), "Embedding should contain finite values" + + # Check L2 normalization + norm = np.linalg.norm(vector_np) + assert abs(norm - 1.0) < 0.01, f"Embedding should be L2 normalized, got norm={norm}" + + +@pytest.mark.heavy +def test_batch_image_embedding(embedding_model, test_image): + """Test embedding multiple images at once.""" + embeddings = embedding_model.embed(test_image, test_image, test_image) + + assert isinstance(embeddings, list), "Batch embedding should return list" + assert len(embeddings) == 3, "Should return 3 embeddings" + + # Check all embeddings are similar (same image) + sim_01 = embeddings[0] @ embeddings[1] + sim_02 = embeddings[0] @ embeddings[2] + + print(f"\nSimilarity between same images: {sim_01:.6f}, {sim_02:.6f}") + + assert sim_01 > 0.99, f"Same image embeddings should be very similar, got {sim_01}" + assert sim_02 > 0.99, f"Same image embeddings should be very similar, got {sim_02}" + + +@pytest.mark.heavy +def test_single_text_embedding(embedding_model): + """Test embedding a single text string.""" + import torch + + if not hasattr(embedding_model, "embed_text"): + pytest.skip("Model does not support text embeddings") + + embedding = embedding_model.embed_text("a cafe") + + # Should be torch.Tensor + assert isinstance(embedding.vector, torch.Tensor), "Text embedding should be torch.Tensor" + + vector_np = embedding.to_numpy() + print(f"\nText embedding shape: {vector_np.shape}") + print(f"Text embedding norm: {np.linalg.norm(vector_np):.4f}") + + assert vector_np.shape[0] > 0, "Text embedding should have features" + assert np.isfinite(vector_np).all(), "Text embedding should contain finite values" + + # Check L2 normalization + norm = np.linalg.norm(vector_np) + assert abs(norm - 1.0) < 0.01, f"Text embedding should be L2 normalized, got norm={norm}" + + +@pytest.mark.heavy +def test_batch_text_embedding(embedding_model): + """Test embedding multiple text strings at once.""" + import torch + + if not hasattr(embedding_model, "embed_text"): + pytest.skip("Model does not support text embeddings") + + embeddings = embedding_model.embed_text("a cafe", "a person", "a dog") + + assert isinstance(embeddings, list), "Batch text embedding should return list" + assert len(embeddings) == 3, "Should return 3 text embeddings" + + # All should be torch.Tensor and normalized + for i, emb in enumerate(embeddings): + assert isinstance(emb.vector, torch.Tensor), f"Embedding {i} should be torch.Tensor" + norm = np.linalg.norm(emb.to_numpy()) + assert abs(norm - 1.0) < 0.01, f"Text embedding {i} should be L2 normalized" + + +@pytest.mark.heavy +def test_text_image_similarity(embedding_model, test_image): + """Test cross-modal text-image similarity using @ operator.""" + if not hasattr(embedding_model, "embed_text"): + pytest.skip("Model does not support text embeddings") + + img_embedding = embedding_model.embed(test_image) + + # Embed text queries + queries = ["a cafe", "a person", "a car", "a dog", "potato", "food"] + text_embeddings = embedding_model.embed_text(*queries) + + # Compute similarities using @ operator + similarities = {} + for query, text_emb in zip(queries, text_embeddings): + similarity = img_embedding @ text_emb + similarities[query] = similarity + print(f"\n'{query}': {similarity:.4f}") + + # Cafe image should match "a cafe" better than "a dog" + assert similarities["a cafe"] > similarities["a dog"], "Should recognize cafe scene" + assert similarities["a person"] > similarities["a car"], "Should detect people in cafe" + + +@pytest.mark.heavy +def test_cosine_distance(embedding_model, test_image): + """Test cosine distance computation (1 - similarity).""" + emb1 = embedding_model.embed(test_image) + emb2 = embedding_model.embed(test_image) + + # Similarity using @ operator + similarity = emb1 @ emb2 + + # Distance is 1 - similarity + distance = 1.0 - similarity + + print(f"\nSimilarity (same image): {similarity:.6f}") + print(f"Distance (same image): {distance:.6f}") + + assert similarity > 0.99, f"Same image should have high similarity, got {similarity}" + assert distance < 0.01, f"Same image should have low distance, got {distance}" + + +@pytest.mark.heavy +def test_query_functionality(embedding_model, test_image): + """Test query method for top-k retrieval.""" + if not hasattr(embedding_model, "embed_text"): + pytest.skip("Model does not support text embeddings") + + # Create a query and some candidates + query_text = embedding_model.embed_text("a cafe") + + # Create candidate embeddings + candidate_texts = ["a cafe", "a restaurant", "a person", "a dog", "a car"] + candidates = embedding_model.embed_text(*candidate_texts) + + # Query for top-3 + results = embedding_model.query(query_text, candidates, top_k=3) + + print("\nTop-3 results:") + for idx, sim in results: + print(f" {candidate_texts[idx]}: {sim:.4f}") + + assert len(results) == 3, "Should return top-3 results" + assert results[0][0] == 0, "Top match should be 'a cafe' itself" + assert results[0][1] > results[1][1], "Results should be sorted by similarity" + assert results[1][1] > results[2][1], "Results should be sorted by similarity" + + +@pytest.mark.heavy +def test_embedding_operator(embedding_model, test_image): + """Test that @ operator works on embeddings.""" + emb1 = embedding_model.embed(test_image) + emb2 = embedding_model.embed(test_image) + + # Use @ operator + similarity = emb1 @ emb2 + + assert isinstance(similarity, float), "@ operator should return float" + assert 0.0 <= similarity <= 1.0, "Cosine similarity should be in [0, 1]" + assert similarity > 0.99, "Same image should have similarity near 1.0" + + +@pytest.mark.heavy +def test_warmup(embedding_model): + """Test that warmup runs without error.""" + # Warmup is already called in fixture, but test it explicitly + embedding_model.warmup() + # Just verify no exceptions raised + assert True + + +@pytest.mark.heavy +def test_compare_one_to_many(embedding_model, test_image): + """Test GPU-accelerated one-to-many comparison.""" + import torch + + # Create query and gallery + query_emb = embedding_model.embed(test_image) + gallery_embs = embedding_model.embed(test_image, test_image, test_image) + + # Compare on GPU + similarities = embedding_model.compare_one_to_many(query_emb, gallery_embs) + + print(f"\nOne-to-many similarities: {similarities}") + + # Should return torch.Tensor + assert isinstance(similarities, torch.Tensor), "Should return torch.Tensor" + assert similarities.shape == (3,), "Should have 3 similarities" + assert similarities.device.type in ["cuda", "cpu"], "Should be on device" + + # All should be ~1.0 (same image) + similarities_np = similarities.cpu().numpy() + assert np.all(similarities_np > 0.99), "Same images should have similarity ~1.0" + + +@pytest.mark.heavy +def test_compare_many_to_many(embedding_model): + """Test GPU-accelerated many-to-many comparison.""" + import torch + + if not hasattr(embedding_model, "embed_text"): + pytest.skip("Model does not support text embeddings") + + # Create queries and candidates + queries = embedding_model.embed_text("a cafe", "a person") + candidates = embedding_model.embed_text("a cafe", "a restaurant", "a dog") + + # Compare on GPU + similarities = embedding_model.compare_many_to_many(queries, candidates) + + print(f"\nMany-to-many similarities:\n{similarities}") + + # Should return torch.Tensor + assert isinstance(similarities, torch.Tensor), "Should return torch.Tensor" + assert similarities.shape == (2, 3), "Should be (2, 3) similarity matrix" + assert similarities.device.type in ["cuda", "cpu"], "Should be on device" + + # First query should match first candidate best + similarities_np = similarities.cpu().numpy() + assert similarities_np[0, 0] > similarities_np[0, 2], "Cafe should match cafe better than dog" + + +@pytest.mark.heavy +def test_gpu_query_performance(embedding_model, test_image): + """Test that query method uses GPU acceleration.""" + # Create a larger gallery + gallery_size = 20 + gallery_images = [test_image] * gallery_size + gallery_embs = embedding_model.embed(*gallery_images) + + query_emb = embedding_model.embed(test_image) + + # Query should use GPU-accelerated comparison + results = embedding_model.query(query_emb, gallery_embs, top_k=5) + + print(f"\nTop-5 results from gallery of {gallery_size}") + for idx, sim in results: + print(f" Index {idx}: {sim:.4f}") + + assert len(results) == 5, "Should return top-5 results" + # All should be high similarity (same image, allow some variation for image preprocessing) + for idx, sim in results: + assert sim > 0.90, f"Same images should have high similarity, got {sim}" + + +@pytest.mark.heavy +def test_embedding_performance(embedding_model): + """Measure embedding performance over multiple real video frames.""" + import time + + from dimos.utils.testing import TimedSensorReplay + + # Load actual video frames + data_dir = "unitree_go2_lidar_corrected" + get_data(data_dir) + + video_replay = TimedSensorReplay(f"{data_dir}/video") + + # Collect 10 real frames from the video + test_images = [] + for ts, frame in video_replay.iterate_ts(duration=1.0): + test_images.append(frame.to_rgb()) + if len(test_images) >= 10: + break + + if len(test_images) < 10: + pytest.skip(f"Not enough video frames found (got {len(test_images)})") + + # Measure single image embedding time + times = [] + for img in test_images: + start = time.perf_counter() + _ = embedding_model.embed(img) + end = time.perf_counter() + elapsed_ms = (end - start) * 1000 + times.append(elapsed_ms) + + # Calculate statistics + avg_time = sum(times) / len(times) + min_time = min(times) + max_time = max(times) + std_time = (sum((t - avg_time) ** 2 for t in times) / len(times)) ** 0.5 + + print("\n" + "=" * 60) + print("Embedding Performance Statistics:") + print("=" * 60) + print(f"Number of images: {len(test_images)}") + print(f"Average time: {avg_time:.2f} ms") + print(f"Min time: {min_time:.2f} ms") + print(f"Max time: {max_time:.2f} ms") + print(f"Std dev: {std_time:.2f} ms") + print(f"Throughput: {1000 / avg_time:.1f} images/sec") + print("=" * 60) + + # Also test batch embedding performance + start = time.perf_counter() + batch_embeddings = embedding_model.embed(*test_images) + end = time.perf_counter() + batch_time = (end - start) * 1000 + batch_per_image = batch_time / len(test_images) + + print("\nBatch Embedding Performance:") + print(f"Total batch time: {batch_time:.2f} ms") + print(f"Time per image (batched): {batch_per_image:.2f} ms") + print(f"Batch throughput: {1000 / batch_per_image:.1f} images/sec") + print(f"Speedup vs single: {avg_time / batch_per_image:.2f}x") + print("=" * 60) + + # Verify embeddings are valid + assert len(batch_embeddings) == len(test_images) + assert all(e.vector is not None for e in batch_embeddings) + + # Sanity check: verify embeddings are meaningful by testing text-image similarity + # Skip for models that don't support text embeddings + if hasattr(embedding_model, "embed_text"): + print("\n" + "=" * 60) + print("Sanity Check: Text-Image Similarity on First Frame") + print("=" * 60) + first_frame_emb = batch_embeddings[0] + + # Test common object/scene queries + test_queries = [ + "indoor scene", + "outdoor scene", + "a person", + "a dog", + "a robot", + "grass and trees", + "furniture", + "a car", + ] + + text_embeddings = embedding_model.embed_text(*test_queries) + similarities = [] + for query, text_emb in zip(test_queries, text_embeddings): + sim = first_frame_emb @ text_emb + similarities.append((query, sim)) + + # Sort by similarity + similarities.sort(key=lambda x: x[1], reverse=True) + + print("Top matching concepts:") + for query, sim in similarities[:5]: + print(f" '{query}': {sim:.4f}") + print("=" * 60) diff --git a/dimos/models/embedding/mobileclip.py b/dimos/models/embedding/mobileclip.py new file mode 100644 index 0000000000..c0295a78ef --- /dev/null +++ b/dimos/models/embedding/mobileclip.py @@ -0,0 +1,112 @@ +# Copyright 2025 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 pathlib import Path + +import open_clip +import torch +import torch.nn.functional as F +from PIL import Image as PILImage + +from dimos.models.embedding.base import Embedding, EmbeddingModel +from dimos.msgs.sensor_msgs import Image + + +class MobileCLIPEmbedding(Embedding): ... + + +class MobileCLIPModel(EmbeddingModel[MobileCLIPEmbedding]): + """MobileCLIP embedding model for vision-language re-identification.""" + + def __init__( + self, + model_name: str = "MobileCLIP2-S4", + model_path: Path | str | None = None, + device: str | None = None, + normalize: bool = True, + ): + """ + Initialize MobileCLIP model. + + Args: + model_name: Name of the model architecture + model_path: Path to pretrained weights + device: Device to run on (cuda/cpu), auto-detects if None + normalize: Whether to L2 normalize embeddings + """ + if not OPEN_CLIP_AVAILABLE: + raise ImportError( + "open_clip is required for MobileCLIPModel. " + "Install it with: pip install open-clip-torch" + ) + + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.normalize = normalize + + # Load model + pretrained = str(model_path) if model_path else None + self.model, _, self.preprocess = open_clip.create_model_and_transforms( + model_name, pretrained=pretrained + ) + self.tokenizer = open_clip.get_tokenizer(model_name) + self.model = self.model.eval().to(self.device) + + def embed(self, *images: Image) -> MobileCLIPEmbedding | list[MobileCLIPEmbedding]: + """Embed one or more images. + + 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] + + # Preprocess and batch + with torch.inference_mode(): + batch = torch.stack([self.preprocess(img) for img in pil_images]).to(self.device) + feats = self.model.encode_image(batch) + if self.normalize: + feats = F.normalize(feats, dim=-1) + + # Create embeddings (keep as torch.Tensor on device) + embeddings = [] + for i, feat in enumerate(feats): + timestamp = images[i].ts + embeddings.append(MobileCLIPEmbedding(vector=feat, timestamp=timestamp)) + + return embeddings[0] if len(images) == 1 else embeddings + + def embed_text(self, *texts: str) -> MobileCLIPEmbedding | list[MobileCLIPEmbedding]: + """Embed one or more text strings. + + Returns embeddings as torch.Tensor on device for efficient GPU comparisons. + """ + with torch.inference_mode(): + text_tokens = self.tokenizer(list(texts)).to(self.device) + feats = self.model.encode_text(text_tokens) + if self.normalize: + feats = F.normalize(feats, dim=-1) + + # Create embeddings (keep as torch.Tensor on device) + embeddings = [] + for feat in feats: + embeddings.append(MobileCLIPEmbedding(vector=feat)) + + return embeddings[0] if len(texts) == 1 else embeddings + + def warmup(self) -> None: + """Warmup the model with a dummy forward pass.""" + dummy_image = torch.randn(1, 3, 224, 224).to(self.device) + dummy_text = self.tokenizer(["warmup"]).to(self.device) + with torch.inference_mode(): + self.model.encode_image(dummy_image) + self.model.encode_text(dummy_text) diff --git a/dimos/models/embedding/treid.py b/dimos/models/embedding/treid.py new file mode 100644 index 0000000000..bdd00627a0 --- /dev/null +++ b/dimos/models/embedding/treid.py @@ -0,0 +1,125 @@ +# Copyright 2025 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 pathlib import Path + +import torch +import torch.nn.functional as F +from torchreid import utils as torchreid_utils + +from dimos.models.embedding.base import Embedding, EmbeddingModel +from dimos.msgs.sensor_msgs import Image + +_CUDA_INITIALIZED = False + + +class TorchReIDEmbedding(Embedding): ... + + +class TorchReIDModel(EmbeddingModel[TorchReIDEmbedding]): + """TorchReID embedding model for person re-identification.""" + + def __init__( + self, + model_name: str = "se_resnext101_32x4d", + model_path: Path | str | None = None, + device: str | None = None, + normalize: bool = False, + ): + """ + Initialize TorchReID model. + + Args: + model_name: Name of the model architecture (e.g., "osnet_x1_0", "osnet_x0_75") + model_path: Path to pretrained weights (.pth.tar file) + device: Device to run on (cuda/cpu), auto-detects if None + normalize: Whether to L2 normalize embeddings + """ + if not TORCHREID_AVAILABLE: + raise ImportError( + "torchreid is required for TorchReIDModel. Install it with: pip install torchreid" + ) + + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.normalize = normalize + + # Load model using torchreid's FeatureExtractor + model_path_str = str(model_path) if model_path else "" + self.extractor = torchreid_utils.FeatureExtractor( + model_name=model_name, + model_path=model_path_str, + device=self.device, + ) + + def embed(self, *images: Image) -> TorchReIDEmbedding | list[TorchReIDEmbedding]: + """Embed one or more images. + + Returns embeddings as torch.Tensor on device for efficient GPU comparisons. + """ + # Convert to numpy arrays - torchreid expects numpy arrays or file paths + np_images = [img.to_opencv() for img in images] + + # Extract features + with torch.inference_mode(): + features = self.extractor(np_images) + + # torchreid may return either numpy array or torch tensor depending on configuration + if isinstance(features, torch.Tensor): + features_tensor = features.to(self.device) + else: + features_tensor = torch.from_numpy(features).to(self.device) + + if self.normalize: + features_tensor = F.normalize(features_tensor, dim=-1) + + # Create embeddings (keep as torch.Tensor on device) + embeddings = [] + for i, feat in enumerate(features_tensor): + timestamp = images[i].ts + embeddings.append(TorchReIDEmbedding(vector=feat, timestamp=timestamp)) + + return embeddings[0] if len(images) == 1 else embeddings + + def embed_text(self, *texts: str) -> TorchReIDEmbedding | list[TorchReIDEmbedding]: + """Text embedding not supported for ReID models. + + TorchReID models are vision-only person re-identification models + and do not support text embeddings. + """ + raise NotImplementedError( + "TorchReID models are vision-only and do not support text embeddings. " + "Use CLIP or MobileCLIP for text-image similarity." + ) + + def warmup(self) -> None: + """Warmup the model with a dummy forward pass.""" + # WORKAROUND: TorchReID can fail with CUBLAS errors when it's the first model to use CUDA. + # Initialize CUDA context with a dummy operation. This only needs to happen once per process. + global _CUDA_INITIALIZED + if self.device == "cuda" and not _CUDA_INITIALIZED: + try: + # Initialize CUDA with a small matmul operation to setup cuBLAS properly + _ = torch.zeros(1, 1, device="cuda") @ torch.zeros(1, 1, device="cuda") + torch.cuda.synchronize() + _CUDA_INITIALIZED = True + except Exception: + # If initialization fails, continue anyway - the warmup might still work + pass + + # Create a dummy 256x128 image (typical person ReID input size) as numpy array + import numpy as np + + dummy_image = np.random.randint(0, 256, (256, 128, 3), dtype=np.uint8) + with torch.inference_mode(): + _ = self.extractor([dummy_image]) diff --git a/dimos/models/vl/base.py b/dimos/models/vl/base.py index faab96363d..cde41bd8fc 100644 --- a/dimos/models/vl/base.py +++ b/dimos/models/vl/base.py @@ -1,10 +1,106 @@ +import json +import logging from abc import ABC, abstractmethod -import numpy as np - from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.type import Detection2DBBox, ImageDetections2D +from dimos.utils.data import get_data +from dimos.utils.decorators import retry +from dimos.utils.llm_utils import extract_json + +logger = logging.getLogger(__name__) + + +def vlm_detection_to_detection2d( + vlm_detection: list, track_id: int, image: Image +) -> Detection2DBBox | None: + """Convert a single VLM detection [label, x1, y1, x2, y2] to Detection2DBBox. + + Args: + vlm_detection: Single detection list containing [label, x1, y1, x2, y2] + track_id: Track ID to assign to this detection + image: Source image for the detection + + Returns: + Detection2DBBox instance or None if invalid + """ + # Validate list structure + if not isinstance(vlm_detection, list): + logger.debug(f"VLM detection is not a list: {type(vlm_detection)}") + return None + + if len(vlm_detection) != 5: + logger.debug( + f"Invalid VLM detection length: {len(vlm_detection)}, expected 5. Got: {vlm_detection}" + ) + return None + + # Extract label + name = str(vlm_detection[0]) + + # Validate and convert coordinates + try: + coords = [float(x) for x in vlm_detection[1:]] + except (ValueError, TypeError) as e: + logger.debug(f"Invalid VLM detection coordinates: {vlm_detection[1:]}. Error: {e}") + return None + + bbox = tuple(coords) + + # Use -1 for class_id since VLM doesn't provide it + # confidence defaults to 1.0 for VLM + return Detection2DBBox( + bbox=bbox, + track_id=track_id, + class_id=-1, + confidence=1.0, + name=name, + ts=image.ts, + image=image, + ) class VlModel(ABC): @abstractmethod - def query(self, image: Image | np.ndarray, query: str) -> str: ... + def query(self, image: Image, query: str, **kwargs) -> str: ... + + def warmup(self) -> None: + try: + image = Image.from_file(get_data("cafe-smol.jpg")).to_rgb() + self._model.detect(image, "person", settings={"max_objects": 1}) + except Exception: + pass + + # requery once if JSON parsing fails + @retry(max_retries=2, on_exception=json.JSONDecodeError, delay=0.0) + def query_json(self, image: Image, query: str) -> dict: + response = self.query(image, query) + return extract_json(response) + + def query_detections(self, image: Image, query: str, **kwargs) -> ImageDetections2D: + full_query = f"""show me bounding boxes in pixels for this query: `{query}` + + format should be: + `[ + [label, x1, y1, x2, y2] + ... + ]` + + (etc, multiple matches are possible) + + If there's no match return `[]`. Label is whatever you think is appropriate + Only respond with the coordinates, no other text.""" + + image_detections = ImageDetections2D(image) + + try: + detection_tuples = self.query_json(image, full_query) + except Exception: + return image_detections + + for track_id, detection_tuple in enumerate(detection_tuples): + detection2d = vlm_detection_to_detection2d(detection_tuple, track_id, image) + if detection2d is not None and detection2d.is_valid(): + image_detections.detections.append(detection2d) + + return image_detections diff --git a/dimos/models/vl/moondream.py b/dimos/models/vl/moondream.py new file mode 100644 index 0000000000..a3b9f5fcca --- /dev/null +++ b/dimos/models/vl/moondream.py @@ -0,0 +1,114 @@ +import warnings +from functools import cached_property +from typing import Optional + +import numpy as np +import torch +from PIL import Image as PILImage +from transformers import AutoModelForCausalLM + +from dimos.models.vl.base import VlModel +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.type import Detection2DBBox, ImageDetections2D + + +class MoondreamVlModel(VlModel): + _model_name: str + _device: str + _dtype: torch.dtype + + def __init__( + self, + model_name: str = "vikhyatk/moondream2", + device: Optional[str] = None, + dtype: torch.dtype = torch.bfloat16, + ): + self._model_name = model_name + self._device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self._dtype = dtype + + @cached_property + def _model(self) -> AutoModelForCausalLM: + model = AutoModelForCausalLM.from_pretrained( + self._model_name, + trust_remote_code=True, + torch_dtype=self._dtype, + ) + model = model.to(self._device) + model.compile() + + return model + + def query(self, image: Image | np.ndarray, query: str, **kwargs) -> str: + if isinstance(image, np.ndarray): + warnings.warn( + "MoondreamVlModel.query should receive standard dimos Image type, not a numpy array", + DeprecationWarning, + stacklevel=2, + ) + image = Image.from_numpy(image) + + # Convert dimos Image to PIL Image + # dimos Image stores data in RGB/BGR format, convert to RGB for PIL + rgb_image = image.to_rgb() + pil_image = PILImage.fromarray(rgb_image.data) + + # Query the model + result = self._model.query(image=pil_image, question=query, reasoning=False) + + # Handle both dict and string responses + if isinstance(result, dict): + return result.get("answer", str(result)) + + return str(result) + + def query_detections(self, image: Image, query: str, **kwargs) -> ImageDetections2D: + """Detect objects using Moondream's native detect method. + + Args: + image: Input image + query: Object query (e.g., "person", "car") + max_objects: Maximum number of objects to detect + + Returns: + ImageDetections2D containing detected bounding boxes + """ + pil_image = PILImage.fromarray(image.data) + + settings = {"max_objects": kwargs.get("max_objects", 5)} + result = self._model.detect(pil_image, query, settings=settings) + + # Convert to ImageDetections2D + image_detections = ImageDetections2D(image) + + # Get image dimensions for converting normalized coords to pixels + height, width = image.height, image.width + + for track_id, obj in enumerate(result.get("objects", [])): + # Convert normalized coordinates (0-1) to pixel coordinates + x_min_norm = obj["x_min"] + y_min_norm = obj["y_min"] + x_max_norm = obj["x_max"] + y_max_norm = obj["y_max"] + + x1 = x_min_norm * width + y1 = y_min_norm * height + x2 = x_max_norm * width + y2 = y_max_norm * height + + bbox = (x1, y1, x2, y2) + + detection = Detection2DBBox( + bbox=bbox, + track_id=track_id, + class_id=-1, # Moondream doesn't provide class IDs + confidence=1.0, # Moondream doesn't provide confidence scores + name=query, # Use the query as the object name + ts=image.ts, + image=image, + ) + + if detection.is_valid(): + image_detections.detections.append(detection) + + return image_detections diff --git a/dimos/models/vl/qwen.py b/dimos/models/vl/qwen.py index 05ad4715c5..c34f6f7964 100644 --- a/dimos/models/vl/qwen.py +++ b/dimos/models/vl/qwen.py @@ -1,4 +1,5 @@ import os +from functools import cached_property from typing import Optional import numpy as np @@ -9,19 +10,22 @@ class QwenVlModel(VlModel): - _client: OpenAI _model_name: str + _api_key: Optional[str] def __init__(self, api_key: Optional[str] = None, model_name: str = "qwen2.5-vl-72b-instruct"): self._model_name = model_name + self._api_key = api_key - api_key = api_key or os.getenv("ALIBABA_API_KEY") + @cached_property + def _client(self) -> OpenAI: + api_key = self._api_key or os.getenv("ALIBABA_API_KEY") if not api_key: raise ValueError( "Alibaba API key must be provided or set in ALIBABA_API_KEY environment variable" ) - self._client = OpenAI( + return OpenAI( base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1", api_key=api_key, ) diff --git a/dimos/models/vl/test_base.py b/dimos/models/vl/test_base.py new file mode 100644 index 0000000000..302a588721 --- /dev/null +++ b/dimos/models/vl/test_base.py @@ -0,0 +1,105 @@ +import os +from unittest.mock import MagicMock + +import pytest + +from dimos.models.vl.qwen import QwenVlModel +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.type import ImageDetections2D +from dimos.utils.data import get_data + +# Captured actual response from Qwen API for cafe.jpg with query "humans" +# Added garbage around JSON to ensure we are robustly extracting it +MOCK_QWEN_RESPONSE = """ + Locating humans for you šŸ˜ŠšŸ˜Š + + [ + ["humans", 76, 368, 219, 580], + ["humans", 354, 372, 512, 525], + ["humans", 409, 370, 615, 748], + ["humans", 628, 350, 762, 528], + ["humans", 785, 323, 960, 650] + ] + + Here is some trash at the end of the response :) + Let me know if you need anything else šŸ˜€šŸ˜Š + """ + + +def test_query_detections_mocked(): + """Test query_detections with mocked API response (no API key required).""" + # Load test image + image = Image.from_file(get_data("cafe.jpg")) + + # Create model and mock the query method + model = QwenVlModel() + model.query = MagicMock(return_value=MOCK_QWEN_RESPONSE) + + # Query for humans in the image + query = "humans" + detections = model.query_detections(image, query) + + # Verify the return type + assert isinstance(detections, ImageDetections2D) + + # Should have 5 detections based on our mock data + assert len(detections.detections) == 5, ( + f"Expected 5 detections, got {len(detections.detections)}" + ) + + # Verify each detection + img_height, img_width = image.shape[:2] + + for i, detection in enumerate(detections.detections): + # Verify attributes + assert detection.name == "humans" + assert detection.confidence == 1.0 + assert detection.class_id == -1 # VLM detections use -1 for class_id + assert detection.track_id == i + assert len(detection.bbox) == 4 + + assert detection.is_valid() + + # Verify bbox coordinates are valid (out-of-bounds detections are discarded) + x1, y1, x2, y2 = detection.bbox + assert x2 > x1, f"Detection {i}: Invalid x coordinates: x1={x1}, x2={x2}" + assert y2 > y1, f"Detection {i}: Invalid y coordinates: y1={y1}, y2={y2}" + + # Check bounds (out-of-bounds detections would have been discarded) + assert 0 <= x1 <= img_width, f"Detection {i}: x1={x1} out of bounds" + assert 0 <= x2 <= img_width, f"Detection {i}: x2={x2} out of bounds" + assert 0 <= y1 <= img_height, f"Detection {i}: y1={y1} out of bounds" + assert 0 <= y2 <= img_height, f"Detection {i}: y2={y2} out of bounds" + + print(f"āœ“ Successfully processed {len(detections.detections)} mocked detections") + + +@pytest.mark.tool +@pytest.mark.skipif(not os.getenv("ALIBABA_API_KEY"), reason="ALIBABA_API_KEY not set") +def test_query_detections_real(): + """Test query_detections with real API calls (requires API key).""" + # Load test image + image = Image.from_file(get_data("cafe.jpg")) + + # Initialize the model (will use real API) + model = QwenVlModel() + + # Query for humans in the image + query = "humans" + detections = model.query_detections(image, query) + + assert isinstance(detections, ImageDetections2D) + print(detections) + + # Check that detections were found + if detections.detections: + for detection in detections.detections: + # Verify each detection has expected attributes + assert detection.bbox is not None + assert len(detection.bbox) == 4 + assert detection.name + assert detection.confidence == 1.0 + assert detection.class_id == -1 # VLM detections use -1 for class_id + assert detection.is_valid() + + print(f"Found {len(detections.detections)} detections for query '{query}'") diff --git a/dimos/models/vl/test_models.py b/dimos/models/vl/test_models.py new file mode 100644 index 0000000000..3871626ae1 --- /dev/null +++ b/dimos/models/vl/test_models.py @@ -0,0 +1,89 @@ +import time + +import pytest +from dimos_lcm.foxglove_msgs.ImageAnnotations import ImageAnnotations + +from dimos.core import LCMTransport +from dimos.models.vl.base import VlModel +from dimos.models.vl.moondream import MoondreamVlModel +from dimos.models.vl.qwen import QwenVlModel +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.detectors.yolo import Yolo2DDetector +from dimos.perception.detection.type import ImageDetections2D +from dimos.utils.data import get_data + + +@pytest.mark.parametrize( + "model_class,model_name", + [ + (MoondreamVlModel, "Moondream"), + (QwenVlModel, "Qwen"), + ], + ids=["moondream", "qwen"], +) +@pytest.mark.gpu +def test_vlm(model_class, model_name): + image = Image.from_file(get_data("cafe.jpg")).to_rgb() + + print(f"Testing {model_name}") + + # Initialize model + print(f"Loading {model_name} model...") + model: VlModel = model_class() + model.warmup() + + queries = [ + "glasses", + "blue shirt", + "bulb", + "cigarette", + "reflection of a car", + "knee", + "flowers on the left table", + "shoes", + "leftmost persons ear", + "rightmost arm", + ] + + all_detections = ImageDetections2D(image) + query_times = [] + + # # First, run YOLO detection + # print("\nRunning YOLO detection...") + # yolo_detector = Yolo2DDetector() + # yolo_detections = yolo_detector.process_image(image) + # print(f" YOLO found {len(yolo_detections.detections)} objects") + # all_detections.detections.extend(yolo_detections.detections) + # annotations_transport.publish(all_detections.to_foxglove_annotations()) + + # Publish to LCM with model-specific channel names + annotations_transport: LCMTransport[ImageAnnotations] = LCMTransport( + "/annotations", ImageAnnotations + ) + + image_transport: LCMTransport[Image] = LCMTransport("/image", Image) + + image_transport.publish(image) + + # Then run VLM queries + for query in queries: + print(f"\nQuerying for: {query}") + start_time = time.time() + detections = model.query_detections(image, query, max_objects=5) + query_time = time.time() - start_time + query_times.append(query_time) + + print(f" Found {len(detections)} detections in {query_time:.3f}s") + all_detections.detections.extend(detections.detections) + annotations_transport.publish(all_detections.to_foxglove_annotations()) + + avg_time = sum(query_times) / len(query_times) if query_times else 0 + print(f"\n{model_name} Results:") + print(f" Average query time: {avg_time:.3f}s") + print(f" Total detections: {len(all_detections)}") + print(all_detections) + + annotations_transport.publish(all_detections.to_foxglove_annotations()) + + annotations_transport.lcm.stop() + image_transport.lcm.stop() diff --git a/dimos/msgs/foxglove_msgs/Color.py b/dimos/msgs/foxglove_msgs/Color.py index 30362f837a..59d60ccc35 100644 --- a/dimos/msgs/foxglove_msgs/Color.py +++ b/dimos/msgs/foxglove_msgs/Color.py @@ -22,12 +22,13 @@ class Color(LCMColor): """Color with convenience methods.""" @classmethod - def from_string(cls, name: str, alpha: float = 0.2) -> Color: + def from_string(cls, name: str, alpha: float = 0.2, brightness: float = 1.0) -> Color: """Generate a consistent color from a string using hash function. Args: name: String to generate color from alpha: Transparency value (0.0-1.0) + brightness: Brightness multiplier (0.0-2.0). Values > 1.0 lighten towards white. Returns: Color instance with deterministic RGB values @@ -41,10 +42,23 @@ def from_string(cls, name: str, alpha: float = 0.2) -> Color: g = hash_bytes[1] / 255.0 b = hash_bytes[2] / 255.0 + # Apply brightness adjustment + # If brightness > 1.0, mix with white to lighten + if brightness > 1.0: + mix_factor = brightness - 1.0 # 0.0 to 1.0 + r = r + (1.0 - r) * mix_factor + g = g + (1.0 - g) * mix_factor + b = b + (1.0 - b) * mix_factor + else: + # If brightness < 1.0, darken by scaling + r *= brightness + g *= brightness + b *= brightness + # Create and return color instance color = cls() - color.r = r - color.g = g - color.b = b + color.r = min(1.0, r) + color.g = min(1.0, g) + color.b = min(1.0, b) color.a = alpha return color diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index 5cf2ed6da8..36f6f1d545 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -22,22 +22,22 @@ import cv2 import numpy as np import reactivex as rx +from dimos_lcm.sensor_msgs.Image import Image as LCMImage +from dimos_lcm.std_msgs.Header import Header +from reactivex import operators as ops +from reactivex.observable import Observable + from dimos.msgs.sensor_msgs.image_impls.AbstractImage import ( - AbstractImage, HAS_CUDA, HAS_NVIMGCODEC, - ImageFormat, NVIMGCODEC_LAST_USED, + AbstractImage, + ImageFormat, ) from dimos.msgs.sensor_msgs.image_impls.CudaImage import CudaImage from dimos.msgs.sensor_msgs.image_impls.NumpyImage import NumpyImage -from dimos_lcm.sensor_msgs.Image import Image as LCMImage -from dimos_lcm.std_msgs.Header import Header from dimos.types.timestamped import Timestamped, TimestampedBufferCollection, to_human_readable from dimos.utils.reactive import quality_barrier -from reactivex import operators as ops -from reactivex.observable import Observable -from reactivex.scheduler import ThreadPoolScheduler try: import cupy as cp # type: ignore diff --git a/dimos/msgs/sensor_msgs/image_impls/test_image_backends.py b/dimos/msgs/sensor_msgs/image_impls/test_image_backends.py index 931a30ea5f..0e19a24167 100644 --- a/dimos/msgs/sensor_msgs/image_impls/test_image_backends.py +++ b/dimos/msgs/sensor_msgs/image_impls/test_image_backends.py @@ -18,7 +18,7 @@ import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat, HAS_CUDA +from dimos.msgs.sensor_msgs.Image import HAS_CUDA, Image, ImageFormat from dimos.utils.data import get_data IMAGE_PATH = get_data("chair-image.png") @@ -416,6 +416,9 @@ def test_perf_solvepnp(alloc_timer): print(f"solvePnP (avg per call) cpu={cpu_t:.6f}s") +# this test is failing with +# raise RuntimeError("OpenCV CSRT tracker not available") +@pytest.mark.skip def test_perf_tracker(alloc_timer): """Test tracker performance with NumpyImage always, add CudaImage when available.""" # Don't check - just let it fail if CSRT isn't available @@ -461,6 +464,9 @@ def test_perf_tracker(alloc_timer): print(f"tracker (avg per call) cpu={cpu_t:.6f}s") +# this test is failing with +# raise RuntimeError("OpenCV CSRT tracker not available") +@pytest.mark.skip def test_csrt_tracker(alloc_timer): """Test CSRT tracker with NumpyImage always, add CudaImage parity when available.""" # Don't check - just let it fail if CSRT isn't available diff --git a/dimos/msgs/vision_msgs/Detection2DArray.py b/dimos/msgs/vision_msgs/Detection2DArray.py index 133893b9f0..79c84f7609 100644 --- a/dimos/msgs/vision_msgs/Detection2DArray.py +++ b/dimos/msgs/vision_msgs/Detection2DArray.py @@ -11,12 +11,17 @@ # 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 dimos_lcm.vision_msgs.Detection2DArray import Detection2DArray as LCMDetection2DArray +from dimos.types.timestamped import to_timestamp + class Detection2DArray(LCMDetection2DArray): msg_name = "vision_msgs.Detection2DArray" # for _get_field_type() to work when decoding in _decode_one() __annotations__ = LCMDetection2DArray.__annotations__ + + @property + def ts(self) -> float: + return to_timestamp(self.header.stamp) diff --git a/dimos/perception/detection/__init__.py b/dimos/perception/detection/__init__.py new file mode 100644 index 0000000000..72663a69b0 --- /dev/null +++ b/dimos/perception/detection/__init__.py @@ -0,0 +1,7 @@ +from dimos.perception.detection.detectors import * +from dimos.perception.detection.module2D import ( + Detection2DModule, +) +from dimos.perception.detection.module3D import ( + Detection3DModule, +) diff --git a/dimos/perception/detection2d/conftest.py b/dimos/perception/detection/conftest.py similarity index 54% rename from dimos/perception/detection2d/conftest.py rename to dimos/perception/detection/conftest.py index 8ada4ec356..cdd15c1f92 100644 --- a/dimos/perception/detection2d/conftest.py +++ b/dimos/perception/detection/conftest.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Optional, TypedDict, Union +import functools +from typing import Callable, Generator, Optional, TypedDict, Union import pytest from dimos_lcm.foxglove_msgs.ImageAnnotations import ImageAnnotations @@ -23,15 +24,14 @@ from dimos.msgs.geometry_msgs import Transform from dimos.msgs.sensor_msgs import CameraInfo, Image, PointCloud2 from dimos.msgs.vision_msgs import Detection2DArray -from dimos.perception.detection2d.module2D import Detection2DModule -from dimos.perception.detection2d.module3D import Detection3DModule -from dimos.perception.detection2d.moduleDB import ObjectDBModule -from dimos.perception.detection2d.type import ( +from dimos.perception.detection.module2D import Detection2DModule +from dimos.perception.detection.module3D import Detection3DModule +from dimos.perception.detection.moduleDB import ObjectDBModule +from dimos.perception.detection.type import ( Detection2D, Detection3D, Detection3DPC, ImageDetections2D, - ImageDetections3D, ImageDetections3DPC, ) from dimos.protocol.tf import TF @@ -60,44 +60,60 @@ class Moment2D(Moment): class Moment3D(Moment): - detections3dpc: ImageDetections3D + detections3dpc: ImageDetections3DPC -@pytest.fixture +@pytest.fixture(scope="session") def tf(): t = TF() yield t t.stop() -@pytest.fixture +@pytest.fixture(scope="session") def get_moment(tf): + @functools.lru_cache(maxsize=1) def moment_provider(**kwargs) -> Moment: + print("MOMENT PROVIDER ARGS:", kwargs) seek = kwargs.get("seek", 10.0) data_dir = "unitree_go2_lidar_corrected" get_data(data_dir) - lidar_frame = TimedSensorReplay(f"{data_dir}/lidar").find_closest_seek(seek) + lidar_frame_result = TimedSensorReplay(f"{data_dir}/lidar").find_closest_seek(seek) + if lidar_frame_result is None: + raise ValueError("No lidar frame found") + lidar_frame: LidarMessage = lidar_frame_result image_frame = TimedSensorReplay( f"{data_dir}/video", ).find_closest(lidar_frame.ts) + if image_frame is None: + raise ValueError("No image frame found") + image_frame.frame_id = "camera_optical" odom_frame = TimedSensorReplay(f"{data_dir}/odom", autocast=Odometry.from_msg).find_closest( lidar_frame.ts ) + if odom_frame is None: + raise ValueError("No odom frame found") + transforms = ConnectionModule._odom_to_tf(odom_frame) tf.receive_transform(*transforms) + camera_info_out = ConnectionModule._camera_info() + # ConnectionModule._camera_info() returns Out[CameraInfo], extract the value + from typing import cast + + camera_info = cast(CameraInfo, camera_info_out) return { "odom_frame": odom_frame, "lidar_frame": lidar_frame, "image_frame": image_frame, - "camera_info": ConnectionModule._camera_info(), + "camera_info": camera_info, "transforms": transforms, "tf": tf, } @@ -105,40 +121,56 @@ def moment_provider(**kwargs) -> Moment: return moment_provider -@pytest.fixture +@pytest.fixture(scope="session") def publish_moment(): def publisher(moment: Moment | Moment2D | Moment3D): - if moment.get("detections2d"): + detections2d_val = moment.get("detections2d") + if detections2d_val: # 2d annotations - annotations = LCMTransport("/annotations", ImageAnnotations) - annotations.publish(moment.get("detections2d").to_foxglove_annotations()) + annotations: LCMTransport[ImageAnnotations] = LCMTransport( + "/annotations", ImageAnnotations + ) + assert isinstance(detections2d_val, ImageDetections2D) + annotations.publish(detections2d_val.to_foxglove_annotations()) - detections = LCMTransport("/detections", Detection2DArray) - detections.publish(moment.get("detections2d").to_ros_detection2d_array()) + detections: LCMTransport[Detection2DArray] = LCMTransport( + "/detections", Detection2DArray + ) + detections.publish(detections2d_val.to_ros_detection2d_array()) annotations.lcm.stop() detections.lcm.stop() - if moment.get("detections3dpc"): - scene_update = LCMTransport("/scene_update", SceneUpdate) + detections3dpc_val = moment.get("detections3dpc") + if detections3dpc_val: + scene_update: LCMTransport[SceneUpdate] = LCMTransport("/scene_update", SceneUpdate) # 3d scene update - scene_update.publish(moment.get("detections3dpc").to_foxglove_scene_update()) + assert isinstance(detections3dpc_val, ImageDetections3DPC) + scene_update.publish(detections3dpc_val.to_foxglove_scene_update()) scene_update.lcm.stop() - lidar = LCMTransport("/lidar", PointCloud2) - lidar.publish(moment.get("lidar_frame")) - lidar.lcm.stop() + lidar_frame = moment.get("lidar_frame") + if lidar_frame: + lidar: LCMTransport[PointCloud2] = LCMTransport("/lidar", PointCloud2) + lidar.publish(lidar_frame) + lidar.lcm.stop() - image = LCMTransport("/image", Image) - image.publish(moment.get("image_frame")) - image.lcm.stop() + image_frame = moment.get("image_frame") + if image_frame: + image: LCMTransport[Image] = LCMTransport("/image", Image) + image.publish(image_frame) + image.lcm.stop() - camera_info = LCMTransport("/camera_info", CameraInfo) - camera_info.publish(moment.get("camera_info")) - camera_info.lcm.stop() + camera_info_val = moment.get("camera_info") + if camera_info_val: + camera_info: LCMTransport[CameraInfo] = LCMTransport("/camera_info", CameraInfo) + camera_info.publish(camera_info_val) + camera_info.lcm.stop() tf = moment.get("tf") - tf.publish(*moment.get("transforms")) + transforms = moment.get("transforms") + if tf is not None and transforms is not None: + tf.publish(*transforms) # moduleDB.scene_update.transport = LCMTransport("/scene_update", SceneUpdate) # moduleDB.target.transport = LCMTransport("/target", PoseStamped) @@ -146,24 +178,39 @@ def publisher(moment: Moment | Moment2D | Moment3D): return publisher -@pytest.fixture +@pytest.fixture(scope="session") +def imageDetections2d(get_moment_2d) -> ImageDetections2D: + moment = get_moment_2d() + assert len(moment["detections2d"]) > 0, "No detections found in the moment" + return moment["detections2d"] + + +@pytest.fixture(scope="session") def detection2d(get_moment_2d) -> Detection2D: - moment = get_moment_2d(seek=10.0) + moment = get_moment_2d() assert len(moment["detections2d"]) > 0, "No detections found in the moment" return moment["detections2d"][0] -@pytest.fixture -def detection3dpc(get_moment_3dpc) -> Detection3DPC: +@pytest.fixture(scope="session") +def detections3dpc(get_moment_3dpc) -> Detection3DPC: moment = get_moment_3dpc(seek=10.0) assert len(moment["detections3dpc"]) > 0, "No detections found in the moment" - return moment["detections3dpc"][0] + return moment["detections3dpc"] + + +@pytest.fixture(scope="session") +def detection3dpc(detections3dpc) -> Detection3DPC: + return detections3dpc[0] -@pytest.fixture -def get_moment_2d(get_moment) -> Callable[[], Moment2D]: - module = Detection2DModule() +@pytest.fixture(scope="session") +def get_moment_2d(get_moment) -> Generator[Callable[[], Moment2D], None, None]: + from dimos.perception.detection.detectors import Yolo2DDetector + module = Detection2DModule(detector=lambda: Yolo2DDetector(device="cpu")) + + @functools.lru_cache(maxsize=1) def moment_provider(**kwargs) -> Moment2D: moment = get_moment(**kwargs) detections = module.process_image_frame(moment.get("image_frame")) @@ -174,39 +221,50 @@ def moment_provider(**kwargs) -> Moment2D: } yield moment_provider + module._close_module() -@pytest.fixture -def get_moment_3dpc(get_moment_2d) -> Callable[[], Moment2D]: - module = None +@pytest.fixture(scope="session") +def get_moment_3dpc(get_moment_2d) -> Generator[Callable[[], Moment3D], None, None]: + module: Optional[Detection3DModule] = None - def moment_provider(**kwargs) -> Moment2D: + @functools.lru_cache(maxsize=1) + def moment_provider(**kwargs) -> Moment3D: nonlocal module moment = get_moment_2d(**kwargs) if not module: module = Detection3DModule(camera_info=moment["camera_info"]) - camera_transform = moment["tf"].get("camera_optical", moment.get("lidar_frame").frame_id) + lidar_frame = moment.get("lidar_frame") + if lidar_frame is None: + raise ValueError("No lidar frame found") + + camera_transform = moment["tf"].get("camera_optical", lidar_frame.frame_id) if camera_transform is None: raise ValueError("No camera_optical transform in tf") + + detections3dpc = module.process_frame( + moment["detections2d"], moment["lidar_frame"], camera_transform + ) + return { **moment, - "detections3dpc": module.process_frame( - moment["detections2d"], moment["lidar_frame"], camera_transform - ), + "detections3dpc": detections3dpc, } yield moment_provider - print("Closing 3D detection module", module) - module._close_module() + if module is not None: + module._close_module() -@pytest.fixture +@pytest.fixture(scope="session") def object_db_module(get_moment): """Create and populate an ObjectDBModule with detections from multiple frames.""" - module2d = Detection2DModule() + from dimos.perception.detection.detectors import Yolo2DDetector + + module2d = Detection2DModule(detector=lambda: Yolo2DDetector(device="cpu")) module3d = Detection3DModule(camera_info=ConnectionModule._camera_info()) moduleDB = ObjectDBModule( camera_info=ConnectionModule._camera_info(), @@ -233,12 +291,13 @@ def object_db_module(get_moment): moduleDB.add_detections(imageDetections3d) yield moduleDB + module2d._close_module() module3d._close_module() moduleDB._close_module() -@pytest.fixture +@pytest.fixture(scope="session") def first_object(object_db_module): """Get the first object from the database.""" objects = list(object_db_module.objects.values()) @@ -246,7 +305,7 @@ def first_object(object_db_module): return objects[0] -@pytest.fixture +@pytest.fixture(scope="session") def all_objects(object_db_module): """Get all objects from the database.""" return list(object_db_module.objects.values()) diff --git a/dimos/perception/detection/detectors/__init__.py b/dimos/perception/detection/detectors/__init__.py new file mode 100644 index 0000000000..d6383d084e --- /dev/null +++ b/dimos/perception/detection/detectors/__init__.py @@ -0,0 +1,3 @@ +# from dimos.perception.detection.detectors.detic import Detic2DDetector +from dimos.perception.detection.detectors.types import Detector +from dimos.perception.detection.detectors.yolo import Yolo2DDetector diff --git a/dimos/perception/detection2d/detectors/config/custom_tracker.yaml b/dimos/perception/detection/detectors/config/custom_tracker.yaml similarity index 100% rename from dimos/perception/detection2d/detectors/config/custom_tracker.yaml rename to dimos/perception/detection/detectors/config/custom_tracker.yaml diff --git a/dimos/perception/detection/detectors/conftest.py b/dimos/perception/detection/detectors/conftest.py new file mode 100644 index 0000000000..7caca818c9 --- /dev/null +++ b/dimos/perception/detection/detectors/conftest.py @@ -0,0 +1,38 @@ +# Copyright 2025 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. + +import pytest + +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector +from dimos.perception.detection.detectors.yolo import Yolo2DDetector +from dimos.utils.data import get_data + + +@pytest.fixture(scope="session") +def test_image(): + """Load the test image used for detector tests.""" + return Image.from_file(get_data("cafe.jpg")) + + +@pytest.fixture(scope="session") +def person_detector(): + """Create a YoloPersonDetector instance.""" + return YoloPersonDetector() + + +@pytest.fixture(scope="session") +def bbox_detector(): + """Create a Yolo2DDetector instance for general object detection.""" + return Yolo2DDetector() diff --git a/dimos/perception/detection2d/detectors/detic.py b/dimos/perception/detection/detectors/detic.py similarity index 98% rename from dimos/perception/detection2d/detectors/detic.py rename to dimos/perception/detection/detectors/detic.py index 0b7b63276f..db2d8bb634 100644 --- a/dimos/perception/detection2d/detectors/detic.py +++ b/dimos/perception/detection/detectors/detic.py @@ -18,16 +18,16 @@ import numpy as np from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.types import Detector +from dimos.perception.detection.detectors.types import Detector from dimos.perception.detection2d.utils import plot_results # Add Detic to Python path from dimos.constants import DIMOS_PROJECT_ROOT detic_path = DIMOS_PROJECT_ROOT / "dimos/models/Detic" -if detic_path not in sys.path: - sys.path.append(detic_path) - sys.path.append(os.path.join(detic_path, "third_party/CenterNet2")) +if str(detic_path) not in sys.path: + sys.path.append(str(detic_path)) + sys.path.append(str(detic_path / "third_party/CenterNet2")) # PIL patch for compatibility import PIL.Image diff --git a/dimos/perception/detection2d/detectors/person/test_yolo.py b/dimos/perception/detection/detectors/person/test_person_detectors.py similarity index 58% rename from dimos/perception/detection2d/detectors/person/test_yolo.py rename to dimos/perception/detection/detectors/person/test_person_detectors.py index 454997ca27..bca39acbcd 100644 --- a/dimos/perception/detection2d/detectors/person/test_yolo.py +++ b/dimos/perception/detection/detectors/person/test_person_detectors.py @@ -14,25 +14,17 @@ import pytest -from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.person.yolo import YoloPersonDetector -from dimos.perception.detection2d.type.person import Person -from dimos.utils.data import get_data +from dimos.perception.detection.type import Detection2DPerson, ImageDetections2D -@pytest.fixture() -def detector(): - return YoloPersonDetector() +@pytest.fixture(scope="session") +def people(person_detector, test_image): + return person_detector.process_image(test_image) -@pytest.fixture() -def test_image(): - return Image.from_file(get_data("cafe.jpg")) - - -@pytest.fixture() -def people(detector, test_image): - return detector.detect_people(test_image) +@pytest.fixture(scope="session") +def person(people): + return people[0] def test_person_detection(people): @@ -41,7 +33,7 @@ def test_person_detection(people): # Check first person person = people[0] - assert isinstance(person, Person) + assert isinstance(person, Detection2DPerson) assert person.confidence > 0 assert len(person.bbox) == 4 # bbox is now a tuple assert person.keypoints.shape == (17, 2) @@ -49,7 +41,7 @@ def test_person_detection(people): def test_person_properties(people): - """Test Person object properties and methods.""" + """Test Detection2DPerson object properties and methods.""" person = people[0] # Test bounding box properties @@ -101,12 +93,19 @@ def test_multiple_people(people): print(f" {name}: ({xy[0]:.1f}, {xy[1]:.1f}) conf={conf:.3f}") +def test_image_detections2d_structure(people): + """Test that process_image returns ImageDetections2D.""" + assert isinstance(people, ImageDetections2D) + assert len(people.detections) > 0 + assert all(isinstance(d, Detection2DPerson) for d in people.detections) + + def test_invalid_keypoint(test_image): """Test error handling for invalid keypoint names.""" - # Create a dummy person + # Create a dummy Detection2DPerson import numpy as np - person = Person( + person = Detection2DPerson( # Detection2DBBox fields bbox=(0.0, 0.0, 100.0, 100.0), track_id=0, @@ -115,10 +114,47 @@ def test_invalid_keypoint(test_image): name="person", ts=test_image.ts, image=test_image, - # Person fields + # Detection2DPerson fields keypoints=np.zeros((17, 2)), keypoint_scores=np.zeros(17), ) with pytest.raises(ValueError): person.get_keypoint("invalid_keypoint") + + +def test_person_annotations(person): + # Test text annotations + text_anns = person.to_text_annotation() + print(f"\nText annotations: {len(text_anns)}") + for i, ann in enumerate(text_anns): + print(f" {i}: {ann.text}") + assert len(text_anns) == 3 # confidence, name/track_id, keypoints count + assert any("keypoints:" in ann.text for ann in text_anns) + + # Test points annotations + points_anns = person.to_points_annotation() + print(f"\nPoints annotations: {len(points_anns)}") + + # Count different types (use actual LCM constants) + from dimos_lcm.foxglove_msgs.ImageAnnotations import PointsAnnotation + + bbox_count = sum(1 for ann in points_anns if ann.type == PointsAnnotation.LINE_LOOP) # 2 + keypoint_count = sum(1 for ann in points_anns if ann.type == PointsAnnotation.POINTS) # 1 + skeleton_count = sum(1 for ann in points_anns if ann.type == PointsAnnotation.LINE_LIST) # 4 + + print(f" - Bounding boxes: {bbox_count}") + print(f" - Keypoint circles: {keypoint_count}") + print(f" - Skeleton lines: {skeleton_count}") + + assert bbox_count >= 1 # At least the person bbox + assert keypoint_count >= 1 # At least some visible keypoints + assert skeleton_count >= 1 # At least some skeleton connections + + # Test full image annotations + img_anns = person.to_image_annotations() + assert img_anns.texts_length == len(text_anns) + assert img_anns.points_length == len(points_anns) + + print(f"\nāœ“ Person annotations working correctly!") + print(f" - {len(person.get_visible_keypoints(0.5))}/17 visible keypoints") diff --git a/dimos/perception/detection/detectors/person/yolo.py b/dimos/perception/detection/detectors/person/yolo.py new file mode 100644 index 0000000000..05e79fa22f --- /dev/null +++ b/dimos/perception/detection/detectors/person/yolo.py @@ -0,0 +1,75 @@ +# Copyright 2025 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 ultralytics import YOLO + +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.detectors.types import Detector +from dimos.perception.detection.type 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("dimos.perception.detection.yolo.person") + + +class YoloPersonDetector(Detector): + def __init__(self, model_path="models_yolo", model_name="yolo11n-pose.pt", device: str = None): + self.model = YOLO(get_data(model_path) / model_name, task="track") + + self.tracker = get_data(model_path) / "botsort.yaml" + + if device: + self.device = device + return + + if is_cuda_available(): + self.device = "cuda" + logger.info("Using CUDA for YOLO person detector") + else: + self.device = "cpu" + logger.info("Using CPU for YOLO person detector") + + def process_image(self, image: Image) -> ImageDetections2D: + """Process image and return detection results. + + Args: + image: Input image + + Returns: + ImageDetections2D containing Detection2DPerson objects with pose keypoints + """ + results = self.model.track( + source=image.to_opencv(), + verbose=False, + conf=0.5, + tracker=self.tracker, + persist=True, + device=self.device, + ) + return ImageDetections2D.from_ultralytics_result(image, results) + + def stop(self): + """ + Clean up resources used by the detector, including tracker threads. + """ + if hasattr(self.model, "predictor") and self.model.predictor is not None: + predictor = self.model.predictor + if hasattr(predictor, "trackers") and predictor.trackers: + for tracker in predictor.trackers: + if hasattr(tracker, "tracker") and hasattr(tracker.tracker, "gmc"): + gmc = tracker.tracker.gmc + if hasattr(gmc, "executor") and gmc.executor is not None: + gmc.executor.shutdown(wait=True) + self.model.predictor = None diff --git a/dimos/perception/detection/detectors/test_bbox_detectors.py b/dimos/perception/detection/detectors/test_bbox_detectors.py new file mode 100644 index 0000000000..d246ded8a3 --- /dev/null +++ b/dimos/perception/detection/detectors/test_bbox_detectors.py @@ -0,0 +1,158 @@ +# Copyright 2025 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. + +import pytest + +from dimos.perception.detection.type import Detection2D, ImageDetections2D + + +@pytest.fixture(params=["bbox_detector", "person_detector"], scope="session") +def detector(request): + """Parametrized fixture that provides both bbox and person detectors.""" + return request.getfixturevalue(request.param) + + +@pytest.fixture(scope="session") +def detections(detector, test_image): + """Get ImageDetections2D from any detector.""" + return detector.process_image(test_image) + + +def test_detection_basic(detections): + """Test that we can detect objects with all detectors.""" + assert len(detections.detections) > 0 + + # Check first detection + detection = detections.detections[0] + assert isinstance(detection, Detection2D) + assert detection.confidence > 0 + assert len(detection.bbox) == 4 # bbox is a tuple (x1, y1, x2, y2) + assert detection.class_id >= 0 + assert detection.name is not None + + +def test_detection_bbox_properties(detections): + """Test Detection2D bbox properties work for all detectors.""" + detection = detections.detections[0] + + # Test bounding box is valid + x1, y1, x2, y2 = detection.bbox + assert x2 > x1, "x2 should be greater than x1" + assert y2 > y1, "y2 should be greater than y1" + assert all(coord >= 0 for coord in detection.bbox), "Coordinates should be non-negative" + + # Test bbox volume + volume = detection.bbox_2d_volume() + assert volume > 0 + expected_volume = (x2 - x1) * (y2 - y1) + assert abs(volume - expected_volume) < 0.01 + + # Test center calculation + center_x, center_y, width, height = detection.get_bbox_center() + assert center_x == (x1 + x2) / 2.0 + assert center_y == (y1 + y2) / 2.0 + assert width == x2 - x1 + assert height == y2 - y1 + + +def test_detection_cropped_image(detections, test_image): + """Test cropping image to detection bbox.""" + detection = detections.detections[0] + + # Test cropped image + cropped = detection.cropped_image(padding=20) + assert cropped is not None + + # Cropped image should be smaller than original (usually) + if test_image.shape: + assert cropped.shape[0] <= test_image.shape[0] + assert cropped.shape[1] <= test_image.shape[1] + + +def test_detection_annotations(detections): + """Test annotation generation for detections.""" + detection = detections.detections[0] + + # Test text annotations - all detections should have at least 2 + text_annotations = detection.to_text_annotation() + assert len(text_annotations) >= 2 # confidence and name/track_id (person has keypoints too) + + # Test points annotations - at least bbox + points_annotations = detection.to_points_annotation() + assert len(points_annotations) >= 1 # At least the bbox polygon + + # Test image annotations + annotations = detection.to_image_annotations() + assert annotations.texts_length >= 2 + assert annotations.points_length >= 1 + + +def test_detection_ros_conversion(detections): + """Test conversion to ROS Detection2D message.""" + detection = detections.detections[0] + + ros_det = detection.to_ros_detection2d() + + # Check bbox conversion + center_x, center_y, width, height = detection.get_bbox_center() + assert abs(ros_det.bbox.center.position.x - center_x) < 0.01 + assert abs(ros_det.bbox.center.position.y - center_y) < 0.01 + assert abs(ros_det.bbox.size_x - width) < 0.01 + assert abs(ros_det.bbox.size_y - height) < 0.01 + + # Check confidence and class_id + assert len(ros_det.results) > 0 + assert ros_det.results[0].hypothesis.score == detection.confidence + assert ros_det.results[0].hypothesis.class_id == detection.class_id + + +def test_detection_is_valid(detections): + """Test bbox validation.""" + detection = detections.detections[0] + + # Detection from real detector should be valid + assert detection.is_valid() + + +def test_image_detections2d_structure(detections): + """Test that process_image returns ImageDetections2D.""" + assert isinstance(detections, ImageDetections2D) + assert len(detections.detections) > 0 + assert all(isinstance(d, Detection2D) for d in detections.detections) + + +def test_multiple_detections(detections): + """Test that multiple objects can be detected.""" + print(f"\nDetected {len(detections.detections)} objects in test image") + + for i, detection in enumerate(detections.detections[:5]): # Show first 5 + print(f"\nDetection {i}:") + print(f" Class: {detection.name} (id: {detection.class_id})") + print(f" Confidence: {detection.confidence:.3f}") + print( + f" Bbox: ({detection.bbox[0]:.1f}, {detection.bbox[1]:.1f}, {detection.bbox[2]:.1f}, {detection.bbox[3]:.1f})" + ) + print(f" Track ID: {detection.track_id}") + + +def test_detection_string_representation(detections): + """Test string representation of detections.""" + detection = detections.detections[0] + str_repr = str(detection) + + # Should contain class name (either Detection2DBBox or Detection2DPerson) + assert "Detection2D" in str_repr + + # Should show object name + assert detection.name in str_repr or f"class_{detection.class_id}" in str_repr diff --git a/dimos/perception/detection2d/detectors/types.py b/dimos/perception/detection/detectors/types.py similarity index 81% rename from dimos/perception/detection2d/detectors/types.py rename to dimos/perception/detection/detectors/types.py index 639fc09247..1a3b0b5471 100644 --- a/dimos/perception/detection2d/detectors/types.py +++ b/dimos/perception/detection/detectors/types.py @@ -15,11 +15,9 @@ from abc import ABC, abstractmethod from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.type import ( - InconvinientDetectionFormat, -) +from dimos.perception.detection.type import ImageDetections2D class Detector(ABC): @abstractmethod - def process_image(self, image: Image) -> InconvinientDetectionFormat: ... + def process_image(self, image: Image) -> ImageDetections2D: ... diff --git a/dimos/perception/detection/detectors/yolo.py b/dimos/perception/detection/detectors/yolo.py new file mode 100644 index 0000000000..a338d3c8de --- /dev/null +++ b/dimos/perception/detection/detectors/yolo.py @@ -0,0 +1,78 @@ +# Copyright 2025 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 ultralytics import YOLO + +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.detectors.types import Detector +from dimos.perception.detection.type 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("dimos.perception.detection.yolo_2d_det") + + +class Yolo2DDetector(Detector): + def __init__(self, model_path="models_yolo", model_name="yolo11n.pt", device: str = None): + self.model = YOLO( + get_data(model_path) / model_name, + task="detect", + ) + + if device: + self.device = device + return + + if is_cuda_available(): + self.device = "cuda" + logger.debug("Using CUDA for YOLO 2d detector") + else: + self.device = "cpu" + logger.debug("Using CPU for YOLO 2d detector") + + def process_image(self, image: Image) -> ImageDetections2D: + """ + Process an image and return detection results. + + Args: + image: Input image + + Returns: + ImageDetections2D containing all detected objects + """ + results = self.model.track( + source=image.to_opencv(), + device=self.device, + conf=0.5, + iou=0.6, + persist=True, + verbose=False, + ) + + return ImageDetections2D.from_ultralytics_result(image, results) + + def stop(self): + """ + Clean up resources used by the detector, including tracker threads. + """ + if hasattr(self.model, "predictor") and self.model.predictor is not None: + predictor = self.model.predictor + if hasattr(predictor, "trackers") and predictor.trackers: + for tracker in predictor.trackers: + if hasattr(tracker, "tracker") and hasattr(tracker.tracker, "gmc"): + gmc = tracker.tracker.gmc + if hasattr(gmc, "executor") and gmc.executor is not None: + gmc.executor.shutdown(wait=True) + self.model.predictor = None diff --git a/dimos/perception/detection/module2D.py b/dimos/perception/detection/module2D.py new file mode 100644 index 0000000000..c4b0ba5a43 --- /dev/null +++ b/dimos/perception/detection/module2D.py @@ -0,0 +1,172 @@ +# Copyright 2025 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 dataclasses import dataclass +from typing import Any, Callable, Optional, Tuple + +from dimos_lcm.foxglove_msgs.ImageAnnotations import ( + ImageAnnotations, +) +from dimos_lcm.sensor_msgs import CameraInfo +from reactivex import operators as ops +from reactivex.observable import Observable +from reactivex.subject import Subject + +from dimos.core import In, Module, Out, rpc +from dimos.core.module import ModuleConfig +from dimos.msgs.geometry_msgs import Transform, Vector3 +from dimos.msgs.sensor_msgs import Image +from dimos.msgs.sensor_msgs.Image import sharpness_barrier +from dimos.msgs.vision_msgs import Detection2DArray +from dimos.perception.detection.detectors import Detector +from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector +from dimos.perception.detection.detectors.yolo import Yolo2DDetector +from dimos.perception.detection.type import ( + ImageDetections2D, +) +from dimos.utils.decorators.decorators import simple_mcache +from dimos.utils.reactive import backpressure + + +@dataclass +class Config(ModuleConfig): + max_freq: float = 10 + detector: Optional[Callable[[Any], Detector]] = YoloPersonDetector + camera_info: CameraInfo = CameraInfo() + + +class Detection2DModule(Module): + default_config = Config + config: Config + detector: Detector + + image: In[Image] = None # type: ignore + + detections: Out[Detection2DArray] = None # type: ignore + annotations: Out[ImageAnnotations] = None # type: ignore + + detected_image_0: Out[Image] = None # type: ignore + detected_image_1: Out[Image] = None # type: ignore + detected_image_2: Out[Image] = None # type: ignore + + cnt: int = 0 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.config: Config = Config(**kwargs) + self.detector = self.config.detector() + self.vlm_detections_subject = Subject() + self.previous_detection_count = 0 + + def process_image_frame(self, image: Image) -> ImageDetections2D: + return self.detector.process_image(image) + + @simple_mcache + def sharp_image_stream(self) -> Observable[Image]: + return backpressure( + self.image.pure_observable().pipe( + sharpness_barrier(self.config.max_freq), + ) + ) + + @simple_mcache + def detection_stream_2d(self) -> Observable[ImageDetections2D]: + return backpressure(self.image.observable().pipe(ops.map(self.process_image_frame))) + + def pixel_to_3d( + self, + pixel: Tuple[int, int], + camera_info: CameraInfo, + assumed_depth: float = 1.0, + ) -> Vector3: + """Unproject 2D pixel coordinates to 3D position in camera optical frame. + + Args: + camera_info: Camera calibration information + assumed_depth: Assumed depth in meters (default 1.0m from camera) + + Returns: + Vector3 position in camera optical frame coordinates + """ + # Extract camera intrinsics + fx, fy = camera_info.K[0], camera_info.K[4] + cx, cy = camera_info.K[2], camera_info.K[5] + + # Unproject pixel to normalized camera coordinates + x_norm = (pixel[0] - cx) / fx + y_norm = (pixel[1] - cy) / fy + + # Create 3D point at assumed depth in camera optical frame + # Camera optical frame: X right, Y down, Z forward + return Vector3(x_norm * assumed_depth, y_norm * assumed_depth, assumed_depth) + + def track(self, detections: ImageDetections2D): + sensor_frame = self.tf.get("sensor", "camera_optical", detections.image.ts, 5.0) + + if not sensor_frame: + return + + if not detections.detections: + return + + sensor_frame.child_frame_id = "sensor_frame" + transforms = [sensor_frame] + + current_count = len(detections.detections) + max_count = max(current_count, self.previous_detection_count) + + # Publish transforms for all detection slots up to max_count + for index in range(max_count): + if index < current_count: + # Active detection - compute real position + detection = detections.detections[index] + position_3d = self.pixel_to_3d( + detection.center_bbox, self.config.camera_info, assumed_depth=1.0 + ) + else: + # No detection at this index - publish zero transform + position_3d = Vector3(0.0, 0.0, 0.0) + + transforms.append( + Transform( + frame_id=sensor_frame.child_frame_id, + child_frame_id=f"det_{index}", + ts=detections.image.ts, + translation=position_3d, + ) + ) + + self.previous_detection_count = current_count + self.tf.publish(*transforms) + + @rpc + def start(self): + self.detection_stream_2d().subscribe(self.track) + + self.detection_stream_2d().subscribe( + lambda det: self.detections.publish(det.to_ros_detection2d_array()) + ) + + self.detection_stream_2d().subscribe( + lambda det: self.annotations.publish(det.to_foxglove_annotations()) + ) + + def publish_cropped_images(detections: ImageDetections2D): + for index, detection in enumerate(detections[:3]): + image_topic = getattr(self, "detected_image_" + str(index)) + image_topic.publish(detection.cropped_image()) + + self.detection_stream_2d().subscribe(publish_cropped_images) + + @rpc + def stop(self): ... diff --git a/dimos/perception/detection2d/module3D.py b/dimos/perception/detection/module3D.py similarity index 54% rename from dimos/perception/detection2d/module3D.py rename to dimos/perception/detection/module3D.py index 3b1dc040a1..b8fe42da9a 100644 --- a/dimos/perception/detection2d/module3D.py +++ b/dimos/perception/detection/module3D.py @@ -13,61 +13,97 @@ # limitations under the License. -from dimos_lcm.sensor_msgs import CameraInfo +from typing import Optional + +from dimos_lcm.foxglove_msgs.ImageAnnotations import ImageAnnotations +from lcm_msgs.foxglove_msgs import SceneUpdate from reactivex import operators as ops from reactivex.observable import Observable +from dimos.agents2 import skill from dimos.core import In, Out, rpc from dimos.msgs.geometry_msgs import Transform from dimos.msgs.sensor_msgs import Image, PointCloud2 -from dimos.perception.detection2d.module2D import Detection2DModule -from dimos.perception.detection2d.type import ( +from dimos.msgs.vision_msgs import Detection2DArray +from dimos.perception.detection.module2D import Config as Module2DConfig +from dimos.perception.detection.module2D import Detection2DModule +from dimos.perception.detection.type import ( ImageDetections2D, - ImageDetections3D, ImageDetections3DPC, ) -from dimos.perception.detection2d.type.detection3dpc import Detection3DPC +from dimos.perception.detection.type.detection3d import Detection3DPC from dimos.types.timestamped import align_timestamped from dimos.utils.reactive import backpressure -class Detection3DModule(Detection2DModule): - camera_info: CameraInfo +class Config(Module2DConfig): ... + +class Detection3DModule(Detection2DModule): image: In[Image] = None # type: ignore pointcloud: In[PointCloud2] = None # type: ignore + detections: Out[Detection2DArray] = None # type: ignore + annotations: Out[ImageAnnotations] = None # type: ignore + scene_update: Out[SceneUpdate] = None # type: ignore + + # just for visualization, + # emits latest pointclouds of detected objects in a frame detected_pointcloud_0: Out[PointCloud2] = None # type: ignore detected_pointcloud_1: Out[PointCloud2] = None # type: ignore detected_pointcloud_2: Out[PointCloud2] = None # type: ignore - detection_3d_stream: Observable[ImageDetections3DPC] = None + # just for visualization, emits latest top 3 detections in a frame + detected_image_0: Out[Image] = None # type: ignore + detected_image_1: Out[Image] = None # type: ignore + detected_image_2: Out[Image] = None # type: ignore - def __init__(self, camera_info: CameraInfo, *args, **kwargs): - super().__init__(*args, **kwargs) - self.camera_info = camera_info + detection_3d_stream: Optional[Observable[ImageDetections3DPC]] = None def process_frame( self, detections: ImageDetections2D, pointcloud: PointCloud2, transform: Transform, - ) -> ImageDetections3D: + ) -> ImageDetections3DPC: if not transform: - return ImageDetections3D(detections.image, []) + return ImageDetections3DPC(detections.image, []) - detection3d_list = [] + detection3d_list: list[Detection3DPC] = [] for detection in detections: detection3d = Detection3DPC.from_2d( detection, world_pointcloud=pointcloud, - camera_info=self.camera_info, + camera_info=self.config.camera_info, world_to_optical_transform=transform, ) if detection3d is not None: detection3d_list.append(detection3d) - return ImageDetections3D(detections.image, detection3d_list) + return ImageDetections3DPC(detections.image, detection3d_list) + + @skill # type: ignore[arg-type] + def ask_vlm(self, question: str) -> str | ImageDetections3DPC: + """ + query visual model about the view in front of the camera + you can ask to mark objects like: + + "red cup on the table left of the pencil" + "laptop on the desk" + "a person wearing a red shirt" + """ + from dimos.models.vl.qwen import QwenVlModel + + model = QwenVlModel() + result = model.query(self.image.get_next(), question) + + if isinstance(result, str) or not result or not len(result): + return "No detections" + + detections: ImageDetections2D = result + pc = self.pointcloud.get_next() + transform = self.tf.get("camera_optical", pc.frame_id, detections.image.ts, 5.0) + return self.process_frame(detections, pc, transform) @rpc def start(self): @@ -85,17 +121,13 @@ def detection2d_to_3d(args): buffer_size=20.0, ).pipe(ops.map(detection2d_to_3d)) - # self.detection_stream_3d = backpressure(self.detection_stream_2d()).pipe( - # ops.with_latest_from(self.pointcloud.observable()), ops.map(detection2d_to_3d) - # ) - self.detection_stream_3d.subscribe(self._publish_detections) @rpc def stop(self) -> None: super().stop() - def _publish_detections(self, detections: ImageDetections3D): + def _publish_detections(self, detections: ImageDetections3DPC): if not detections: return diff --git a/dimos/perception/detection2d/moduleDB.py b/dimos/perception/detection/moduleDB.py similarity index 71% rename from dimos/perception/detection2d/moduleDB.py rename to dimos/perception/detection/moduleDB.py index 456b1d8c87..ccc14d96f5 100644 --- a/dimos/perception/detection2d/moduleDB.py +++ b/dimos/perception/detection/moduleDB.py @@ -25,28 +25,35 @@ from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Transform, Vector3 from dimos.msgs.sensor_msgs import Image, PointCloud2 from dimos.msgs.vision_msgs import Detection2DArray -from dimos.perception.detection2d.module3D import Detection3DModule -from dimos.perception.detection2d.type import Detection3D, ImageDetections3D, TableStr +from dimos.perception.detection.module3D import Detection3DModule +from dimos.perception.detection.type import Detection3D, ImageDetections3DPC, TableStr +from dimos.perception.detection.type.detection3d import Detection3DPC from dimos.protocol.skill.skill import skill from dimos.protocol.skill.type import Output, Reducer, Stream from dimos.types.timestamped import to_datetime # Represents an object in space, as collection of 3d detections over time -class Object3D(Detection3D): - best_detection: Detection3D = None - center: Vector3 = None - track_id: str = None +class Object3D(Detection3DPC): + best_detection: Optional[Detection3DPC] = None # type: ignore + center: Optional[Vector3] = None # type: ignore + track_id: Optional[str] = None # type: ignore detections: int = 0 def to_repr_dict(self) -> Dict[str, Any]: + if self.center is None: + center_str = "None" + else: + center_str = ( + "[" + ", ".join(list(map(lambda n: f"{n:1f}", self.center.to_list()))) + "]" + ) return { "object_id": self.track_id, "detections": self.detections, - "center": "[" + ", ".join(list(map(lambda n: f"{n:1f}", self.center.to_list()))) + "]", + "center": center_str, } - def __init__(self, track_id: str, detection: Optional[Detection3D] = None, *args, **kwargs): + def __init__(self, track_id: str, detection: Optional[Detection3DPC] = None, *args, **kwargs): if detection is None: return self.ts = detection.ts @@ -62,7 +69,9 @@ def __init__(self, track_id: str, detection: Optional[Detection3D] = None, *args self.detections = self.detections + 1 self.best_detection = detection - def __add__(self, detection: Detection3D) -> "Object3D": + def __add__(self, detection: Detection3DPC) -> "Object3D": + if self.track_id is None: + raise ValueError("Cannot add detection to object with None track_id") new_object = Object3D(self.track_id) new_object.bbox = detection.bbox new_object.confidence = max(self.confidence, detection.confidence) @@ -83,9 +92,8 @@ def __add__(self, detection: Detection3D) -> "Object3D": return new_object - @property - def image(self) -> Image: - return self.best_detection.image + def get_image(self) -> Optional[Image]: + return self.best_detection.image if self.best_detection else None def scene_entity_label(self) -> str: return f"{self.name} ({self.detections})" @@ -100,6 +108,9 @@ def agent_encode(self): } def to_pose(self) -> PoseStamped: + if self.best_detection is None or self.center is None: + raise ValueError("Cannot compute pose without best_detection and center") + optical_inverse = Transform( translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), @@ -126,9 +137,9 @@ def to_pose(self) -> PoseStamped: class ObjectDBModule(Detection3DModule, TableStr): cnt: int = 0 objects: dict[str, Object3D] - object_stream: Observable[Object3D] = None + object_stream: Optional[Observable[Object3D]] = None - goto: Callable[[PoseStamped], Any] = None + goto: Optional[Callable[[PoseStamped], Any]] = None image: In[Image] = None # type: ignore pointcloud: In[PointCloud2] = None # type: ignore @@ -156,7 +167,7 @@ def __init__(self, goto: Callable[[PoseStamped], Any], *args, **kwargs): self.objects = {} self.remembered_locations = {} - def closest_object(self, detection: Detection3D) -> Optional[Object3D]: + def closest_object(self, detection: Detection3DPC) -> Optional[Object3D]: # Filter objects to only those with matching names matching_objects = [obj for obj in self.objects.values() if obj.name == detection.name] @@ -168,12 +179,12 @@ def closest_object(self, detection: Detection3D) -> Optional[Object3D]: return distances[0] - def add_detections(self, detections: List[Detection3D]) -> List[Object3D]: + def add_detections(self, detections: List[Detection3DPC]) -> List[Object3D]: return [ detection for detection in map(self.add_detection, detections) if detection is not None ] - def add_detection(self, detection: Detection3D): + def add_detection(self, detection: Detection3DPC): """Add detection to existing object or create new one.""" closest = self.closest_object(detection) if closest and closest.bounding_box_intersects(detection): @@ -181,18 +192,20 @@ def add_detection(self, detection: Detection3D): else: return self.create_new_object(detection) - def add_to_object(self, closest: Object3D, detection: Detection3D): + def add_to_object(self, closest: Object3D, detection: Detection3DPC): new_object = closest + detection - self.objects[closest.track_id] = new_object + if closest.track_id is not None: + self.objects[closest.track_id] = new_object return new_object - def create_new_object(self, detection: Detection3D): + def create_new_object(self, detection: Detection3DPC): new_object = Object3D(f"obj_{self.cnt}", detection) - self.objects[new_object.track_id] = new_object + if new_object.track_id is not None: + self.objects[new_object.track_id] = new_object self.cnt += 1 return new_object - def agent_encode(self) -> List[Any]: + def agent_encode(self) -> str: ret = [] for obj in copy(self.objects).values(): # we need at least 3 detectieons to consider it a valid object @@ -204,8 +217,8 @@ def agent_encode(self) -> List[Any]: return "No objects detected yet." return "\n".join(ret) - def vlm_query(self, description: str) -> str: - imageDetections2D = super().vlm_query(description) + def vlm_query(self, description: str) -> Optional[Object3D]: # type: ignore[override] + imageDetections2D = super().ask_vlm(description) print("VLM query found", imageDetections2D, "detections") time.sleep(3) @@ -234,68 +247,7 @@ def vlm_query(self, description: str) -> str: return ret[0] if ret else None - @skill() - def remember_location(self, name: str) -> str: - """Remember the current location with a name.""" - transform = self.tf.get("map", "sensor", time_point=time.time(), time_tolerance=1.0) - if not transform: - return f"Could not get current location transform from map to sensor" - - pose = transform.to_pose() - pose.frame_id = "map" - self.remembered_locations[name] = pose - return f"Location '{name}' saved at position: {pose.position}" - - @skill() - def goto_remembered_location(self, name: str) -> str: - """Go to a remembered location by name.""" - pose = self.remembered_locations.get(name, None) - if not pose: - return f"Location {name} not found. Known locations: {list(self.remembered_locations.keys())}" - self.goto(pose) - return f"Navigating to remembered location {name} and pose {pose}" - - @skill() - def list_remembered_locations(self) -> List[str]: - """List all remembered locations.""" - return str(list(self.remembered_locations.keys())) - - def nav_to(self, target_pose) -> str: - target_pose.orientation = Quaternion(0.0, 0.0, 0.0, 0.0) - self.target.publish(target_pose) - time.sleep(0.1) - self.target.publish(target_pose) - self.goto(target_pose) - - @skill() - def navigate_to_object_in_view(self, query: str) -> str: - """Navigate to an object in your current image view via natural language query using vision-language model to find it.""" - target_obj = self.vlm_query(query) - if not target_obj: - return f"No objects found matching '{query}'" - return self.navigate_to_object_by_id(target_obj.track_id) - - @skill(reducer=Reducer.all) - def list_objects(self): - """List all detected objects that the system remembers and can navigate to.""" - data = self.agent_encode() - return data - - @skill() - def navigate_to_object_by_id(self, object_id: str): - """Navigate to an object by an object id""" - target_obj = self.objects.get(object_id, None) - if not target_obj: - return f"Object {object_id} not found\nHere are the known objects:\n{str(self.agent_encode())}" - target_pose = target_obj.to_pose() - target_pose.frame_id = "map" - self.target.publish(target_pose) - time.sleep(0.1) - self.target.publish(target_pose) - self.nav_to(target_pose) - return f"Navigating to f{object_id} f{target_obj.name}" - - def lookup(self, label: str) -> List[Detection3D]: + def lookup(self, label: str) -> List[Detection3DPC]: """Look up a detection by label.""" return [] @@ -303,7 +255,7 @@ def lookup(self, label: str) -> List[Detection3D]: def start(self): Detection3DModule.start(self) - def update_objects(imageDetections: ImageDetections3D): + def update_objects(imageDetections: ImageDetections3DPC): for detection in imageDetections.detections: # print(detection) return self.add_detection(detection) diff --git a/dimos/perception/detection/person_tracker.py b/dimos/perception/detection/person_tracker.py new file mode 100644 index 0000000000..fe69fbc15e --- /dev/null +++ b/dimos/perception/detection/person_tracker.py @@ -0,0 +1,116 @@ +# Copyright 2025 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 typing import Tuple + +from reactivex import operators as ops +from reactivex.observable import Observable + +from dimos.core import In, Module, Out, rpc +from dimos.msgs.geometry_msgs import PoseStamped, Transform, Vector3 +from dimos.msgs.sensor_msgs import CameraInfo, Image +from dimos.msgs.vision_msgs import Detection2DArray +from dimos.perception.detection.type import ImageDetections2D +from dimos.types.timestamped import align_timestamped +from dimos.utils.reactive import backpressure + + +class PersonTracker(Module): + detections: In[Detection2DArray] = None # type: ignore + image: In[Image] = None # type: ignore + target: Out[PoseStamped] = None # type: ignore + + camera_info: CameraInfo + + def __init__(self, cameraInfo: CameraInfo, **kwargs): + super().__init__(**kwargs) + self.camera_info = cameraInfo + + def center_to_3d( + self, + pixel: Tuple[int, int], + camera_info: CameraInfo, + assumed_depth: float = 1.0, + ) -> Vector3: + """Unproject 2D pixel coordinates to 3D position in camera_link frame. + + Args: + camera_info: Camera calibration information + assumed_depth: Assumed depth in meters (default 1.0m from camera) + + Returns: + Vector3 position in camera_link frame coordinates (Z up, X forward) + """ + # Extract camera intrinsics + fx, fy = camera_info.K[0], camera_info.K[4] + cx, cy = camera_info.K[2], camera_info.K[5] + + # Unproject pixel to normalized camera coordinates + x_norm = (pixel[0] - cx) / fx + y_norm = (pixel[1] - cy) / fy + + # Create 3D point at assumed depth in camera optical frame + # Camera optical frame: X right, Y down, Z forward + x_optical = x_norm * assumed_depth + y_optical = y_norm * assumed_depth + z_optical = assumed_depth + + # Transform from camera optical frame to camera_link frame + # Optical: X right, Y down, Z forward + # Link: X forward, Y left, Z up + # Transformation: x_link = z_optical, y_link = -x_optical, z_link = -y_optical + return Vector3(z_optical, -x_optical, -y_optical) + + def detections_stream(self) -> Observable[ImageDetections2D]: + return backpressure( + align_timestamped( + self.image.pure_observable(), + self.detections.pure_observable().pipe( + ops.filter(lambda d: d.detections_length > 0) # type: ignore[attr-defined] + ), + match_tolerance=0.0, + buffer_size=2.0, + ).pipe(ops.map(lambda pair: ImageDetections2D.from_ros_detection2d_array(*pair))) + ) + + @rpc + def start(self): + self.detections_stream().subscribe(self.track) + + @rpc + def stop(self): + super().stop() + + def track(self, detections2D: ImageDetections2D): + if len(detections2D) == 0: + return + + target = max(detections2D.detections, key=lambda det: det.bbox_2d_volume()) + vector = self.center_to_3d(target.center_bbox, self.camera_info, 2.0) + + pose_in_camera = PoseStamped( + ts=detections2D.ts, + position=vector, + frame_id="camera_link", + ) + + tf_world_to_camera = self.tf.get("world", "camera_link", detections2D.ts, 5.0) + if not tf_world_to_camera: + return + + tf_camera_to_target = Transform.from_pose("target", pose_in_camera) + tf_world_to_target = tf_world_to_camera + tf_camera_to_target + pose_in_world = tf_world_to_target.to_pose(ts=detections2D.ts) + + self.target.publish(pose_in_world) diff --git a/dimos/perception/detection/reid/__init__.py b/dimos/perception/detection/reid/__init__.py new file mode 100644 index 0000000000..b76741a7eb --- /dev/null +++ b/dimos/perception/detection/reid/__init__.py @@ -0,0 +1,13 @@ +from dimos.perception.detection.reid.module import Config, ReidModule +from dimos.perception.detection.reid.embedding_id_system import EmbeddingIDSystem +from dimos.perception.detection.reid.type import IDSystem, PassthroughIDSystem + +__all__ = [ + # ID Systems + "IDSystem", + "PassthroughIDSystem", + "EmbeddingIDSystem", + # Module + "ReidModule", + "Config", +] diff --git a/dimos/perception/detection/reid/embedding_id_system.py b/dimos/perception/detection/reid/embedding_id_system.py new file mode 100644 index 0000000000..7fb0a2ba40 --- /dev/null +++ b/dimos/perception/detection/reid/embedding_id_system.py @@ -0,0 +1,263 @@ +# Copyright 2025 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 typing import Callable, Dict, List, Literal, Set + +import numpy as np + +from dimos.models.embedding.base import Embedding, EmbeddingModel +from dimos.perception.detection.reid.type import IDSystem +from dimos.perception.detection.type import Detection2DBBox + + +class EmbeddingIDSystem(IDSystem): + """Associates short-term track_ids to long-term unique detection IDs via embedding similarity. + + Maintains: + - All embeddings per track_id (as numpy arrays) for robust group comparison + - Negative constraints from co-occurrence (tracks in same frame = different objects) + - Mapping from track_id to unique long-term ID + """ + + def __init__( + self, + model: Callable[[], EmbeddingModel[Embedding]], + padding: int = 0, + similarity_threshold: float = 0.63, + comparison_mode: Literal["max", "mean", "top_k_mean"] = "top_k_mean", + top_k: int = 30, + max_embeddings_per_track: int = 500, + min_embeddings_for_matching: int = 10, + ): + """Initialize track associator. + + Args: + model: Callable (class or function) that returns an embedding model for feature extraction + padding: Padding to add around detection bbox when cropping (default: 0) + similarity_threshold: Minimum similarity for associating tracks (0-1) + comparison_mode: How to aggregate similarities between embedding groups + - "max": Use maximum similarity between any pair + - "mean": Use mean of all pairwise similarities + - "top_k_mean": Use mean of top-k similarities + top_k: Number of top similarities to average (if using top_k_mean) + max_embeddings_per_track: Maximum number of embeddings to keep per track + min_embeddings_for_matching: Minimum embeddings before attempting to match tracks + """ + # Call model factory (class or function) to get model instance + self.model = model() + + # Call warmup if available + if hasattr(self.model, "warmup"): + self.model.warmup() + + self.padding = padding + self.similarity_threshold = similarity_threshold + self.comparison_mode = comparison_mode + self.top_k = top_k + self.max_embeddings_per_track = max_embeddings_per_track + self.min_embeddings_for_matching = min_embeddings_for_matching + + # Track embeddings (list of all embeddings as numpy arrays) + self.track_embeddings: Dict[int, List[np.ndarray]] = {} + + # Negative constraints (track_ids that co-occurred = different objects) + self.negative_pairs: Dict[int, Set[int]] = {} + + # Track ID to long-term unique ID mapping + self.track_to_long_term: Dict[int, int] = {} + self.long_term_counter: int = 0 + + # Similarity history for optional adaptive thresholding + self.similarity_history: List[float] = [] + + def register_detection(self, detection: Detection2DBBox) -> int: + """ + Register detection and return long-term ID. + + Args: + detection: Detection to register + + Returns: + Long-term unique ID for this detection + """ + # Extract embedding from detection's cropped image + cropped_image = detection.cropped_image(padding=self.padding) + embedding = self.model.embed(cropped_image) + assert not isinstance(embedding, list), "Expected single embedding for single image" + # Move embedding to CPU immediately to free GPU memory + embedding = embedding.to_cpu() + + # Update and associate track + self.update_embedding(detection.track_id, embedding) + return self.associate(detection.track_id) + + def update_embedding(self, track_id: int, new_embedding: Embedding) -> None: + """Add new embedding to track's embedding collection. + + Args: + track_id: Short-term track ID from detector + new_embedding: New embedding to add to collection + """ + # Convert to numpy array (already on CPU from feature extractor) + new_vec = new_embedding.to_numpy() + + # Ensure normalized for cosine similarity + norm = np.linalg.norm(new_vec) + if norm > 0: + new_vec = new_vec / norm + + if track_id not in self.track_embeddings: + self.track_embeddings[track_id] = [] + + embeddings = self.track_embeddings[track_id] + embeddings.append(new_vec) + + # Keep only most recent embeddings if limit exceeded + if len(embeddings) > self.max_embeddings_per_track: + embeddings.pop(0) # Remove oldest + + def _compute_group_similarity( + self, query_embeddings: List[np.ndarray], candidate_embeddings: List[np.ndarray] + ) -> float: + """Compute similarity between two groups of embeddings. + + Args: + query_embeddings: List of embeddings for query track + candidate_embeddings: List of embeddings for candidate track + + Returns: + Aggregated similarity score + """ + # Compute all pairwise similarities efficiently + query_matrix = np.stack(query_embeddings) # [M, D] + candidate_matrix = np.stack(candidate_embeddings) # [N, D] + + # Cosine similarity via matrix multiplication (already normalized) + similarities = query_matrix @ candidate_matrix.T # [M, N] + + if self.comparison_mode == "max": + # Maximum similarity across all pairs + return float(np.max(similarities)) + + elif self.comparison_mode == "mean": + # Mean of all pairwise similarities + return float(np.mean(similarities)) + + elif self.comparison_mode == "top_k_mean": + # Mean of top-k similarities + flat_sims = similarities.flatten() + k = min(self.top_k, len(flat_sims)) + top_k_sims = np.partition(flat_sims, -k)[-k:] + return float(np.mean(top_k_sims)) + + else: + raise ValueError(f"Unknown comparison mode: {self.comparison_mode}") + + def add_negative_constraints(self, track_ids: List[int]) -> None: + """Record that these track_ids co-occurred in same frame (different objects). + + Args: + track_ids: List of track_ids present in current frame + """ + # All pairs of track_ids in same frame can't be same object + for i, tid1 in enumerate(track_ids): + for tid2 in track_ids[i + 1 :]: + self.negative_pairs.setdefault(tid1, set()).add(tid2) + self.negative_pairs.setdefault(tid2, set()).add(tid1) + + def associate(self, track_id: int) -> int: + """Associate track_id to long-term unique detection ID. + + Args: + track_id: Short-term track ID to associate + + Returns: + Long-term unique detection ID + """ + # Already has assignment + if track_id in self.track_to_long_term: + return self.track_to_long_term[track_id] + + # Need embeddings to compare + if track_id not in self.track_embeddings or not self.track_embeddings[track_id]: + # Create new ID if no embeddings yet + new_id = self.long_term_counter + self.long_term_counter += 1 + self.track_to_long_term[track_id] = new_id + return new_id + + # Get query embeddings + query_embeddings = self.track_embeddings[track_id] + + # Don't attempt matching until we have enough embeddings for the query track + if len(query_embeddings) < self.min_embeddings_for_matching: + # Not ready yet - return -1 + return -1 + + # Build candidate list (only tracks with assigned long_term_ids) + best_similarity = -1.0 + best_track_id = None + + for other_tid, other_embeddings in self.track_embeddings.items(): + # Skip self + if other_tid == track_id: + continue + + # Skip if negative constraint (co-occurred) + if other_tid in self.negative_pairs.get(track_id, set()): + continue + + # Skip if no long_term_id yet + if other_tid not in self.track_to_long_term: + continue + + # Skip if not enough embeddings + if len(other_embeddings) < self.min_embeddings_for_matching: + continue + + # Compute group similarity + similarity = self._compute_group_similarity(query_embeddings, other_embeddings) + + if similarity > best_similarity: + best_similarity = similarity + best_track_id = other_tid + + # Check if best match exceeds threshold + if best_track_id is not None and best_similarity >= self.similarity_threshold: + matched_long_term_id = self.track_to_long_term[best_track_id] + print( + f"Track {track_id}: matched with track {best_track_id} " + f"(long_term_id={matched_long_term_id}, similarity={best_similarity:.4f}, " + f"mode={self.comparison_mode}, embeddings: {len(query_embeddings)} vs {len(self.track_embeddings[best_track_id])}), threshold: {self.similarity_threshold}" + ) + + # Track similarity history + self.similarity_history.append(best_similarity) + + # Associate with existing long_term_id + self.track_to_long_term[track_id] = matched_long_term_id + return matched_long_term_id + + # Create new unique detection ID + new_id = self.long_term_counter + self.long_term_counter += 1 + self.track_to_long_term[track_id] = new_id + + if best_track_id is not None: + print( + f"Track {track_id}: creating new ID {new_id} " + f"(best similarity={best_similarity:.4f} with id={self.track_to_long_term[best_track_id]} below threshold={self.similarity_threshold})" + ) + + return new_id diff --git a/dimos/perception/detection/reid/module.py b/dimos/perception/detection/reid/module.py new file mode 100644 index 0000000000..64769b1038 --- /dev/null +++ b/dimos/perception/detection/reid/module.py @@ -0,0 +1,112 @@ +# Copyright 2025 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 dimos_lcm.foxglove_msgs.ImageAnnotations import ( + ImageAnnotations, + TextAnnotation, +) +from dimos_lcm.foxglove_msgs.Point2 import Point2 +from reactivex import operators as ops +from reactivex.observable import Observable + +from dimos.core import In, Module, ModuleConfig, Out, rpc +from dimos.msgs.foxglove_msgs.Color import Color +from dimos.msgs.sensor_msgs import Image +from dimos.msgs.vision_msgs import Detection2DArray +from dimos.perception.detection.reid.embedding_id_system import EmbeddingIDSystem +from dimos.perception.detection.reid.type import IDSystem +from dimos.perception.detection.type import ImageDetections2D +from dimos.types.timestamped import align_timestamped, to_ros_stamp +from dimos.utils.reactive import backpressure + + +class Config(ModuleConfig): + idsystem: IDSystem + + +class ReidModule(Module): + default_config = Config + + detections: In[Detection2DArray] = None # type: ignore + image: In[Image] = None # type: ignore + annotations: Out[ImageAnnotations] = None # type: ignore + + def __init__(self, idsystem: IDSystem | None = None, **kwargs): + super().__init__(**kwargs) + if idsystem is None: + try: + from dimos.models.embedding import TorchReIDModel + + idsystem = EmbeddingIDSystem(model=TorchReIDModel, padding=0) + except Exception as e: + raise RuntimeError( + "TorchReIDModel not available. Please install with: pip install dimos[torchreid]" + ) from e + + self.idsystem = idsystem + + def detections_stream(self) -> Observable[ImageDetections2D]: + return backpressure( + align_timestamped( + self.image.pure_observable(), + self.detections.pure_observable().pipe( + ops.filter(lambda d: d.detections_length > 0) # type: ignore[attr-defined] + ), + match_tolerance=0.0, + buffer_size=2.0, + ).pipe(ops.map(lambda pair: ImageDetections2D.from_ros_detection2d_array(*pair))) # type: ignore[misc] + ) + + @rpc + def start(self): + self.detections_stream().subscribe(self.ingress) + + @rpc + def stop(self): + super().stop() + + def ingress(self, imageDetections: ImageDetections2D): + text_annotations = [] + + for detection in imageDetections: + # Register detection and get long-term ID + long_term_id = self.idsystem.register_detection(detection) + + # Skip annotation if not ready yet (long_term_id == -1) + if long_term_id == -1: + continue + + # Create text annotation for long_term_id above the detection + x1, y1, _, _ = detection.bbox + font_size = imageDetections.image.width / 60 + + text_annotations.append( + TextAnnotation( + timestamp=to_ros_stamp(detection.ts), + position=Point2(x=x1, y=y1 - font_size * 1.5), + text=f"PERSON: {long_term_id}", + font_size=font_size, + text_color=Color(r=0.0, g=1.0, b=1.0, a=1.0), # Cyan + background_color=Color(r=0.0, g=0.0, b=0.0, a=0.8), + ) + ) + + # Publish annotations (even if empty to clear previous annotations) + annotations = ImageAnnotations( + texts=text_annotations, + texts_length=len(text_annotations), + points=[], + points_length=0, + ) + self.annotations.publish(annotations) diff --git a/dimos/perception/detection/reid/test_embedding_id_system.py b/dimos/perception/detection/reid/test_embedding_id_system.py new file mode 100644 index 0000000000..b2bc84bc55 --- /dev/null +++ b/dimos/perception/detection/reid/test_embedding_id_system.py @@ -0,0 +1,270 @@ +# Copyright 2025 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. + +import pytest +import torch + +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.reid.embedding_id_system import EmbeddingIDSystem +from dimos.utils.data import get_data + + +@pytest.fixture(scope="session") +def mobileclip_model(): + """Load MobileCLIP model once for all tests.""" + from dimos.models.embedding.mobileclip import MobileCLIPModel + + model_path = get_data("models_mobileclip") / "mobileclip2_s0.pt" + model = MobileCLIPModel(model_name="MobileCLIP2-S0", model_path=model_path) + model.warmup() + return model + + +@pytest.fixture +def track_associator(mobileclip_model): + """Create fresh EmbeddingIDSystem for each test.""" + return EmbeddingIDSystem(model=lambda: mobileclip_model, similarity_threshold=0.75) + + +@pytest.fixture(scope="session") +def test_image(): + """Load test image.""" + return Image.from_file(get_data("cafe.jpg")).to_rgb() + + +@pytest.mark.gpu +def test_update_embedding_single(track_associator, mobileclip_model, test_image): + """Test updating embedding for a single track.""" + embedding = mobileclip_model.embed(test_image) + + # First update + track_associator.update_embedding(track_id=1, new_embedding=embedding) + + assert 1 in track_associator.track_embeddings + assert track_associator.embedding_counts[1] == 1 + + # Verify embedding is on device and normalized + emb_vec = track_associator.track_embeddings[1] + assert isinstance(emb_vec, torch.Tensor) + assert emb_vec.device.type in ["cuda", "cpu"] + norm = torch.norm(emb_vec).item() + assert abs(norm - 1.0) < 0.01, "Embedding should be normalized" + + +@pytest.mark.gpu +def test_update_embedding_running_average(track_associator, mobileclip_model, test_image): + """Test running average of embeddings.""" + embedding1 = mobileclip_model.embed(test_image) + embedding2 = mobileclip_model.embed(test_image) + + # Add first embedding + track_associator.update_embedding(track_id=1, new_embedding=embedding1) + first_vec = track_associator.track_embeddings[1].clone() + + # Add second embedding (same image, should be very similar) + track_associator.update_embedding(track_id=1, new_embedding=embedding2) + avg_vec = track_associator.track_embeddings[1] + + assert track_associator.embedding_counts[1] == 2 + + # Average should still be normalized + norm = torch.norm(avg_vec).item() + assert abs(norm - 1.0) < 0.01, "Average embedding should be normalized" + + # Average should be similar to both originals (same image) + similarity1 = (first_vec @ avg_vec).item() + assert similarity1 > 0.99, "Average should be very similar to original" + + +@pytest.mark.gpu +def test_negative_constraints(track_associator): + """Test negative constraint recording.""" + # Simulate frame with 3 tracks + track_ids = [1, 2, 3] + track_associator.add_negative_constraints(track_ids) + + # Check that all pairs are recorded + assert 2 in track_associator.negative_pairs[1] + assert 3 in track_associator.negative_pairs[1] + assert 1 in track_associator.negative_pairs[2] + assert 3 in track_associator.negative_pairs[2] + assert 1 in track_associator.negative_pairs[3] + assert 2 in track_associator.negative_pairs[3] + + +@pytest.mark.gpu +def test_associate_new_track(track_associator, mobileclip_model, test_image): + """Test associating a new track creates new long_term_id.""" + embedding = mobileclip_model.embed(test_image) + track_associator.update_embedding(track_id=1, new_embedding=embedding) + + # First association should create new long_term_id + long_term_id = track_associator.associate(track_id=1) + + assert long_term_id == 0, "First track should get long_term_id=0" + assert track_associator.track_to_long_term[1] == 0 + assert track_associator.long_term_counter == 1 + + +@pytest.mark.gpu +def test_associate_similar_tracks(track_associator, mobileclip_model, test_image): + """Test associating similar tracks to same long_term_id.""" + # Create embeddings from same image (should be very similar) + embedding1 = mobileclip_model.embed(test_image) + embedding2 = mobileclip_model.embed(test_image) + + # Add first track + track_associator.update_embedding(track_id=1, new_embedding=embedding1) + long_term_id_1 = track_associator.associate(track_id=1) + + # Add second track with similar embedding + track_associator.update_embedding(track_id=2, new_embedding=embedding2) + long_term_id_2 = track_associator.associate(track_id=2) + + # Should get same long_term_id (similarity > 0.75) + assert long_term_id_1 == long_term_id_2, "Similar tracks should get same long_term_id" + assert track_associator.long_term_counter == 1, "Only one long_term_id should be created" + + +@pytest.mark.gpu +def test_associate_with_negative_constraint(track_associator, mobileclip_model, test_image): + """Test that negative constraints prevent association.""" + # Create similar embeddings + embedding1 = mobileclip_model.embed(test_image) + embedding2 = mobileclip_model.embed(test_image) + + # Add first track + track_associator.update_embedding(track_id=1, new_embedding=embedding1) + long_term_id_1 = track_associator.associate(track_id=1) + + # Add negative constraint (tracks co-occurred) + track_associator.add_negative_constraints([1, 2]) + + # Add second track with similar embedding + track_associator.update_embedding(track_id=2, new_embedding=embedding2) + long_term_id_2 = track_associator.associate(track_id=2) + + # Should get different long_term_ids despite high similarity + assert long_term_id_1 != long_term_id_2, ( + "Co-occurring tracks should get different long_term_ids" + ) + assert track_associator.long_term_counter == 2, "Two long_term_ids should be created" + + +@pytest.mark.gpu +def test_associate_different_objects(track_associator, mobileclip_model, test_image): + """Test that dissimilar embeddings get different long_term_ids.""" + # Create embeddings for image and text (very different) + image_emb = mobileclip_model.embed(test_image) + text_emb = mobileclip_model.embed_text("a dog") + + # Add first track (image) + track_associator.update_embedding(track_id=1, new_embedding=image_emb) + long_term_id_1 = track_associator.associate(track_id=1) + + # Add second track (text - very different embedding) + track_associator.update_embedding(track_id=2, new_embedding=text_emb) + long_term_id_2 = track_associator.associate(track_id=2) + + # Should get different long_term_ids (similarity < 0.75) + assert long_term_id_1 != long_term_id_2, "Different objects should get different long_term_ids" + assert track_associator.long_term_counter == 2 + + +@pytest.mark.gpu +def test_associate_returns_cached(track_associator, mobileclip_model, test_image): + """Test that repeated calls return same long_term_id.""" + embedding = mobileclip_model.embed(test_image) + track_associator.update_embedding(track_id=1, new_embedding=embedding) + + # First call + long_term_id_1 = track_associator.associate(track_id=1) + + # Second call should return cached result + long_term_id_2 = track_associator.associate(track_id=1) + + assert long_term_id_1 == long_term_id_2 + assert track_associator.long_term_counter == 1, "Should not create new ID" + + +@pytest.mark.gpu +def test_associate_not_ready(track_associator): + """Test that associate returns -1 for track without embedding.""" + long_term_id = track_associator.associate(track_id=999) + assert long_term_id == -1, "Should return -1 for track without embedding" + + +@pytest.mark.gpu +def test_gpu_performance(track_associator, mobileclip_model, test_image): + """Test that embeddings stay on GPU for performance.""" + embedding = mobileclip_model.embed(test_image) + track_associator.update_embedding(track_id=1, new_embedding=embedding) + + # Embedding should stay on device + emb_vec = track_associator.track_embeddings[1] + assert isinstance(emb_vec, torch.Tensor) + # Device comparison (handle "cuda" vs "cuda:0") + expected_device = mobileclip_model.device + assert emb_vec.device.type == torch.device(expected_device).type + + # Running average should happen on GPU + embedding2 = mobileclip_model.embed(test_image) + track_associator.update_embedding(track_id=1, new_embedding=embedding2) + + avg_vec = track_associator.track_embeddings[1] + assert avg_vec.device.type == torch.device(expected_device).type + + +@pytest.mark.gpu +def test_similarity_threshold_configurable(mobileclip_model): + """Test that similarity threshold is configurable.""" + associator_strict = EmbeddingIDSystem(model=lambda: mobileclip_model, similarity_threshold=0.95) + associator_loose = EmbeddingIDSystem(model=lambda: mobileclip_model, similarity_threshold=0.50) + + assert associator_strict.similarity_threshold == 0.95 + assert associator_loose.similarity_threshold == 0.50 + + +@pytest.mark.gpu +def test_multi_track_scenario(track_associator, mobileclip_model, test_image): + """Test realistic scenario with multiple tracks across frames.""" + # Frame 1: Track 1 appears + emb1 = mobileclip_model.embed(test_image) + track_associator.update_embedding(1, emb1) + track_associator.add_negative_constraints([1]) + lt1 = track_associator.associate(1) + + # Frame 2: Track 1 and Track 2 appear (different objects) + text_emb = mobileclip_model.embed_text("a dog") + track_associator.update_embedding(1, emb1) # Update average + track_associator.update_embedding(2, text_emb) + track_associator.add_negative_constraints([1, 2]) # Co-occur = different + lt2 = track_associator.associate(2) + + # Track 2 should get different ID despite any similarity + assert lt1 != lt2 + + # Frame 3: Track 1 disappears, Track 3 appears (same as Track 1) + emb3 = mobileclip_model.embed(test_image) + track_associator.update_embedding(3, emb3) + track_associator.add_negative_constraints([2, 3]) + lt3 = track_associator.associate(3) + + # Track 3 should match Track 1 (not co-occurring, similar embedding) + assert lt3 == lt1 + + print("\nMulti-track scenario results:") + print(f" Track 1 -> long_term_id {lt1}") + print(f" Track 2 -> long_term_id {lt2} (different object, co-occurred)") + print(f" Track 3 -> long_term_id {lt3} (re-identified as Track 1)") diff --git a/dimos/perception/detection/reid/test_module.py b/dimos/perception/detection/reid/test_module.py new file mode 100644 index 0000000000..6c977e13a5 --- /dev/null +++ b/dimos/perception/detection/reid/test_module.py @@ -0,0 +1,44 @@ +# Copyright 2025 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. + +import pytest + +from dimos.core import LCMTransport +from dimos.msgs.foxglove_msgs import ImageAnnotations +from dimos.perception.detection.reid.embedding_id_system import EmbeddingIDSystem +from dimos.perception.detection.reid.module import ReidModule + + +@pytest.mark.tool +def test_reid_ingress(imageDetections2d): + try: + from dimos.models.embedding import TorchReIDModel + except Exception: + pytest.skip("TorchReIDModel not available") + + # Create TorchReID-based IDSystem for testing + reid_model = TorchReIDModel(model_name="osnet_x1_0") + reid_model.warmup() + idsystem = EmbeddingIDSystem( + model=lambda: reid_model, + padding=20, + similarity_threshold=0.75, + ) + + reid_module = ReidModule(idsystem=idsystem, warmup=False) + print("Processing detections through ReidModule...") + reid_module.annotations._transport = LCMTransport("/annotations", ImageAnnotations) + reid_module.ingress(imageDetections2d) + reid_module._close_module() + print("āœ“ ReidModule ingress test completed successfully") diff --git a/dimos/perception/detection/reid/type.py b/dimos/perception/detection/reid/type.py new file mode 100644 index 0000000000..0ef2da961c --- /dev/null +++ b/dimos/perception/detection/reid/type.py @@ -0,0 +1,50 @@ +# Copyright 2025 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 abc import ABC, abstractmethod + +from dimos.perception.detection.type import Detection2DBBox, ImageDetections2D + + +class IDSystem(ABC): + """Abstract base class for ID assignment systems.""" + + def register_detections(self, detections: ImageDetections2D) -> None: + """Register multiple detections.""" + for detection in detections.detections: + if isinstance(detection, Detection2DBBox): + self.register_detection(detection) + + @abstractmethod + def register_detection(self, detection: Detection2DBBox) -> int: + """ + Register a single detection, returning assigned (long term) ID. + + Args: + detection: Detection to register + + Returns: + Long-term unique ID for this detection + """ + ... + + +class PassthroughIDSystem(IDSystem): + """Simple ID system that returns track_id with no object permanence.""" + + def register_detection(self, detection: Detection2DBBox) -> int: + """Return detection's track_id as long-term ID (no permanence).""" + return detection.track_id diff --git a/dimos/perception/detection2d/test_moduleDB.py b/dimos/perception/detection/test_moduleDB.py similarity index 97% rename from dimos/perception/detection2d/test_moduleDB.py rename to dimos/perception/detection/test_moduleDB.py index a3a1b003fd..1ede53f172 100644 --- a/dimos/perception/detection2d/test_moduleDB.py +++ b/dimos/perception/detection/test_moduleDB.py @@ -21,7 +21,7 @@ from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import Image, PointCloud2 from dimos.msgs.vision_msgs import Detection2DArray -from dimos.perception.detection2d.moduleDB import ObjectDBModule +from dimos.perception.detection.moduleDB import ObjectDBModule from dimos.protocol.service import lcmservice as lcm from dimos.robot.unitree_webrtc.modular import deploy_connection, deploy_navigation from dimos.robot.unitree_webrtc.modular.connection_module import ConnectionModule diff --git a/dimos/perception/detection/type/__init__.py b/dimos/perception/detection/type/__init__.py new file mode 100644 index 0000000000..d8f36d79dc --- /dev/null +++ b/dimos/perception/detection/type/__init__.py @@ -0,0 +1,41 @@ +from dimos.perception.detection.type.detection2d import ( + Detection2D, + Detection2DBBox, + Detection2DPerson, + ImageDetections2D, +) +from dimos.perception.detection.type.detection3d import ( + Detection3D, + Detection3DBBox, + Detection3DPC, + ImageDetections3DPC, + PointCloudFilter, + height_filter, + radius_outlier, + raycast, + statistical, +) +from dimos.perception.detection.type.imageDetections import ImageDetections +from dimos.perception.detection.type.utils import TableStr + +__all__ = [ + # 2D Detection types + "Detection2D", + "Detection2DBBox", + "Detection2DPerson", + "ImageDetections2D", + # 3D Detection types + "Detection3D", + "Detection3DBBox", + "Detection3DPC", + "ImageDetections3DPC", + # Point cloud filters + "PointCloudFilter", + "height_filter", + "radius_outlier", + "raycast", + "statistical", + # Base types + "ImageDetections", + "TableStr", +] diff --git a/dimos/perception/detection/type/detection2d/__init__.py b/dimos/perception/detection/type/detection2d/__init__.py new file mode 100644 index 0000000000..1096abda9c --- /dev/null +++ b/dimos/perception/detection/type/detection2d/__init__.py @@ -0,0 +1,25 @@ +# Copyright 2025 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 dimos.perception.detection.type.detection2d.base import Detection2D +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D +from dimos.perception.detection.type.detection2d.person import Detection2DPerson + +__all__ = [ + "Detection2D", + "Detection2DBBox", + "ImageDetections2D", + "Detection2DPerson", +] diff --git a/dimos/perception/detection/type/detection2d/base.py b/dimos/perception/detection/type/detection2d/base.py new file mode 100644 index 0000000000..e89bf65409 --- /dev/null +++ b/dimos/perception/detection/type/detection2d/base.py @@ -0,0 +1,52 @@ +# Copyright 2025 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 abc import abstractmethod +from typing import List + +from dimos_lcm.foxglove_msgs.ImageAnnotations import PointsAnnotation, TextAnnotation +from dimos_lcm.vision_msgs import Detection2D as ROSDetection2D + +from dimos.msgs.foxglove_msgs import ImageAnnotations +from dimos.msgs.sensor_msgs import Image +from dimos.types.timestamped import Timestamped + + +class Detection2D(Timestamped): + """Abstract base class for 2D detections.""" + + @abstractmethod + def cropped_image(self, padding: int = 20) -> Image: + """Return a cropped version of the image focused on the detection area.""" + ... + + @abstractmethod + def to_image_annotations(self) -> ImageAnnotations: + """Convert detection to Foxglove ImageAnnotations for visualization.""" + ... + + @abstractmethod + def to_text_annotation(self) -> List[TextAnnotation]: + """Return text annotations for visualization.""" + ... + + @abstractmethod + def to_points_annotation(self) -> List[PointsAnnotation]: + """Return points/shape annotations for visualization.""" + ... + + @abstractmethod + def to_ros_detection2d(self) -> ROSDetection2D: + """Convert detection to ROS Detection2D message.""" + ... diff --git a/dimos/perception/detection2d/type/detection2d.py b/dimos/perception/detection/type/detection2d/bbox.py similarity index 64% rename from dimos/perception/detection2d/type/detection2d.py rename to dimos/perception/detection/type/detection2d/bbox.py index 48e1a5191d..223e1bc018 100644 --- a/dimos/perception/detection2d/type/detection2d.py +++ b/dimos/perception/detection/type/detection2d/bbox.py @@ -15,9 +15,11 @@ from __future__ import annotations import hashlib -from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, List, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Sequence, Tuple, Union + +if TYPE_CHECKING: + from dimos.perception.detection.type.detection2d.person import Detection2DPerson from dimos_lcm.foxglove_msgs.ImageAnnotations import ( PointsAnnotation, @@ -36,26 +38,19 @@ ) from rich.console import Console from rich.text import Text +from ultralytics.engine.results import Results from dimos.msgs.foxglove_msgs import ImageAnnotations from dimos.msgs.foxglove_msgs.Color import Color from dimos.msgs.sensor_msgs import Image from dimos.msgs.std_msgs import Header -from dimos.perception.detection2d.type.imageDetections import ImageDetections -from dimos.types.timestamped import Timestamped, to_ros_stamp, to_timestamp - -if TYPE_CHECKING: - from dimos.perception.detection2d.type.person import Person +from dimos.perception.detection.type.detection2d.base import Detection2D +from dimos.types.timestamped import to_ros_stamp, to_timestamp +from dimos.utils.decorators.decorators import simple_mcache Bbox = Tuple[float, float, float, float] CenteredBbox = Tuple[float, float, float, float] -# yolo and detic have bad output formats -InconvinientDetectionFormat = Tuple[List[Bbox], List[int], List[int], List[float], List[str]] - -Detection = Tuple[Bbox, int, int, float, str] -Detections = List[Detection] - def _hash_to_color(name: str) -> str: """Generate a consistent color for a given name using hash.""" @@ -83,25 +78,6 @@ def _hash_to_color(name: str) -> str: return colors[hash_value % len(colors)] -# yolo and detic have bad formats this translates into list of detections -def better_detection_format(inconvinient_detections: InconvinientDetectionFormat) -> Detections: - bboxes, track_ids, class_ids, confidences, names = inconvinient_detections - return [ - (bbox, track_id, class_id, confidence, name if name else "") - for bbox, track_id, class_id, confidence, name in zip( - bboxes, track_ids, class_ids, confidences, names - ) - ] - - -class Detection2D(Timestamped): - @abstractmethod - def cropped_image(self, padding: int = 20) -> Image: ... - - @abstractmethod - def to_image_annotations(self) -> ImageAnnotations: ... - - @dataclass class Detection2DBBox(Detection2D): bbox: Bbox @@ -123,6 +99,33 @@ def to_repr_dict(self) -> Dict[str, Any]: "bbox": f"[{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}]", } + def center_to_3d( + self, + pixel: Tuple[int, int], + camera_info: CameraInfo, + assumed_depth: float = 1.0, + ) -> PoseStamped: + """Unproject 2D pixel coordinates to 3D position in camera optical frame. + + Args: + camera_info: Camera calibration information + assumed_depth: Assumed depth in meters (default 1.0m from camera) + + Returns: + Vector3 position in camera optical frame coordinates + """ + # Extract camera intrinsics + fx, fy = camera_info.K[0], camera_info.K[4] + cx, cy = camera_info.K[2], camera_info.K[5] + + # Unproject pixel to normalized camera coordinates + x_norm = (pixel[0] - cx) / fx + y_norm = (pixel[1] - cy) / fy + + # Create 3D point at assumed depth in camera optical frame + # Camera optical frame: X right, Y down, Z forward + return Vector3(x_norm * assumed_depth, y_norm * assumed_depth, assumed_depth) + # return focused image, only on the bbox def cropped_image(self, padding: int = 20) -> Image: """Return a cropped version of the image focused on the bounding box. @@ -162,23 +165,82 @@ def __str__(self): console.print(*parts, end="") return capture.get().strip() + @property + def center_bbox(self) -> Tuple[float, float]: + """Get center point of bounding box.""" + x1, y1, x2, y2 = self.bbox + return ((x1 + x2) / 2, (y1 + y2) / 2) + def bbox_2d_volume(self) -> float: x1, y1, x2, y2 = self.bbox width = max(0.0, x2 - x1) height = max(0.0, y2 - y1) return width * height - @classmethod - def from_detector( - cls, raw_detections: InconvinientDetectionFormat, **kwargs - ) -> List["Detection2D"]: - return [ - cls.from_detection(raw, **kwargs) for raw in better_detection_format(raw_detections) - ] + @simple_mcache + def is_valid(self) -> bool: + """Check if detection bbox is valid. + + Validates that: + - Bounding box has positive dimensions + - Bounding box is within image bounds (if image has shape) + + Returns: + True if bbox is valid, False otherwise + """ + x1, y1, x2, y2 = self.bbox + + # Check positive dimensions + if x2 <= x1 or y2 <= y1: + return False + + # Check if within image bounds (if image has shape) + if self.image.shape: + h, w = self.image.shape[:2] + if not (0 <= x1 <= w and 0 <= y1 <= h and 0 <= x2 <= w and 0 <= y2 <= h): + return False + + return True @classmethod - def from_detection(cls, raw_detection: Detection, **kwargs) -> "Detection2D": - bbox, track_id, class_id, confidence, name = raw_detection + def from_ultralytics_result(cls, result: Results, idx: int, image: Image) -> "Detection2DBBox": + """Create Detection2DBBox from ultralytics Results object. + + Args: + result: Ultralytics Results object containing detection data + idx: Index of the detection in the results + image: Source image + + Returns: + Detection2DBBox instance + """ + if result.boxes is None: + raise ValueError("Result has no boxes") + + # Extract bounding box coordinates + bbox_array = result.boxes.xyxy[idx].cpu().numpy() + bbox: Bbox = ( + float(bbox_array[0]), + float(bbox_array[1]), + float(bbox_array[2]), + float(bbox_array[3]), + ) + + # Extract confidence + confidence = float(result.boxes.conf[idx].cpu()) + + # Extract class ID and name + class_id = int(result.boxes.cls[idx].cpu()) + name = ( + result.names.get(class_id, f"class_{class_id}") + if hasattr(result, "names") + else f"class_{class_id}" + ) + + # Extract track ID if available + track_id = -1 + if hasattr(result.boxes, "id") and result.boxes.id is not None: + track_id = int(result.boxes.id[idx].cpu()) return cls( bbox=bbox, @@ -186,7 +248,8 @@ def from_detection(cls, raw_detection: Detection, **kwargs) -> "Detection2D": class_id=class_id, confidence=confidence, name=name, - **kwargs, + ts=image.ts, + image=image, ) def get_bbox_center(self) -> CenteredBbox: @@ -214,37 +277,53 @@ def lcm_encode(self): def to_text_annotation(self) -> List[TextAnnotation]: x1, y1, x2, y2 = self.bbox - font_size = 20 + font_size = self.image.width / 80 - return [ - TextAnnotation( - timestamp=to_ros_stamp(self.ts), - position=Point2(x=x1, y=y2 + font_size), - text=f"confidence: {self.confidence:.3f}", - font_size=font_size, - text_color=Color(r=1.0, g=1.0, b=1.0, a=1), - background_color=Color(r=0, g=0, b=0, a=1), - ), + # Build label text - exclude class_id if it's -1 (VLM detection) + if self.class_id == -1: + label_text = f"{self.name}_{self.track_id}" + else: + label_text = f"{self.name}_{self.class_id}_{self.track_id}" + + annotations = [ TextAnnotation( timestamp=to_ros_stamp(self.ts), position=Point2(x=x1, y=y1), - text=f"{self.name}_{self.class_id}_{self.track_id}", + text=label_text, font_size=font_size, text_color=Color(r=1.0, g=1.0, b=1.0, a=1), background_color=Color(r=0, g=0, b=0, a=1), ), ] + # Only show confidence if it's not 1.0 + if self.confidence != 1.0: + annotations.append( + TextAnnotation( + timestamp=to_ros_stamp(self.ts), + position=Point2(x=x1, y=y2 + font_size), + text=f"confidence: {self.confidence:.3f}", + font_size=font_size, + text_color=Color(r=1.0, g=1.0, b=1.0, a=1), + background_color=Color(r=0, g=0, b=0, a=1), + ) + ) + + return annotations + def to_points_annotation(self) -> List[PointsAnnotation]: x1, y1, x2, y2 = self.bbox thickness = 1 + # Use consistent color based on object name, brighter for outline + outline_color = Color.from_string(self.name, alpha=1.0, brightness=1.25) + return [ PointsAnnotation( timestamp=to_ros_stamp(self.ts), - outline_color=Color(r=0.0, g=0.0, b=0.0, a=1.0), - fill_color=Color.from_string(self.name, alpha=0.15), + outline_color=outline_color, + fill_color=Color.from_string(self.name, alpha=0.2), thickness=thickness, points_length=4, points=[ @@ -301,8 +380,6 @@ def from_ros_detection2d(cls, ros_det: ROSDetection2D, **kwargs) -> "Detection2D # Extract timestamp ts = to_timestamp(ros_det.header.stamp) - # Name is not stored in ROS Detection2D, so we'll use a placeholder - # Remove 'name' from kwargs if present to avoid duplicate name = kwargs.pop("name", f"class_{class_id}") return cls( @@ -329,30 +406,3 @@ def to_ros_detection2d(self) -> ROSDetection2D: ], id=str(self.track_id), ) - - -class ImageDetections2D(ImageDetections[Detection2D]): - @classmethod - def from_bbox_detector( - cls, image: Image, raw_detections: InconvinientDetectionFormat, **kwargs - ) -> "ImageDetections2D": - return cls( - image=image, - detections=Detection2DBBox.from_detector(raw_detections, image=image, ts=image.ts), - ) - - @classmethod - def from_pose_detector( - cls, image: Image, people: List["Person"], **kwargs - ) -> "ImageDetections2D": - """Create ImageDetections2D from a list of Person detections. - Args: - image: Source image - people: List of Person objects with pose keypoints - Returns: - ImageDetections2D containing the pose detections - """ - return cls( - image=image, - detections=people, # Person objects are already Detection2D subclasses - ) diff --git a/dimos/perception/detection/type/detection2d/imageDetections2D.py b/dimos/perception/detection/type/detection2d/imageDetections2D.py new file mode 100644 index 0000000000..74854dae47 --- /dev/null +++ b/dimos/perception/detection/type/detection2d/imageDetections2D.py @@ -0,0 +1,79 @@ +# Copyright 2025 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 typing import List + +from dimos_lcm.vision_msgs import Detection2DArray +from ultralytics.engine.results import Results + +from dimos.msgs.sensor_msgs import Image +from dimos.perception.detection.type.detection2d.base import Detection2D +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +from dimos.perception.detection.type.imageDetections import ImageDetections + + +class ImageDetections2D(ImageDetections[Detection2D]): + @classmethod + def from_ros_detection2d_array( + cls, image: Image, ros_detections: Detection2DArray, **kwargs + ) -> "ImageDetections2D": + """Convert from ROS Detection2DArray message to ImageDetections2D object.""" + detections: List[Detection2D] = [] + for ros_det in ros_detections.detections: + detection = Detection2DBBox.from_ros_detection2d(ros_det, image=image, **kwargs) + if detection.is_valid(): # type: ignore[attr-defined] + detections.append(detection) + + return cls(image=image, detections=detections) + + @classmethod + def from_ultralytics_result( + cls, image: Image, results: List[Results], **kwargs + ) -> "ImageDetections2D": + """Create ImageDetections2D from ultralytics Results. + + Dispatches to appropriate Detection2D subclass based on result type: + - If keypoints present: creates Detection2DPerson + - Otherwise: creates Detection2DBBox + + Args: + image: Source image + results: List of ultralytics Results objects + **kwargs: Additional arguments passed to detection constructors + + Returns: + ImageDetections2D containing appropriate detection types + """ + from dimos.perception.detection.type.detection2d.person import Detection2DPerson + + detections: List[Detection2D] = [] + for result in results: + if result.boxes is None: + continue + + num_detections = len(result.boxes.xyxy) + for i in range(num_detections): + detection: Detection2D + if result.keypoints is not None: + # Pose detection with keypoints + detection = Detection2DPerson.from_ultralytics_result(result, i, image) + else: + # Regular bbox detection + detection = Detection2DBBox.from_ultralytics_result(result, i, image) + if detection.is_valid(): + detections.append(detection) + + return cls(image=image, detections=detections) diff --git a/dimos/perception/detection2d/type/person.py b/dimos/perception/detection/type/detection2d/person.py similarity index 65% rename from dimos/perception/detection2d/type/person.py rename to dimos/perception/detection/type/detection2d/person.py index b61045f48c..1c6fee5cae 100644 --- a/dimos/perception/detection2d/type/person.py +++ b/dimos/perception/detection/type/detection2d/person.py @@ -23,15 +23,16 @@ from dimos.msgs.foxglove_msgs.Color import Color from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.type.detection2d import Bbox, Detection2DBBox +from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox from dimos.types.timestamped import to_ros_stamp +from dimos.utils.decorators.decorators import simple_mcache if TYPE_CHECKING: from ultralytics.engine.results import Results @dataclass -class Person(Detection2DBBox): +class Detection2DPerson(Detection2DBBox): """Represents a detected person with pose keypoints.""" # Pose keypoints - additional fields beyond Detection2DBBox @@ -68,16 +69,48 @@ class Person(Detection2DBBox): ] @classmethod - def from_yolo(cls, result: "Results", person_idx: int, image: Image) -> "Person": - """Create Person instance from YOLO results. + def from_ultralytics_result( + cls, result: "Results", idx: int, image: Image + ) -> "Detection2DPerson": + """Create Detection2DPerson from ultralytics Results object with pose keypoints. Args: - result: Single Results object from YOLO - person_idx: Index of the person in the detection results - image: Original image for the detection + result: Ultralytics Results object containing detection and keypoint data + idx: Index of the detection in the results + image: Source image + + Returns: + Detection2DPerson instance + + Raises: + ValueError: If the result doesn't contain keypoints or is not a person detection """ + # Validate that this is a pose detection result + if not hasattr(result, "keypoints") or result.keypoints is None: + raise ValueError( + f"Cannot create Detection2DPerson from result without keypoints. " + f"This appears to be a regular detection result, not a pose detection. " + f"Use Detection2DBBox.from_ultralytics_result() instead." + ) + + if not hasattr(result, "boxes") or result.boxes is None: + raise ValueError("Cannot create Detection2DPerson from result without bounding boxes") + + # Check if this is actually a person detection (class 0 in COCO) + class_id = int(result.boxes.cls[idx].cpu()) + if class_id != 0: # Person is class 0 in COCO + class_name = ( + result.names.get(class_id, f"class_{class_id}") + if hasattr(result, "names") + else f"class_{class_id}" + ) + raise ValueError( + f"Cannot create Detection2DPerson from non-person detection. " + f"Got class {class_id} ({class_name}), expected class 0 (person)." + ) + # Extract bounding box as tuple for Detection2DBBox - bbox_array = result.boxes.xyxy[person_idx].cpu().numpy() + bbox_array = result.boxes.xyxy[idx].cpu().numpy() bbox: Bbox = ( float(bbox_array[0]), @@ -87,31 +120,42 @@ def from_yolo(cls, result: "Results", person_idx: int, image: Image) -> "Person" ) bbox_norm = ( - result.boxes.xyxyn[person_idx].cpu().numpy() if hasattr(result.boxes, "xyxyn") else None + result.boxes.xyxyn[idx].cpu().numpy() if hasattr(result.boxes, "xyxyn") else None ) - confidence = float(result.boxes.conf[person_idx].cpu()) - class_id = int(result.boxes.cls[person_idx].cpu()) + confidence = float(result.boxes.conf[idx].cpu()) + class_id = int(result.boxes.cls[idx].cpu()) # Extract keypoints - keypoints = result.keypoints.xy[person_idx].cpu().numpy() - keypoint_scores = result.keypoints.conf[person_idx].cpu().numpy() + if result.keypoints.xy is None or result.keypoints.conf is None: + raise ValueError("Keypoints xy or conf data is missing from the result") + + keypoints = result.keypoints.xy[idx].cpu().numpy() + keypoint_scores = result.keypoints.conf[idx].cpu().numpy() keypoints_norm = ( - result.keypoints.xyn[person_idx].cpu().numpy() - if hasattr(result.keypoints, "xyn") + result.keypoints.xyn[idx].cpu().numpy() + if hasattr(result.keypoints, "xyn") and result.keypoints.xyn is not None else None ) # Get image dimensions height, width = result.orig_shape + # Extract track ID if available + track_id = idx # Use index as default + if hasattr(result.boxes, "id") and result.boxes.id is not None: + track_id = int(result.boxes.id[idx].cpu()) + + # Get class name + name = result.names.get(class_id, "person") if hasattr(result, "names") else "person" + return cls( # Detection2DBBox fields bbox=bbox, - track_id=person_idx, # Use person index as track_id for now + track_id=track_id, class_id=class_id, confidence=confidence, - name="person", + name=name, ts=image.ts, image=image, # Person specific fields @@ -123,6 +167,30 @@ def from_yolo(cls, result: "Results", person_idx: int, image: Image) -> "Person" image_height=height, ) + @classmethod + def from_yolo(cls, result: "Results", idx: int, image: Image) -> "Detection2DPerson": + """Alias for from_ultralytics_result for backward compatibility.""" + return cls.from_ultralytics_result(result, idx, image) + + @classmethod + def from_ros_detection2d(cls, *args, **kwargs) -> "Detection2DPerson": + """Conversion from ROS Detection2D is not supported for Detection2DPerson. + + The ROS Detection2D message format does not include keypoint data, + which is required for Detection2DPerson. Use Detection2DBBox for + round-trip ROS conversions, or store keypoints separately. + + Raises: + NotImplementedError: Always raised as this conversion is impossible + """ + raise NotImplementedError( + "Cannot convert from ROS Detection2D to Detection2DPerson. " + "The ROS Detection2D message format does not contain keypoint data " + "(keypoints and keypoint_scores) which are required fields for Detection2DPerson. " + "Consider using Detection2DBBox for ROS conversions, or implement a custom " + "message format that includes pose keypoints." + ) + def get_keypoint(self, name: str) -> Tuple[np.ndarray, float]: """Get specific keypoint by name. Returns: @@ -145,6 +213,11 @@ def get_visible_keypoints(self, threshold: float = 0.5) -> List[Tuple[str, np.nd visible.append((name, self.keypoints[i], score)) return visible + @simple_mcache + def is_valid(self) -> bool: + valid_keypoints = sum(1 for score in self.keypoint_scores if score > 0.8) + return valid_keypoints >= 5 + @property def width(self) -> float: """Get width of bounding box.""" diff --git a/dimos/perception/detection2d/type/test_detection2d.py b/dimos/perception/detection/type/detection2d/test_bbox.py similarity index 100% rename from dimos/perception/detection2d/type/test_detection2d.py rename to dimos/perception/detection/type/detection2d/test_bbox.py diff --git a/dimos/perception/detection/type/detection2d/test_imageDetections2D.py b/dimos/perception/detection/type/detection2d/test_imageDetections2D.py new file mode 100644 index 0000000000..6731b7b0c7 --- /dev/null +++ b/dimos/perception/detection/type/detection2d/test_imageDetections2D.py @@ -0,0 +1,52 @@ +# Copyright 2025 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. +import pytest + +from dimos.perception.detection.type import ImageDetections2D + + +def test_from_ros_detection2d_array(get_moment_2d): + moment = get_moment_2d() + + detections2d = moment["detections2d"] + + test_image = detections2d.image + + # Convert to ROS detection array + ros_array = detections2d.to_ros_detection2d_array() + + # Convert back to ImageDetections2D + recovered = ImageDetections2D.from_ros_detection2d_array(test_image, ros_array) + + # Verify we got the same number of detections + assert len(recovered.detections) == len(detections2d.detections) + + # Verify the detection matches + original_det = detections2d.detections[0] + recovered_det = recovered.detections[0] + + # Check bbox is approximately the same (allow 1 pixel tolerance due to float conversion) + for orig_val, rec_val in zip(original_det.bbox, recovered_det.bbox): + assert orig_val == pytest.approx(rec_val, abs=1.0) + + # Check other properties + assert recovered_det.track_id == original_det.track_id + assert recovered_det.class_id == original_det.class_id + assert recovered_det.confidence == pytest.approx(original_det.confidence, abs=0.01) + + print(f"\nSuccessfully round-tripped detection through ROS format:") + print(f" Original bbox: {original_det.bbox}") + print(f" Recovered bbox: {recovered_det.bbox}") + print(f" Track ID: {recovered_det.track_id}") + print(f" Confidence: {recovered_det.confidence:.3f}") diff --git a/dimos/perception/detection/type/detection2d/test_person.py b/dimos/perception/detection/type/detection2d/test_person.py new file mode 100644 index 0000000000..ba930fd299 --- /dev/null +++ b/dimos/perception/detection/type/detection2d/test_person.py @@ -0,0 +1,71 @@ +# Copyright 2025 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. +import pytest + + +def test_person_ros_confidence(): + """Test that Detection2DPerson preserves confidence when converting to ROS format.""" + + from dimos.msgs.sensor_msgs import Image + from dimos.perception.detection.detectors.person.yolo import YoloPersonDetector + from dimos.perception.detection.type.detection2d.person import Detection2DPerson + from dimos.utils.data import get_data + + # Load test image + image_path = get_data("cafe.jpg") + image = Image.from_file(image_path) + + # Run pose detection + detector = YoloPersonDetector(device="cpu") + detections = detector.process_image(image) + + # Find a Detection2DPerson (should have at least one person in cafe.jpg) + person_detections = [d for d in detections.detections if isinstance(d, Detection2DPerson)] + assert len(person_detections) > 0, "No person detections found in cafe.jpg" + + # Test each person detection + for person_det in person_detections: + original_confidence = person_det.confidence + assert 0.0 <= original_confidence <= 1.0, "Confidence should be between 0 and 1" + + # Convert to ROS format + ros_det = person_det.to_ros_detection2d() + + # Extract confidence from ROS message + assert len(ros_det.results) > 0, "ROS detection should have results" + ros_confidence = ros_det.results[0].hypothesis.score + + # Verify confidence is preserved (allow small floating point tolerance) + assert original_confidence == pytest.approx(ros_confidence, abs=0.001), ( + f"Confidence mismatch: {original_confidence} != {ros_confidence}" + ) + + print("\nSuccessfully preserved confidence in ROS conversion for Detection2DPerson:") + print(f" Original confidence: {original_confidence:.3f}") + print(f" ROS confidence: {ros_confidence:.3f}") + print(f" Track ID: {person_det.track_id}") + print(f" Visible keypoints: {len(person_det.get_visible_keypoints(threshold=0.3))}/17") + + +def test_person_from_ros_raises(): + """Test that Detection2DPerson.from_ros_detection2d() raises NotImplementedError.""" + from dimos.perception.detection.type.detection2d.person import Detection2DPerson + + with pytest.raises(NotImplementedError) as exc_info: + Detection2DPerson.from_ros_detection2d() + + # Verify the error message is informative + error_msg = str(exc_info.value) + assert "keypoint data" in error_msg.lower() + assert "Detection2DBBox" in error_msg diff --git a/dimos/perception/detection/type/detection3d/__init__.py b/dimos/perception/detection/type/detection3d/__init__.py new file mode 100644 index 0000000000..a8d11ca87f --- /dev/null +++ b/dimos/perception/detection/type/detection3d/__init__.py @@ -0,0 +1,37 @@ +# Copyright 2025 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 dimos.perception.detection.type.detection3d.base import Detection3D +from dimos.perception.detection.type.detection3d.bbox import Detection3DBBox +from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC +from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC +from dimos.perception.detection.type.detection3d.pointcloud_filters import ( + PointCloudFilter, + height_filter, + radius_outlier, + raycast, + statistical, +) + +__all__ = [ + "Detection3D", + "Detection3DBBox", + "Detection3DPC", + "ImageDetections3DPC", + "PointCloudFilter", + "height_filter", + "raycast", + "radius_outlier", + "statistical", +] diff --git a/dimos/perception/detection/type/detection3d/base.py b/dimos/perception/detection/type/detection3d/base.py new file mode 100644 index 0000000000..a82a50d474 --- /dev/null +++ b/dimos/perception/detection/type/detection3d/base.py @@ -0,0 +1,44 @@ +# Copyright 2025 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 abc import abstractmethod +from dataclasses import dataclass +from typing import Optional + +from dimos_lcm.sensor_msgs import CameraInfo + +from dimos.msgs.geometry_msgs import Transform +from dimos.perception.detection.type.detection2d import Detection2DBBox + + +@dataclass +class Detection3D(Detection2DBBox): + """Abstract base class for 3D detections.""" + + transform: Transform + frame_id: str + + @classmethod + @abstractmethod + def from_2d( + cls, + det: Detection2DBBox, + distance: float, + camera_info: CameraInfo, + world_to_optical_transform: Transform, + ) -> Optional["Detection3D"]: + """Create a 3D detection from a 2D detection.""" + ... diff --git a/dimos/perception/detection/type/detection3d/bbox.py b/dimos/perception/detection/type/detection3d/bbox.py new file mode 100644 index 0000000000..2bc0c1c541 --- /dev/null +++ b/dimos/perception/detection/type/detection3d/bbox.py @@ -0,0 +1,76 @@ +# Copyright 2025 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 + +import functools +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, TypeVar + +import numpy as np +from dimos_lcm.sensor_msgs import CameraInfo +from lcm_msgs.builtin_interfaces import Duration +from lcm_msgs.foxglove_msgs import CubePrimitive, SceneEntity, SceneUpdate, TextPrimitive +from lcm_msgs.geometry_msgs import Point, Pose, Quaternion +from lcm_msgs.geometry_msgs import Vector3 as LCMVector3 + +from dimos.msgs.foxglove_msgs.Color import Color +from dimos.msgs.geometry_msgs import PoseStamped, Transform, Vector3 +from dimos.msgs.sensor_msgs import PointCloud2 +from dimos.perception.detection.type.detection2d import Detection2D, Detection2DBBox +from dimos.perception.detection.type.detection3d.base import Detection3D +from dimos.perception.detection.type.imageDetections import ImageDetections +from dimos.types.timestamped import to_ros_stamp + + +@dataclass +class Detection3DBBox(Detection2DBBox): + """3D bounding box detection with center, size, and orientation. + + Represents a 3D detection as an oriented bounding box in world space. + """ + + transform: Transform # Camera to world transform + frame_id: str # Frame ID (e.g., "world", "map") + center: Vector3 # Center point in world frame + size: Vector3 # Width, height, depth + orientation: tuple[float, float, float, float] # Quaternion (x, y, z, w) + + @functools.cached_property + def pose(self) -> PoseStamped: + """Convert detection to a PoseStamped using bounding box center. + + Returns pose in world frame with the detection's orientation. + """ + return PoseStamped( + ts=self.ts, + frame_id=self.frame_id, + position=self.center, + orientation=self.orientation, + ) + + def to_repr_dict(self) -> Dict[str, Any]: + # Calculate distance from camera + camera_pos = self.transform.translation + distance = (self.center - camera_pos).magnitude() + + parent_dict = super().to_repr_dict() + # Remove bbox key if present + parent_dict.pop("bbox", None) + + return { + **parent_dict, + "dist": f"{distance:.2f}m", + "size": f"[{self.size.x:.2f},{self.size.y:.2f},{self.size.z:.2f}]", + } diff --git a/dimos/perception/detection/type/detection3d/imageDetections3DPC.py b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py new file mode 100644 index 0000000000..efad114a2c --- /dev/null +++ b/dimos/perception/detection/type/detection3d/imageDetections3DPC.py @@ -0,0 +1,45 @@ +# Copyright 2025 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 lcm_msgs.foxglove_msgs import SceneUpdate + +from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC +from dimos.perception.detection.type.imageDetections import ImageDetections + + +class ImageDetections3DPC(ImageDetections[Detection3DPC]): + """Specialized class for 3D detections in an image.""" + + def to_foxglove_scene_update(self) -> "SceneUpdate": + """Convert all detections to a Foxglove SceneUpdate message. + + Returns: + SceneUpdate containing SceneEntity objects for all detections + """ + + # Create SceneUpdate message with all detections + scene_update = SceneUpdate() + scene_update.deletions_length = 0 + scene_update.deletions = [] + scene_update.entities = [] + + # Process each detection + for i, detection in enumerate(self.detections): + entity = detection.to_foxglove_scene_entity(entity_id=f"detection_{detection.name}_{i}") + scene_update.entities.append(entity) + + scene_update.entities_length = len(scene_update.entities) + return scene_update diff --git a/dimos/perception/detection2d/type/detection3d.py b/dimos/perception/detection/type/detection3d/pointcloud.py similarity index 56% rename from dimos/perception/detection2d/type/detection3d.py rename to dimos/perception/detection/type/detection3d/pointcloud.py index a203bb1a4b..e5fb82549c 100644 --- a/dimos/perception/detection2d/type/detection3d.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -16,7 +16,7 @@ import functools from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, TypeVar +from typing import Any, Dict, Optional import numpy as np from dimos_lcm.sensor_msgs import CameraInfo @@ -28,25 +28,20 @@ from dimos.msgs.foxglove_msgs.Color import Color from dimos.msgs.geometry_msgs import PoseStamped, Transform, Vector3 from dimos.msgs.sensor_msgs import PointCloud2 -from dimos.perception.detection2d.type.detection2d import Detection2D, Detection2DBBox -from dimos.perception.detection2d.type.imageDetections import ImageDetections +from dimos.perception.detection.type.detection2d import Detection2DBBox +from dimos.perception.detection.type.detection3d.base import Detection3D +from dimos.perception.detection.type.detection3d.pointcloud_filters import ( + PointCloudFilter, + radius_outlier, + raycast, + statistical, +) from dimos.types.timestamped import to_ros_stamp @dataclass -class Detection3D(Detection2DBBox): - transform: Transform - frame_id: str - - @classmethod - def from_2d( - cls, - det: Detection2D, - distance: float, - camera_info: CameraInfo, - world_to_optical_transform: Transform, - ) -> Optional["Detection3D"]: - raise NotImplementedError() +class Detection3DPC(Detection3D): + pointcloud: PointCloud2 @functools.cached_property def center(self) -> Vector3: @@ -78,7 +73,7 @@ def get_bounding_box_dimensions(self) -> tuple[float, float, float]: """Get dimensions (width, height, depth) of the detection's bounding box.""" return self.pointcloud.get_bounding_box_dimensions() - def bounding_box_intersects(self, other: "Detection3D") -> bool: + def bounding_box_intersects(self, other: "Detection3DPC") -> bool: """Check if this detection's bounding box intersects with another's.""" return self.pointcloud.bounding_box_intersects(other.pointcloud) @@ -101,7 +96,7 @@ def to_repr_dict(self) -> Dict[str, Any]: "points": str(len(self.pointcloud)), } - def to_foxglove_scene_entity(self, entity_id: str = None) -> "SceneEntity": + def to_foxglove_scene_entity(self, entity_id: Optional[str] = None) -> "SceneEntity": """Convert detection to a Foxglove SceneEntity with cube primitive and text label. Args: @@ -200,30 +195,131 @@ def to_foxglove_scene_entity(self, entity_id: str = None) -> "SceneEntity": def scene_entity_label(self) -> str: return f"{self.track_id}/{self.name} ({self.confidence:.0%})" + @classmethod + def from_2d( # type: ignore[override] + cls, + det: Detection2DBBox, + world_pointcloud: PointCloud2, + camera_info: CameraInfo, + world_to_optical_transform: Transform, + # filters are to be adjusted based on the sensor noise characteristics if feeding + # sensor data directly + filters: Optional[list[PointCloudFilter]] = None, + ) -> Optional["Detection3DPC"]: + """Create a Detection3D from a 2D detection by projecting world pointcloud. + + This method handles: + 1. Projecting world pointcloud to camera frame + 2. Filtering points within the 2D detection bounding box + 3. Cleaning up the pointcloud (height filter, outlier removal) + 4. Hidden point removal from camera perspective -T = TypeVar("T", bound="Detection2D") - - -class ImageDetections3D(ImageDetections[Detection3D]): - """Specialized class for 3D detections in an image.""" - - def to_foxglove_scene_update(self) -> "SceneUpdate": - """Convert all detections to a Foxglove SceneUpdate message. - + Args: + det: The 2D detection + world_pointcloud: Full pointcloud in world frame + camera_info: Camera calibration info + world_to_camerlka_transform: Transform from world to camera frame + filters: List of functions to apply to the pointcloud for filtering Returns: - SceneUpdate containing SceneEntity objects for all detections + Detection3D with filtered pointcloud, or None if no valid points """ + # Set default filters if none provided + if filters is None: + filters = [ + # height_filter(0.1), + raycast(), + radius_outlier(), + statistical(), + ] + + # Extract camera parameters + fx, fy = camera_info.K[0], camera_info.K[4] + cx, cy = camera_info.K[2], camera_info.K[5] + image_width = camera_info.width + image_height = camera_info.height + + camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) + + # Convert pointcloud to numpy array + world_points = world_pointcloud.as_numpy() + + # Project points to camera frame + points_homogeneous = np.hstack([world_points, np.ones((world_points.shape[0], 1))]) + extrinsics_matrix = world_to_optical_transform.to_matrix() + points_camera = (extrinsics_matrix @ points_homogeneous.T).T + + # Filter out points behind the camera + valid_mask = points_camera[:, 2] > 0 + points_camera = points_camera[valid_mask] + world_points = world_points[valid_mask] + + if len(world_points) == 0: + return None + + # Project to 2D + points_2d_homogeneous = (camera_matrix @ points_camera[:, :3].T).T + points_2d = points_2d_homogeneous[:, :2] / points_2d_homogeneous[:, 2:3] + + # Filter points within image bounds + in_image_mask = ( + (points_2d[:, 0] >= 0) + & (points_2d[:, 0] < image_width) + & (points_2d[:, 1] >= 0) + & (points_2d[:, 1] < image_height) + ) + points_2d = points_2d[in_image_mask] + world_points = world_points[in_image_mask] + + if len(world_points) == 0: + return None + + # Extract bbox from Detection2D + x_min, y_min, x_max, y_max = det.bbox + + # Find points within this detection box (with small margin) + margin = 5 # pixels + in_box_mask = ( + (points_2d[:, 0] >= x_min - margin) + & (points_2d[:, 0] <= x_max + margin) + & (points_2d[:, 1] >= y_min - margin) + & (points_2d[:, 1] <= y_max + margin) + ) + + detection_points = world_points[in_box_mask] - # Create SceneUpdate message with all detections - scene_update = SceneUpdate() - scene_update.deletions_length = 0 - scene_update.deletions = [] - scene_update.entities = [] + if detection_points.shape[0] == 0: + # print(f"No points found in detection bbox after projection. {det.name}") + return None - # Process each detection - for i, detection in enumerate(self.detections): - entity = detection.to_foxglove_scene_entity(entity_id=f"detection_{detection.name}_{i}") - scene_update.entities.append(entity) + # Create initial pointcloud for this detection + initial_pc = PointCloud2.from_numpy( + detection_points, + frame_id=world_pointcloud.frame_id, + timestamp=world_pointcloud.ts, + ) - scene_update.entities_length = len(scene_update.entities) - return scene_update + # Apply filters - each filter gets all arguments + detection_pc = initial_pc + for filter_func in filters: + result = filter_func(det, detection_pc, camera_info, world_to_optical_transform) + if result is None: + return None + detection_pc = result + + # Final check for empty pointcloud + if len(detection_pc.pointcloud.points) == 0: + return None + + # Create Detection3D with filtered pointcloud + return cls( + image=det.image, + bbox=det.bbox, + track_id=det.track_id, + class_id=det.class_id, + confidence=det.confidence, + name=det.name, + ts=det.ts, + pointcloud=detection_pc, + transform=world_to_optical_transform, + frame_id=world_pointcloud.frame_id, + ) diff --git a/dimos/perception/detection/type/detection3d/pointcloud_filters.py b/dimos/perception/detection/type/detection3d/pointcloud_filters.py new file mode 100644 index 0000000000..51cf3d7f33 --- /dev/null +++ b/dimos/perception/detection/type/detection3d/pointcloud_filters.py @@ -0,0 +1,82 @@ +# Copyright 2025 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 typing import Callable, Optional + +from dimos_lcm.sensor_msgs import CameraInfo + +from dimos.msgs.geometry_msgs import Transform +from dimos.msgs.sensor_msgs import PointCloud2 +from dimos.perception.detection.type.detection2d import Detection2DBBox + +# Filters take Detection2DBBox, PointCloud2, CameraInfo, Transform and return filtered PointCloud2 or None +PointCloudFilter = Callable[ + [Detection2DBBox, PointCloud2, CameraInfo, Transform], Optional[PointCloud2] +] + + +def height_filter(height=0.1) -> PointCloudFilter: + return lambda det, pc, ci, tf: pc.filter_by_height(height) + + +def statistical(nb_neighbors=40, std_ratio=0.5) -> PointCloudFilter: + def filter_func( + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + ) -> Optional[PointCloud2]: + try: + statistical, removed = pc.pointcloud.remove_statistical_outlier( + nb_neighbors=nb_neighbors, std_ratio=std_ratio + ) + return PointCloud2(statistical, pc.frame_id, pc.ts) + except Exception as e: + # print("statistical filter failed:", e) + return None + + return filter_func + + +def raycast() -> PointCloudFilter: + def filter_func( + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + ) -> Optional[PointCloud2]: + try: + camera_pos = tf.inverse().translation + camera_pos_np = camera_pos.to_numpy() + _, visible_indices = pc.pointcloud.hidden_point_removal(camera_pos_np, radius=100.0) + visible_pcd = pc.pointcloud.select_by_index(visible_indices) + return PointCloud2(visible_pcd, pc.frame_id, pc.ts) + except Exception as e: + # print("raycast filter failed:", e) + return None + + return filter_func + + +def radius_outlier(min_neighbors: int = 20, radius: float = 0.3) -> PointCloudFilter: + """ + Remove isolated points: keep only points that have at least `min_neighbors` + neighbors within `radius` meters (same units as your point cloud). + """ + + def filter_func( + det: Detection2DBBox, pc: PointCloud2, ci: CameraInfo, tf: Transform + ) -> Optional[PointCloud2]: + filtered_pcd, removed = pc.pointcloud.remove_radius_outlier( + nb_points=min_neighbors, radius=radius + ) + return PointCloud2(filtered_pcd, pc.frame_id, pc.ts) + + return filter_func diff --git a/dimos/perception/detection/type/detection3d/test_imageDetections3DPC.py b/dimos/perception/detection/type/detection3d/test_imageDetections3DPC.py new file mode 100644 index 0000000000..31e44dad91 --- /dev/null +++ b/dimos/perception/detection/type/detection3d/test_imageDetections3DPC.py @@ -0,0 +1,35 @@ +# Copyright 2025 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. + +import pytest + + +@pytest.mark.skip +def test_to_foxglove_scene_update(detections3dpc): + # Convert to scene update + scene_update = detections3dpc.to_foxglove_scene_update() + + # Verify scene update structure + assert scene_update is not None + assert scene_update.deletions_length == 0 + assert len(scene_update.deletions) == 0 + assert scene_update.entities_length == len(detections3dpc.detections) + assert len(scene_update.entities) == len(detections3dpc.detections) + + # Verify each entity corresponds to a detection + for i, (entity, detection) in enumerate(zip(scene_update.entities, detections3dpc.detections)): + assert entity.id == str(detection.track_id) + assert entity.frame_id == detection.frame_id + assert entity.cubes_length == 1 + assert entity.texts_length == 1 diff --git a/dimos/perception/detection2d/type/test_detection3dpc.py b/dimos/perception/detection/type/detection3d/test_pointcloud.py similarity index 99% rename from dimos/perception/detection2d/type/test_detection3dpc.py rename to dimos/perception/detection/type/detection3d/test_pointcloud.py index b9ad04fb3a..308839f8bf 100644 --- a/dimos/perception/detection2d/type/test_detection3dpc.py +++ b/dimos/perception/detection/type/detection3d/test_pointcloud.py @@ -58,7 +58,7 @@ def test_detection3dpc(detection3dpc): # def test_point_cloud_properties(detection3dpc): """Test point cloud data and boundaries.""" pc_points = detection3dpc.pointcloud.points() - assert len(pc_points) in [68, 69, 70] + assert len(pc_points) > 60 assert detection3dpc.pointcloud.frame_id == "world", ( f"Expected frame_id 'world', got '{detection3dpc.pointcloud.frame_id}'" ) diff --git a/dimos/perception/detection/type/imageDetections.py b/dimos/perception/detection/type/imageDetections.py new file mode 100644 index 0000000000..994c939e4d --- /dev/null +++ b/dimos/perception/detection/type/imageDetections.py @@ -0,0 +1,79 @@ +# Copyright 2025 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 typing import TYPE_CHECKING, Generic, List, Optional, TypeVar + +from dimos_lcm.vision_msgs import Detection2DArray + +from dimos.msgs.foxglove_msgs import ImageAnnotations +from dimos.msgs.sensor_msgs import Image +from dimos.msgs.std_msgs import Header +from dimos.perception.detection.type.utils import TableStr + +if TYPE_CHECKING: + from dimos.perception.detection.type.detection2d.base import Detection2D + + T = TypeVar("T", bound=Detection2D) +else: + from dimos.perception.detection.type.detection2d.base import Detection2D + + T = TypeVar("T", bound=Detection2D) + + +class ImageDetections(Generic[T], TableStr): + image: Image + detections: List[T] + + @property + def ts(self) -> float: + return self.image.ts + + def __init__(self, image: Image, detections: Optional[List[T]] = None): + self.image = image + self.detections = detections or [] + for det in self.detections: + if not det.ts: + det.ts = image.ts + + def __len__(self): + return len(self.detections) + + def __iter__(self): + return iter(self.detections) + + def __getitem__(self, index): + return self.detections[index] + + def to_ros_detection2d_array(self) -> Detection2DArray: + return Detection2DArray( + detections_length=len(self.detections), + header=Header(self.image.ts, "camera_optical"), + detections=[det.to_ros_detection2d() for det in self.detections], + ) + + def to_foxglove_annotations(self) -> ImageAnnotations: + def flatten(xss): + return [x for xs in xss for x in xs] + + texts = flatten(det.to_text_annotation() for det in self.detections) + points = flatten(det.to_points_annotation() for det in self.detections) + + return ImageAnnotations( + texts=texts, + texts_length=len(texts), + points=points, + points_length=len(points), + ) diff --git a/dimos/perception/detection2d/type/test_detection3d.py b/dimos/perception/detection/type/test_detection3d.py similarity index 91% rename from dimos/perception/detection2d/type/test_detection3d.py rename to dimos/perception/detection/type/test_detection3d.py index 642e6c7542..44413df1fe 100644 --- a/dimos/perception/detection2d/type/test_detection3d.py +++ b/dimos/perception/detection/type/test_detection3d.py @@ -14,11 +14,11 @@ import time -from dimos.perception.detection2d.type.detection3d import Detection3D +from dimos.perception.detection.type.detection3d import Detection3D def test_guess_projection(get_moment_2d, publish_moment): - moment = get_moment_2d(seek=10.0) + moment = get_moment_2d() for key, value in moment.items(): print(key, "====================================") print(value) diff --git a/dimos/perception/detection2d/type/test_object3d.py b/dimos/perception/detection/type/test_object3d.py similarity index 90% rename from dimos/perception/detection2d/type/test_object3d.py rename to dimos/perception/detection/type/test_object3d.py index b7933e86d5..1dc3cb6bd0 100644 --- a/dimos/perception/detection2d/type/test_object3d.py +++ b/dimos/perception/detection/type/test_object3d.py @@ -14,10 +14,10 @@ import pytest -from dimos.perception.detection2d.module2D import Detection2DModule -from dimos.perception.detection2d.module3D import Detection3DModule -from dimos.perception.detection2d.moduleDB import Object3D, ObjectDBModule -from dimos.perception.detection2d.type.detection3d import ImageDetections3D +from dimos.perception.detection.module2D import Detection2DModule +from dimos.perception.detection.module3D import Detection3DModule +from dimos.perception.detection.moduleDB import Object3D, ObjectDBModule +from dimos.perception.detection.type.detection3d import ImageDetections3DPC from dimos.robot.unitree_webrtc.modular.connection_module import ConnectionModule @@ -86,9 +86,9 @@ def test_object3d_repr_dict(first_object): assert encoded["last_seen"].endswith("s ago") # def test_object3d_image_property(first_object): - """Test image property returns best_detection's image.""" - assert first_object.image is not None - assert first_object.image is first_object.best_detection.image + """Test get_image method returns best_detection's image.""" + assert first_object.get_image() is not None + assert first_object.get_image() is first_object.best_detection.image def test_all_objeects(all_objects): @@ -158,7 +158,7 @@ def test_objectdb_module(object_db_module): assert combined.center is not None # def test_image_detections3d_scene_update(object_db_module): - """Test ImageDetections3D to Foxglove scene update conversion.""" + """Test ImageDetections3DPC to Foxglove scene update conversion.""" # Get some detections objects = list(object_db_module.objects.values()) if not objects: @@ -166,7 +166,7 @@ def test_objectdb_module(object_db_module): detections = [obj.best_detection for obj in objects[:3]] # Take up to 3 - image_detections = ImageDetections3D(image=detections[0].image, detections=detections) + image_detections = ImageDetections3DPC(image=detections[0].image, detections=detections) scene_update = image_detections.to_foxglove_scene_update() diff --git a/dimos/perception/detection2d/type/imageDetections.py b/dimos/perception/detection/type/utils.py similarity index 64% rename from dimos/perception/detection2d/type/imageDetections.py rename to dimos/perception/detection/type/utils.py index edd8449f06..f1e2187015 100644 --- a/dimos/perception/detection2d/type/imageDetections.py +++ b/dimos/perception/detection/type/utils.py @@ -12,26 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - import hashlib -from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Tuple, TypeVar from rich.console import Console from rich.table import Table from rich.text import Text -from dimos.msgs.foxglove_msgs import ImageAnnotations -from dimos.msgs.sensor_msgs import Image -from dimos.msgs.std_msgs import Header -from dimos.msgs.vision_msgs import Detection2DArray from dimos.types.timestamped import to_timestamp -if TYPE_CHECKING: - from dimos.perception.detection2d.type.detection2d import Detection2D - -T = TypeVar("T", bound="Detection2D") - def _hash_to_color(name: str) -> str: """Generate a consistent color for a given name using hash.""" @@ -60,6 +48,8 @@ def _hash_to_color(name: str) -> str: class TableStr: + """Mixin class that provides table-based string representation for detection collections.""" + def __str__(self): console = Console(force_terminal=True, legacy_windows=False) @@ -109,49 +99,3 @@ def __str__(self): with console.capture() as capture: console.print(table) return capture.get().strip() - - -class ImageDetections(Generic[T], TableStr): - image: Image - detections: List[T] - - @property - def ts(self) -> float: - return self.image.ts - - def __init__(self, image: Image, detections: Optional[List[T]] = None): - self.image = image - self.detections = detections or [] - for det in self.detections: - if not det.ts: - det.ts = image.ts - - def __len__(self): - return len(self.detections) - - def __iter__(self): - return iter(self.detections) - - def __getitem__(self, index): - return self.detections[index] - - def to_ros_detection2d_array(self) -> Detection2DArray: - return Detection2DArray( - detections_length=len(self.detections), - header=Header(self.image.ts, "camera_optical"), - detections=[det.to_ros_detection2d() for det in self.detections], - ) - - def to_foxglove_annotations(self) -> ImageAnnotations: - def flatten(xss): - return [x for xs in xss for x in xs] - - texts = flatten(det.to_text_annotation() for det in self.detections) - points = flatten(det.to_points_annotation() for det in self.detections) - - return ImageAnnotations( - texts=texts, - texts_length=len(texts), - points=points, - points_length=len(points), - ) diff --git a/dimos/perception/detection2d/__init__.py b/dimos/perception/detection2d/__init__.py deleted file mode 100644 index 6dc59e7366..0000000000 --- a/dimos/perception/detection2d/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from dimos.perception.detection2d.detectors import * -from dimos.perception.detection2d.module2D import ( - Detection2DModule, -) -from dimos.perception.detection2d.module3D import ( - Detection3DModule, -) -from dimos.perception.detection2d.utils import * diff --git a/dimos/perception/detection2d/detectors/__init__.py b/dimos/perception/detection2d/detectors/__init__.py deleted file mode 100644 index 287fff1a15..0000000000 --- a/dimos/perception/detection2d/detectors/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# from dimos.perception.detection2d.detectors.detic import Detic2DDetector -from dimos.perception.detection2d.detectors.types import Detector -from dimos.perception.detection2d.detectors.yolo import Yolo2DDetector diff --git a/dimos/perception/detection2d/detectors/person/test_annotations.py b/dimos/perception/detection2d/detectors/person/test_annotations.py deleted file mode 100644 index c686c33bd9..0000000000 --- a/dimos/perception/detection2d/detectors/person/test_annotations.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright 2025 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. - -"""Test person annotations work correctly.""" - -import sys - -from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.person.yolo import YoloPersonDetector -from dimos.utils.data import get_data - - -def test_person_annotations(): - """Test that Person annotations include keypoints and skeleton.""" - image = Image.from_file(get_data("cafe.jpg")) - detector = YoloPersonDetector() - people = detector.detect_people(image) - - assert len(people) > 0 - person = people[0] - - # Test text annotations - text_anns = person.to_text_annotation() - print(f"\nText annotations: {len(text_anns)}") - for i, ann in enumerate(text_anns): - print(f" {i}: {ann.text}") - assert len(text_anns) == 3 # confidence, name/track_id, keypoints count - assert any("keypoints:" in ann.text for ann in text_anns) - - # Test points annotations - points_anns = person.to_points_annotation() - print(f"\nPoints annotations: {len(points_anns)}") - - # Count different types (use actual LCM constants) - from dimos_lcm.foxglove_msgs.ImageAnnotations import PointsAnnotation - - bbox_count = sum(1 for ann in points_anns if ann.type == PointsAnnotation.LINE_LOOP) # 2 - keypoint_count = sum(1 for ann in points_anns if ann.type == PointsAnnotation.POINTS) # 1 - skeleton_count = sum(1 for ann in points_anns if ann.type == PointsAnnotation.LINE_LIST) # 4 - - print(f" - Bounding boxes: {bbox_count}") - print(f" - Keypoint circles: {keypoint_count}") - print(f" - Skeleton lines: {skeleton_count}") - - assert bbox_count >= 1 # At least the person bbox - assert keypoint_count >= 1 # At least some visible keypoints - assert skeleton_count >= 1 # At least some skeleton connections - - # Test full image annotations - img_anns = person.to_image_annotations() - assert img_anns.texts_length == len(text_anns) - assert img_anns.points_length == len(points_anns) - - print(f"\nāœ“ Person annotations working correctly!") - print(f" - {len(person.get_visible_keypoints(0.5))}/17 visible keypoints") - - -if __name__ == "__main__": - test_person_annotations() diff --git a/dimos/perception/detection2d/detectors/person/test_detection2d_conformance.py b/dimos/perception/detection2d/detectors/person/test_detection2d_conformance.py deleted file mode 100644 index f7c7cc088c..0000000000 --- a/dimos/perception/detection2d/detectors/person/test_detection2d_conformance.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright 2025 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. - -import pytest -from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.person.yolo import YoloPersonDetector -from dimos.perception.detection2d.type.person import Person -from dimos.utils.data import get_data - - -def test_person_detection2d_bbox_conformance(): - """Test that Person conforms to Detection2DBBox interface.""" - image = Image.from_file(get_data("cafe.jpg")) - detector = YoloPersonDetector() - people = detector.detect_people(image) - - assert len(people) > 0 - person = people[0] - - # Test Detection2DBBox methods - # Test bbox operations - assert hasattr(person, "bbox") - assert len(person.bbox) == 4 - assert all(isinstance(x, float) for x in person.bbox) - - # Test inherited properties - assert hasattr(person, "get_bbox_center") - center_bbox = person.get_bbox_center() - assert len(center_bbox) == 4 # center_x, center_y, width, height - - # Test volume calculation - volume = person.bbox_2d_volume() - assert volume > 0 - - # Test cropped image - cropped = person.cropped_image(padding=10) - assert isinstance(cropped, Image) - - # Test annotation methods - text_annotations = person.to_text_annotation() - assert len(text_annotations) == 3 # confidence, name/track_id, and keypoints count - - points_annotations = person.to_points_annotation() - # Should have: 1 bbox + 1 keypoints + multiple skeleton lines - assert len(points_annotations) > 1 - print(f" - Points annotations: {len(points_annotations)} (bbox + keypoints + skeleton)") - - # Test image annotations - annotations = person.to_image_annotations() - assert annotations.texts_length == 3 - assert annotations.points_length > 1 - - # Test ROS conversion - ros_det = person.to_ros_detection2d() - assert ros_det.bbox.size_x == person.width - assert ros_det.bbox.size_y == person.height - - # Test string representation - str_repr = str(person) - assert "Person" in str_repr - assert "person" in str_repr # name field - - print("\nāœ“ Person class fully conforms to Detection2DBBox interface") - print(f" - Detected {len(people)} people") - print(f" - First person confidence: {person.confidence:.3f}") - print(f" - Bbox volume: {volume:.1f}") - print(f" - Has {len(person.get_visible_keypoints(0.5))} visible keypoints") - - -if __name__ == "__main__": - test_person_detection2d_bbox_conformance() diff --git a/dimos/perception/detection2d/detectors/person/test_imagedetections2d.py b/dimos/perception/detection2d/detectors/person/test_imagedetections2d.py deleted file mode 100644 index 89fd770aa6..0000000000 --- a/dimos/perception/detection2d/detectors/person/test_imagedetections2d.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright 2025 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. - -"""Test ImageDetections2D with pose detections.""" - -from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.person.yolo import YoloPersonDetector -from dimos.perception.detection2d.type import ImageDetections2D -from dimos.utils.data import get_data - - -def test_image_detections_2d_with_person(): - """Test creating ImageDetections2D from person detector.""" - # Load image and detect people - image = Image.from_file(get_data("cafe.jpg")) - detector = YoloPersonDetector() - people = detector.detect_people(image) - - # Create ImageDetections2D using from_pose_detector - image_detections = ImageDetections2D.from_pose_detector(image, people) - - # Verify structure - assert image_detections.image is image - assert len(image_detections.detections) == len(people) - assert all(det in people for det in image_detections.detections) - - # Test image annotations (includes pose keypoints) - annotations = image_detections.to_foxglove_annotations() - print(f"\nImageDetections2D created with {len(people)} people") - print(f"Total text annotations: {annotations.texts_length}") - print(f"Total points annotations: {annotations.points_length}") - - # Points should include: bounding boxes + keypoints + skeleton lines - # At least 3 annotations per person (bbox, keypoints, skeleton) - assert annotations.points_length >= len(people) * 3 - - # Text annotations should include confidence, name/id, and keypoint count - assert annotations.texts_length >= len(people) * 3 - - print("\nāœ“ ImageDetections2D.from_pose_detector working correctly!") - - -if __name__ == "__main__": - test_image_detections_2d_with_person() diff --git a/dimos/perception/detection2d/detectors/person/yolo.py b/dimos/perception/detection2d/detectors/person/yolo.py deleted file mode 100644 index fb4fe4769e..0000000000 --- a/dimos/perception/detection2d/detectors/person/yolo.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright 2025 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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np -import torch -from ultralytics import YOLO -from ultralytics.engine.results import Boxes, Keypoints, Results - -from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.types import Detector -from dimos.utils.data import get_data -from dimos.utils.logging_config import setup_logger - -logger = setup_logger("dimos.perception.detection2d.yolo.person") - - -# Type alias for YOLO person detection results -YoloPersonResults = List[Results] - -""" -YOLO Person Detection Results Structure: - -Each Results object in the list contains: - -1. boxes (Boxes object): - - boxes.xyxy: torch.Tensor [N, 4] - bounding boxes in [x1, y1, x2, y2] format - - boxes.xywh: torch.Tensor [N, 4] - boxes in [x_center, y_center, width, height] format - - boxes.conf: torch.Tensor [N] - confidence scores (0-1) - - boxes.cls: torch.Tensor [N] - class IDs (0 for person) - - boxes.xyxyn: torch.Tensor [N, 4] - normalized xyxy coordinates (0-1) - - boxes.xywhn: torch.Tensor [N, 4] - normalized xywh coordinates (0-1) - -2. keypoints (Keypoints object): - - keypoints.xy: torch.Tensor [N, 17, 2] - absolute x,y coordinates for 17 keypoints - - keypoints.conf: torch.Tensor [N, 17] - confidence/visibility scores for each keypoint - - keypoints.xyn: torch.Tensor [N, 17, 2] - normalized coordinates (0-1) - - Keypoint order (COCO format): - 0: nose, 1: left_eye, 2: right_eye, 3: left_ear, 4: right_ear, - 5: left_shoulder, 6: right_shoulder, 7: left_elbow, 8: right_elbow, - 9: left_wrist, 10: right_wrist, 11: left_hip, 12: right_hip, - 13: left_knee, 14: right_knee, 15: left_ankle, 16: right_ankle - -3. Other attributes: - - names: Dict[int, str] - class names mapping {0: 'person'} - - orig_shape: Tuple[int, int] - original image (height, width) - - speed: Dict[str, float] - timing info {'preprocess': ms, 'inference': ms, 'postprocess': ms} - - path: str - image path - - orig_img: np.ndarray - original image array - -Note: All tensor data is on GPU by default. Use .cpu() to move to CPU. -""" -from dimos.perception.detection2d.type.person import Person - - -class YoloPersonDetector(Detector): - def __init__(self, model_path="models_yolo", model_name="yolo11n-pose.pt"): - self.model = YOLO(get_data(model_path) / model_name, task="pose") - - def process_image(self, image: Image) -> YoloPersonResults: - """Process image and return YOLO person detection results. - - Returns: - List of Results objects, typically one per image. - Each Results object contains: - - boxes: Boxes with xyxy, xywh, conf, cls tensors - - keypoints: Keypoints with xy, conf, xyn tensors - - names: {0: 'person'} class mapping - - orig_shape: original image dimensions - - speed: inference timing - """ - return self.model(source=image.to_opencv()) - - def detect_people(self, image: Image) -> List[Person]: - """Process image and return list of Person objects. - - Returns: - List of Person objects with pose keypoints - """ - results = self.process_image(image) - - people = [] - for result in results: - if result.keypoints is None or result.boxes is None: - continue - - # Create Person object for each detection - num_detections = len(result.boxes.xyxy) - for i in range(num_detections): - person = Person.from_yolo(result, i, image) - people.append(person) - - return people - - -def main(): - image = Image.from_file(get_data("cafe.jpg")) - detector = YoloPersonDetector() - - # Get Person objects - people = detector.detect_people(image) - - print(f"Detected {len(people)} people") - for i, person in enumerate(people): - print(f"\nPerson {i}:") - print(f" Confidence: {person.confidence:.3f}") - print(f" Bounding box: {person.bbox}") - cx, cy = person.center - print(f" Center: ({cx:.1f}, {cy:.1f})") - print(f" Size: {person.width:.1f} x {person.height:.1f}") - - # Get specific keypoints - nose_xy, nose_conf = person.get_keypoint("nose") - print(f" Nose: {nose_xy} (conf: {nose_conf:.3f})") - - # Get all visible keypoints - visible = person.get_visible_keypoints(threshold=0.7) - print(f" Visible keypoints (>0.7): {len(visible)}") - for name, xy, conf in visible[:3]: # Show first 3 - print(f" {name}: {xy} (conf: {conf:.3f})") - - -if __name__ == "__main__": - main() diff --git a/dimos/perception/detection2d/detectors/yolo.py b/dimos/perception/detection2d/detectors/yolo.py deleted file mode 100644 index 2d8681f0ef..0000000000 --- a/dimos/perception/detection2d/detectors/yolo.py +++ /dev/null @@ -1,172 +0,0 @@ -# Copyright 2025 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. - -import os - -import cv2 -import onnxruntime -from ultralytics import YOLO - -from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection2d.detectors.types import Detector -from dimos.perception.detection2d.utils import ( - extract_detection_results, - filter_detections, - plot_results, -) -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("dimos.perception.detection2d.yolo_2d_det") - - -class Yolo2DDetector(Detector): - def __init__(self, model_path="models_yolo", model_name="yolo11n.onnx", device="cpu"): - """ - Initialize the YOLO detector. - - Args: - model_path (str): Path to the YOLO model weights in tests/data LFS directory - model_name (str): Name of the YOLO model weights file - device (str): Device to run inference on ('cuda' or 'cpu') - """ - self.device = device - self.model = YOLO(get_data(model_path) / model_name, task="detect") - - module_dir = os.path.dirname(__file__) - self.tracker_config = os.path.join(module_dir, "config", "custom_tracker.yaml") - if is_cuda_available(): - if hasattr(onnxruntime, "preload_dlls"): # Handles CUDA 11 / onnxruntime-gpu<=1.18 - onnxruntime.preload_dlls(cuda=True, cudnn=True) - self.device = "cuda" - logger.debug("Using CUDA for YOLO 2d detector") - else: - self.device = "cpu" - logger.debug("Using CPU for YOLO 2d detector") - - def process_image(self, image: Image): - """ - Process an image and return detection results. - - Args: - image: Input image in BGR format (OpenCV) - - Returns: - tuple: (bboxes, track_ids, class_ids, confidences, names) - - bboxes: list of [x1, y1, x2, y2] coordinates - - track_ids: list of tracking IDs (or -1 if no tracking) - - class_ids: list of class indices - - confidences: list of detection confidences - - names: list of class names - """ - results = self.model.track( - source=image.to_opencv(), - device=self.device, - conf=0.5, - iou=0.6, - persist=True, - verbose=False, - tracker=self.tracker_config, - ) - - if len(results) > 0: - # Extract detection results - bboxes, track_ids, class_ids, confidences, names = extract_detection_results(results[0]) - return bboxes, track_ids, class_ids, confidences, names - - return [], [], [], [], [] - - def visualize_results(self, image, bboxes, track_ids, class_ids, confidences, names): - """ - Generate visualization of detection results. - - Args: - image: Original input image - bboxes: List of bounding boxes - track_ids: List of tracking IDs - class_ids: List of class indices - confidences: List of detection confidences - names: List of class names - - Returns: - Image with visualized detections - """ - return plot_results(image, bboxes, track_ids, class_ids, confidences, names) - - def stop(self): - """ - Clean up resources used by the detector, including tracker threads. - """ - if hasattr(self.model, "predictor") and self.model.predictor is not None: - predictor = self.model.predictor - if hasattr(predictor, "trackers") and predictor.trackers: - for tracker in predictor.trackers: - if hasattr(tracker, "tracker") and hasattr(tracker.tracker, "gmc"): - gmc = tracker.tracker.gmc - if hasattr(gmc, "executor") and gmc.executor is not None: - gmc.executor.shutdown(wait=True) - self.model.predictor = None - - -def main(): - """Example usage of the Yolo2DDetector class.""" - # Initialize video capture - cap = cv2.VideoCapture(0) - - # Initialize detector - detector = Yolo2DDetector() - - enable_person_filter = True - - try: - while cap.isOpened(): - ret, frame = cap.read() - if not ret: - break - - # Process frame - bboxes, track_ids, class_ids, confidences, names = detector.process_image(frame) - - # Apply person filtering if enabled - if enable_person_filter and len(bboxes) > 0: - # Person is class_id 0 in COCO dataset - bboxes, track_ids, class_ids, confidences, names = filter_detections( - bboxes, - track_ids, - class_ids, - confidences, - names, - class_filter=[0], # 0 is the class_id for person - name_filter=["person"], - ) - - # Visualize results - if len(bboxes) > 0: - frame = detector.visualize_results( - frame, bboxes, track_ids, class_ids, confidences, names - ) - - # Display results - cv2.imshow("YOLO Detection", frame) - if cv2.waitKey(1) & 0xFF == ord("q"): - break - - finally: - cap.release() - cv2.destroyAllWindows() - - -if __name__ == "__main__": - main() diff --git a/dimos/perception/detection2d/module2D.py b/dimos/perception/detection2d/module2D.py deleted file mode 100644 index c20e51cd9c..0000000000 --- a/dimos/perception/detection2d/module2D.py +++ /dev/null @@ -1,111 +0,0 @@ -# Copyright 2025 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. -import functools -from dataclasses import dataclass -from typing import Any, Callable, Optional - -import numpy as np -from dimos_lcm.foxglove_msgs.ImageAnnotations import ( - ImageAnnotations, -) -from reactivex import operators as ops -from reactivex.observable import Observable -from reactivex.subject import Subject - -from dimos.core import In, Module, Out, rpc -from dimos.msgs.sensor_msgs import Image -from dimos.msgs.sensor_msgs.Image import sharpness_barrier -from dimos.msgs.vision_msgs import Detection2DArray -from dimos.perception.detection2d.type import ImageDetections2D -from dimos.perception.detection2d.detectors import Detector, Yolo2DDetector -from dimos.perception.detection2d.detectors.person.yolo import YoloPersonDetector -from dimos.perception.detection2d.type import ImageDetections2D -from dimos.utils.decorators.decorators import simple_mcache -from dimos.utils.reactive import backpressure - - -@dataclass -class Config: - max_freq: float = 5 # hz - detector: Optional[Callable[[Any], Detector]] = lambda: Yolo2DDetector() - - -class Detection2DModule(Module): - config: Config - detector: Detector - - image: In[Image] = None # type: ignore - - detections: Out[Detection2DArray] = None # type: ignore - annotations: Out[ImageAnnotations] = None # type: ignore - - # just for visualization, emits latest top 3 detections in a frame - detected_image_0: Out[Image] = None # type: ignore - detected_image_1: Out[Image] = None # type: ignore - detected_image_2: Out[Image] = None # type: ignore - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.config: Config = Config(**kwargs) - self.detector = self.config.detector() - self.vlm_detections_subject = Subject() - - def process_image_frame(self, image: Image) -> ImageDetections2D: - # Use person detection specifically if it's a YoloPersonDetector - if isinstance(self.detector, YoloPersonDetector): - people = self.detector.detect_people(image) - return ImageDetections2D.from_pose_detector(image, people) - else: - # Fallback to generic dettection for other detectors - return ImageDetections2D.from_bbox_detector(image, self.detector.process_image(image)) - - @simple_mcache - def sharp_image_stream(self) -> Observable[Image]: - return backpressure( - self.image.pure_observable().pipe( - sharpness_barrier(self.config.max_freq), - ) - ) - - @simple_mcache - def detection_stream_2d(self) -> Observable[ImageDetections2D]: - # return self.vlm_detections_subject - # Regular detection stream from the detector - regular_detections = self.sharp_image_stream().pipe(ops.map(self.process_image_frame)) - # Merge with VL model detections - return backpressure(regular_detections.pipe(ops.merge(self.vlm_detections_subject))) - - @rpc - def start(self): - super().start() - unsub = self.detection_stream_2d().subscribe( - lambda det: self.detections.publish(det.to_ros_detection2d_array()) - ) - self._disposables.add(unsub) - - unsub = self.detection_stream_2d().subscribe( - lambda det: self.annotations.publish(det.to_foxglove_annotations()) - ) - self._disposables.add(unsub) - - def publish_cropped_images(detections: ImageDetections2D): - for index, detection in enumerate(detections[:3]): - image_topic = getattr(self, "detected_image_" + str(index)) - image_topic.publish(detection.cropped_image()) - - self.detection_stream_2d().subscribe(publish_cropped_images) - - @rpc - def stop(self) -> None: - super().stop() diff --git a/dimos/perception/detection2d/type/__init__.py b/dimos/perception/detection2d/type/__init__.py deleted file mode 100644 index aee8597d5c..0000000000 --- a/dimos/perception/detection2d/type/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -from dimos.perception.detection2d.type.detection2d import ( - Detection2D, - Detection2DBBox, - ImageDetections2D, - InconvinientDetectionFormat, -) -from dimos.perception.detection2d.type.detection3d import ( - Detection3D, - ImageDetections3D, -) -from dimos.perception.detection2d.type.detection3dpc import ( - Detection3DPC, - ImageDetections3DPC, -) -from dimos.perception.detection2d.type.imageDetections import ImageDetections, TableStr -from dimos.perception.detection2d.type.person import Person diff --git a/dimos/perception/detection2d/type/detection3dpc.py b/dimos/perception/detection2d/type/detection3dpc.py deleted file mode 100644 index 44d242de9e..0000000000 --- a/dimos/perception/detection2d/type/detection3dpc.py +++ /dev/null @@ -1,247 +0,0 @@ -# Copyright 2025 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 - -import functools -from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, TypeVar - -import numpy as np -from dimos_lcm.sensor_msgs import CameraInfo -from lcm_msgs.builtin_interfaces import Duration -from lcm_msgs.foxglove_msgs import CubePrimitive, SceneEntity, SceneUpdate, TextPrimitive -from lcm_msgs.geometry_msgs import Point, Pose, Quaternion -from lcm_msgs.geometry_msgs import Vector3 as LCMVector3 - -from dimos.msgs.foxglove_msgs.Color import Color -from dimos.msgs.geometry_msgs import PoseStamped, Transform, Vector3 -from dimos.msgs.sensor_msgs import PointCloud2 -from dimos.perception.detection2d.type.detection2d import Detection2D -from dimos.perception.detection2d.type.detection3d import Detection3D -from dimos.perception.detection2d.type.imageDetections import ImageDetections -from dimos.types.timestamped import to_ros_stamp - -Detection3DPCFilter = Callable[ - [Detection2D, PointCloud2, CameraInfo, Transform], Optional["Detection3DPC"] -] - - -def height_filter(height=0.1) -> Detection3DPCFilter: - return lambda det, pc, ci, tf: pc.filter_by_height(height) - - -def statistical(nb_neighbors=40, std_ratio=0.5) -> Detection3DPCFilter: - def filter_func( - det: Detection2D, pc: PointCloud2, ci: CameraInfo, tf: Transform - ) -> Optional[PointCloud2]: - try: - statistical, removed = pc.pointcloud.remove_statistical_outlier( - nb_neighbors=nb_neighbors, std_ratio=std_ratio - ) - return PointCloud2(statistical, pc.frame_id, pc.ts) - except Exception as e: - # print("statistical filter failed:", e) - return None - - return filter_func - - -def raycast() -> Detection3DPCFilter: - def filter_func( - det: Detection2D, pc: PointCloud2, ci: CameraInfo, tf: Transform - ) -> Optional[PointCloud2]: - try: - camera_pos = tf.inverse().translation - camera_pos_np = camera_pos.to_numpy() - _, visible_indices = pc.pointcloud.hidden_point_removal(camera_pos_np, radius=100.0) - visible_pcd = pc.pointcloud.select_by_index(visible_indices) - return PointCloud2(visible_pcd, pc.frame_id, pc.ts) - except Exception as e: - # print("raycast filter failed:", e) - return None - - return filter_func - - -def radius_outlier(min_neighbors: int = 20, radius: float = 0.3) -> Detection3DPCFilter: - """ - Remove isolated points: keep only points that have at least `min_neighbors` - neighbors within `radius` meters (same units as your point cloud). - """ - - def filter_func( - det: Detection2D, pc: PointCloud2, ci: CameraInfo, tf: Transform - ) -> Optional[PointCloud2]: - filtered_pcd, removed = pc.pointcloud.remove_radius_outlier( - nb_points=min_neighbors, radius=radius - ) - return PointCloud2(filtered_pcd, pc.frame_id, pc.ts) - - return filter_func - - -@dataclass -class Detection3DPC(Detection3D): - pointcloud: PointCloud2 - - @classmethod - def from_2d( - cls, - det: Detection2D, - world_pointcloud: PointCloud2, - camera_info: CameraInfo, - world_to_optical_transform: Transform, - # filters are to be adjusted based on the sensor noise characteristics if feeding - # sensor data directly - filters: list[Callable[[PointCloud2], PointCloud2]] = [ - # height_filter(0.1), - raycast(), - radius_outlier(), - statistical(), - ], - ) -> Optional["Detection3D"]: - """Create a Detection3D from a 2D detection by projecting world pointcloud. - - This method handles: - 1. Projecting world pointcloud to camera frame - 2. Filtering points within the 2D detection bounding box - 3. Cleaning up the pointcloud (height filter, outlier removal) - 4. Hidden point removal from camera perspective - - Args: - det: The 2D detection - world_pointcloud: Full pointcloud in world frame - camera_info: Camera calibration info - world_to_camerlka_transform: Transform from world to camera frame - filters: List of functions to apply to the pointcloud for filtering - Returns: - Detection3D with filtered pointcloud, or None if no valid points - """ - # Extract camera parameters - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] - image_width = camera_info.width - image_height = camera_info.height - - camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) - - # Convert pointcloud to numpy array - world_points = world_pointcloud.as_numpy() - - # Project points to camera frame - points_homogeneous = np.hstack([world_points, np.ones((world_points.shape[0], 1))]) - extrinsics_matrix = world_to_optical_transform.to_matrix() - points_camera = (extrinsics_matrix @ points_homogeneous.T).T - - # Filter out points behind the camera - valid_mask = points_camera[:, 2] > 0 - points_camera = points_camera[valid_mask] - world_points = world_points[valid_mask] - - if len(world_points) == 0: - return None - - # Project to 2D - points_2d_homogeneous = (camera_matrix @ points_camera[:, :3].T).T - points_2d = points_2d_homogeneous[:, :2] / points_2d_homogeneous[:, 2:3] - - # Filter points within image bounds - in_image_mask = ( - (points_2d[:, 0] >= 0) - & (points_2d[:, 0] < image_width) - & (points_2d[:, 1] >= 0) - & (points_2d[:, 1] < image_height) - ) - points_2d = points_2d[in_image_mask] - world_points = world_points[in_image_mask] - - if len(world_points) == 0: - return None - - # Extract bbox from Detection2D - x_min, y_min, x_max, y_max = det.bbox - - # Find points within this detection box (with small margin) - margin = 5 # pixels - in_box_mask = ( - (points_2d[:, 0] >= x_min - margin) - & (points_2d[:, 0] <= x_max + margin) - & (points_2d[:, 1] >= y_min - margin) - & (points_2d[:, 1] <= y_max + margin) - ) - - detection_points = world_points[in_box_mask] - - if detection_points.shape[0] == 0: - # print(f"No points found in detection bbox after projection. {det.name}") - return None - - # Create initial pointcloud for this detection - initial_pc = PointCloud2.from_numpy( - detection_points, - frame_id=world_pointcloud.frame_id, - timestamp=world_pointcloud.ts, - ) - - # Apply filters - each filter needs all 4 arguments - detection_pc = initial_pc - for filter_func in filters: - result = filter_func(det, detection_pc, camera_info, world_to_optical_transform) - if result is None: - return None - detection_pc = result - - # Final check for empty pointcloud - if len(detection_pc.pointcloud.points) == 0: - return None - - # Create Detection3D with filtered pointcloud - return cls( - image=det.image, - bbox=det.bbox, - track_id=det.track_id, - class_id=det.class_id, - confidence=det.confidence, - name=det.name, - ts=det.ts, - pointcloud=detection_pc, - transform=world_to_optical_transform, - frame_id=world_pointcloud.frame_id, - ) - - -class ImageDetections3DPC(ImageDetections[Detection3DPC]): - """Specialized class for 3D detections in an image.""" - - def to_foxglove_scene_update(self) -> "SceneUpdate": - """Convert all detections to a Foxglove SceneUpdate message. - - Returns: - SceneUpdate containing SceneEntity objects for all detections - """ - - # Create SceneUpdate message with all detections - scene_update = SceneUpdate() - scene_update.deletions_length = 0 - scene_update.deletions = [] - scene_update.entities = [] - - # Process each detection - for i, detection in enumerate(self.detections): - entity = detection.to_foxglove_scene_entity(entity_id=f"detection_{detection.name}_{i}") - scene_update.entities.append(entity) - - scene_update.entities_length = len(scene_update.entities) - return scene_update diff --git a/dimos/protocol/service/lcmservice.py b/dimos/protocol/service/lcmservice.py index bc3f7317b7..2228a671fc 100644 --- a/dimos/protocol/service/lcmservice.py +++ b/dimos/protocol/service/lcmservice.py @@ -21,7 +21,7 @@ import traceback from dataclasses import dataclass from functools import cache -from typing import Any, Callable, Optional, Protocol, runtime_checkable +from typing import Optional, Protocol, runtime_checkable import lcm diff --git a/dimos/robot/unitree_webrtc/connection.py b/dimos/robot/unitree_webrtc/connection.py index f7cf683a2a..8ddc77ac63 100644 --- a/dimos/robot/unitree_webrtc/connection.py +++ b/dimos/robot/unitree_webrtc/connection.py @@ -38,6 +38,7 @@ from dimos.robot.unitree_webrtc.type.lidar import LidarMessage from dimos.robot.unitree_webrtc.type.lowstate import LowStateMsg from dimos.robot.unitree_webrtc.type.odometry import Odometry +from dimos.utils.decorators.decorators import simple_mcache from dimos.utils.reactive import backpressure, callback_to_observable VideoMessage: TypeAlias = np.ndarray[tuple[int, int, Literal[3]], np.uint8] @@ -224,15 +225,15 @@ def publish_request(self, topic: str, data: dict): ) return future.result() - @functools.cache + @simple_mcache def raw_lidar_stream(self) -> Subject[LidarMessage]: return backpressure(self.unitree_sub_stream(RTC_TOPIC["ULIDAR_ARRAY"])) - @functools.cache + @simple_mcache def raw_odom_stream(self) -> Subject[Pose]: return backpressure(self.unitree_sub_stream(RTC_TOPIC["ROBOTODOM"])) - @functools.cache + @simple_mcache def lidar_stream(self) -> Subject[LidarMessage]: return backpressure( self.raw_lidar_stream().pipe( @@ -240,22 +241,23 @@ def lidar_stream(self) -> Subject[LidarMessage]: ) ) - @functools.cache + @simple_mcache def tf_stream(self) -> Subject[Transform]: base_link = functools.partial(Transform.from_pose, "base_link") return backpressure(self.odom_stream().pipe(ops.map(base_link))) - @functools.cache + @simple_mcache def odom_stream(self) -> Subject[Pose]: return backpressure(self.raw_odom_stream().pipe(ops.map(Odometry.from_msg))) - @functools.cache + @simple_mcache def video_stream(self) -> Observable[Image]: return backpressure( self.raw_video_stream().pipe( ops.filter(lambda frame: frame is not None), ops.map( lambda frame: Image.from_numpy( + # np.ascontiguousarray(frame.to_ndarray("rgb24")), frame.to_ndarray(format="rgb24"), frame_id="camera_optical", ) @@ -263,7 +265,7 @@ def video_stream(self) -> Observable[Image]: ) ) - @functools.cache + @simple_mcache def lowstate_stream(self) -> Subject[LowStateMsg]: return backpressure(self.unitree_sub_stream(RTC_TOPIC["LOW_STATE"])) @@ -306,7 +308,7 @@ def color(self, color: VUI_COLOR = VUI_COLOR.RED, colortime: int = 60) -> bool: }, ) - @functools.lru_cache(maxsize=None) + @simple_mcache def raw_video_stream(self) -> Observable[VideoMessage]: subject: Subject[VideoMessage] = Subject() stop_event = threading.Event() diff --git a/dimos/robot/unitree_webrtc/modular/connection_module.py b/dimos/robot/unitree_webrtc/modular/connection_module.py index 30413bf182..5950282f0b 100644 --- a/dimos/robot/unitree_webrtc/modular/connection_module.py +++ b/dimos/robot/unitree_webrtc/modular/connection_module.py @@ -30,7 +30,8 @@ from reactivex.observable import Observable from dimos.agents2 import Agent, Output, Reducer, Stream, skill -from dimos.core import DimosCluster, LCMTransport, Module, ModuleConfig, Out, rpc, In +from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE +from dimos.core import DimosCluster, In, LCMTransport, Module, ModuleConfig, Out, pSHMTransport, rpc from dimos.msgs.foxglove_msgs import ImageAnnotations from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Transform, Twist, Vector3 from dimos.msgs.sensor_msgs.Image import Image, sharpness_window @@ -176,7 +177,7 @@ def start(self): case "webrtc": self.connection = UnitreeWebRTCConnection(**self.connection_config) case "fake": - self.connection = FakeRTC(**self.connection_config) + self.connection = FakeRTC(**self.connection_config, seek=12.0) case "mujoco": from dimos.robot.unitree_webrtc.mujoco_connection import MujocoConnection @@ -235,10 +236,19 @@ def _odom_to_tf(self, odom: PoseStamped) -> List[Transform]: ts=odom.ts, ) + sensor = Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="world", + child_frame_id="sensor", + ts=odom.ts, + ) + return [ Transform.from_pose("base_link", odom), camera_link, camera_optical, + sensor, ] def _publish_tf(self, msg): @@ -314,10 +324,19 @@ def deploy_connection(dimos: DimosCluster, **kwargs): **kwargs, ) - connection.lidar.transport = LCMTransport("/lidar", LidarMessage) connection.odom.transport = LCMTransport("/odom", PoseStamped) + + connection.video.transport = pSHMTransport( + "/image", default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE + ) + + connection.lidar.transport = pSHMTransport( + "/lidar", default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE + ) + connection.video.transport = LCMTransport("/image", Image) - connection.movecmd.transport = LCMTransport("/cmd_vel", Vector3) + connection.lidar.transport = LCMTransport("/lidar", LidarMessage) + connection.movecmd.transport = LCMTransport("/cmd_vel", Twist) connection.camera_info.transport = LCMTransport("/camera_info", CameraInfo) return connection diff --git a/dimos/robot/unitree_webrtc/modular/detect.py b/dimos/robot/unitree_webrtc/modular/detect.py index 7d0ded7ac8..3f6c2c04b2 100644 --- a/dimos/robot/unitree_webrtc/modular/detect.py +++ b/dimos/robot/unitree_webrtc/modular/detect.py @@ -135,7 +135,7 @@ def broadcast( def process_data(): from dimos.msgs.sensor_msgs import Image - from dimos.perception.detection2d.module import Detect2DModule, build_imageannotations + from dimos.perception.detection.module2D import Detection2DModule, build_imageannotations from dimos.robot.unitree_webrtc.type.lidar import LidarMessage from dimos.robot.unitree_webrtc.type.odometry import Odometry from dimos.utils.data import get_data @@ -155,7 +155,7 @@ def attach_frame_id(image: Image) -> Image: video_frame = attach_frame_id(video_store.find_closest(target, tolerance=1)) odom_frame = odom_store.find_closest(target, tolerance=1) - detector = Detect2DModule() + detector = Detection2DModule() detections = detector.detect(video_frame) annotations = build_imageannotations(detections) diff --git a/dimos/robot/unitree_webrtc/modular/ivan_unitree.py b/dimos/robot/unitree_webrtc/modular/ivan_unitree.py index 73927cf248..948dccaa16 100644 --- a/dimos/robot/unitree_webrtc/modular/ivan_unitree.py +++ b/dimos/robot/unitree_webrtc/modular/ivan_unitree.py @@ -15,19 +15,22 @@ import logging import time -from dimos_lcm.sensor_msgs import CameraInfo -from lcm_msgs.foxglove_msgs import SceneUpdate +from dimos_lcm.foxglove_msgs import SceneUpdate from dimos.agents2.spec import Model, Provider from dimos.core import LCMTransport, start # from dimos.msgs.detection2d import Detection2DArray from dimos.msgs.foxglove_msgs import ImageAnnotations +from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import Image, PointCloud2 from dimos.msgs.vision_msgs import Detection2DArray -from dimos.perception.detection2d import Detection3DModule -from dimos.perception.detection2d.moduleDB import ObjectDBModule +from dimos.perception.detection.module2D import Detection2DModule +from dimos.perception.detection.module3D import Detection3DModule +from dimos.perception.detection.person_tracker import PersonTracker +from dimos.perception.detection.reid import ReidModule from dimos.protocol.pubsub import lcm +from dimos.robot.foxglove_bridge import FoxgloveBridge from dimos.robot.unitree_webrtc.modular import deploy_connection, deploy_navigation from dimos.robot.unitree_webrtc.modular.connection_module import ConnectionModule from dimos.utils.logging_config import setup_logger @@ -36,46 +39,64 @@ def detection_unitree(): - dimos = start(6) + dimos = start(8) connection = deploy_connection(dimos) - # mapper = deploy_navigation(dimos, connection) - # mapper.start() def goto(pose): print("NAVIGATION REQUESTED:", pose) return True - module3D = dimos.deploy( - ObjectDBModule, - goto=goto, + detector = dimos.deploy( + Detection2DModule, + # goto=goto, camera_info=ConnectionModule._camera_info(), ) - module3D.image.connect(connection.video) - # module3D.pointcloud.connect(mapper.global_map) - module3D.pointcloud.connect(connection.lidar) + detector.image.connect(connection.video) + # detector.pointcloud.connect(mapper.global_map) + # detector.pointcloud.connect(connection.lidar) - module3D.annotations.transport = LCMTransport("/annotations", ImageAnnotations) - module3D.detections.transport = LCMTransport("/detections", Detection2DArray) + detector.annotations.transport = LCMTransport("/annotations", ImageAnnotations) + detector.detections.transport = LCMTransport("/detections", Detection2DArray) - module3D.detected_pointcloud_0.transport = LCMTransport("/detected/pointcloud/0", PointCloud2) - module3D.detected_pointcloud_1.transport = LCMTransport("/detected/pointcloud/1", PointCloud2) - module3D.detected_pointcloud_2.transport = LCMTransport("/detected/pointcloud/2", PointCloud2) + # detector.detected_pointcloud_0.transport = LCMTransport("/detected/pointcloud/0", PointCloud2) + # detector.detected_pointcloud_1.transport = LCMTransport("/detected/pointcloud/1", PointCloud2) + # detector.detected_pointcloud_2.transport = LCMTransport("/detected/pointcloud/2", PointCloud2) - module3D.detected_image_0.transport = LCMTransport("/detected/image/0", Image) - module3D.detected_image_1.transport = LCMTransport("/detected/image/1", Image) - module3D.detected_image_2.transport = LCMTransport("/detected/image/2", Image) + detector.detected_image_0.transport = LCMTransport("/detected/image/0", Image) + detector.detected_image_1.transport = LCMTransport("/detected/image/1", Image) + detector.detected_image_2.transport = LCMTransport("/detected/image/2", Image) + # detector.scene_update.transport = LCMTransport("/scene_update", SceneUpdate) - module3D.scene_update.transport = LCMTransport("/scene_update", SceneUpdate) + # reidModule = dimos.deploy(ReidModule) - module3D.start() + # reidModule.image.connect(connection.video) + # reidModule.detections.connect(detector.detections) + # reidModule.annotations.transport = LCMTransport("/reid/annotations", ImageAnnotations) + + # nav = deploy_navigation(dimos, connection) + + # person_tracker = dimos.deploy(PersonTracker, cameraInfo=ConnectionModule._camera_info()) + # person_tracker.image.connect(connection.video) + # person_tracker.detections.connect(detector.detections) + # person_tracker.target.transport = LCMTransport("/goal_request", PoseStamped) + + reid = dimos.deploy(ReidModule) + + reid.image.connect(connection.video) + reid.detections.connect(detector.detections) + reid.annotations.transport = LCMTransport("/reid/annotations", ImageAnnotations) + + detector.start() + # person_tracker.start() connection.start() + reid.start() from dimos.agents2 import Agent, Output, Reducer, Stream, skill from dimos.agents2.cli.human import HumanInput agent = Agent( - system_prompt="You are a helpful assistant for controlling a Unitree Go2 robot. ", + system_prompt="You are a helpful assistant for controlling a Unitree Go2 robot.", model=Model.GPT_4O, # Could add CLAUDE models to enum provider=Provider.OPENAI, # Would need ANTHROPIC provider ) @@ -83,13 +104,23 @@ def goto(pose): human_input = dimos.deploy(HumanInput) agent.register_skills(human_input) # agent.register_skills(connection) - agent.register_skills(module3D) + agent.register_skills(detector) + + bridge = FoxgloveBridge( + shm_channels=[ + "/image#sensor_msgs.Image", + "/lidar#sensor_msgs.PointCloud2", + ] + ) + # bridge = FoxgloveBridge() + time.sleep(1) + bridge.start() # agent.run_implicit_skill("video_stream_tool") - agent.run_implicit_skill("human") + # agent.run_implicit_skill("human") - agent.start() - agent.loop_thread() + # agent.start() + # agent.loop_thread() try: while True: diff --git a/dimos/robot/unitree_webrtc/modular/navigation.py b/dimos/robot/unitree_webrtc/modular/navigation.py index c37cac700a..f16fd29816 100644 --- a/dimos/robot/unitree_webrtc/modular/navigation.py +++ b/dimos/robot/unitree_webrtc/modular/navigation.py @@ -15,7 +15,7 @@ from dimos_lcm.std_msgs import Bool, String from dimos.core import LCMTransport -from dimos.msgs.geometry_msgs import PoseStamped, Vector3 +from dimos.msgs.geometry_msgs import PoseStamped, Twist, Vector3 from dimos.msgs.nav_msgs import OccupancyGrid, Path from dimos.navigation.bt_navigator.navigator import BehaviorTreeNavigator from dimos.navigation.frontier_exploration import WavefrontFrontierExplorer @@ -27,7 +27,7 @@ def deploy_navigation(dimos, connection): - mapper = dimos.deploy(Map, voxel_size=0.5, cost_resolution=0.05, global_publish_interval=0.5) + mapper = dimos.deploy(Map, voxel_size=0.5, cost_resolution=0.05, global_publish_interval=2.5) mapper.lidar.connect(connection.lidar) mapper.global_map.transport = LCMTransport("/global_map", LidarMessage) mapper.global_costmap.transport = LCMTransport("/global_costmap", OccupancyGrid) @@ -49,7 +49,7 @@ def deploy_navigation(dimos, connection): navigator.navigation_state.transport = LCMTransport("/navigation_state", String) navigator.global_costmap.transport = LCMTransport("/global_costmap", OccupancyGrid) global_planner.path.transport = LCMTransport("/global_path", Path) - local_planner.cmd_vel.transport = LCMTransport("/cmd_vel", Vector3) + local_planner.cmd_vel.transport = LCMTransport("/cmd_vel", Twist) frontier_explorer.goal_request.transport = LCMTransport("/goal_request", PoseStamped) frontier_explorer.goal_reached.transport = LCMTransport("/goal_reached", Bool) frontier_explorer.explore_cmd.transport = LCMTransport("/explore_cmd", Bool) @@ -83,4 +83,11 @@ def deploy_navigation(dimos, connection): navigator.start() websocket_vis.start() - return mapper + return { + "mapper": mapper, + "global_planner": global_planner, + "local_planner": local_planner, + "navigator": navigator, + "frontier_explorer": frontier_explorer, + "websocket_vis": websocket_vis, + } diff --git a/dimos/robot/unitree_webrtc/unitree_b1/test_connection.py b/dimos/robot/unitree_webrtc/unitree_b1/test_connection.py index a9451acdf0..57227e6e23 100644 --- a/dimos/robot/unitree_webrtc/unitree_b1/test_connection.py +++ b/dimos/robot/unitree_webrtc/unitree_b1/test_connection.py @@ -290,6 +290,9 @@ def test_mode_changes_with_watchdog(self): conn.watchdog_thread = threading.Thread(target=conn._watchdog_loop, daemon=True) conn.watchdog_thread.start() + # Give threads time to initialize + time.sleep(0.05) + # Send walk command twist = TwistStamped( ts=time.time(), @@ -301,8 +304,8 @@ def test_mode_changes_with_watchdog(self): assert conn.current_mode == 2 assert conn._current_cmd.ly == 1.0 - # Wait for timeout first - time.sleep(0.25) + # Wait for timeout first (0.2s timeout + 0.15s margin for reliability) + time.sleep(0.35) assert conn.timeout_active assert conn._current_cmd.ly == 0.0 # Watchdog zeroed it diff --git a/dimos/robot/unitree_webrtc/unitree_g1.py b/dimos/robot/unitree_webrtc/unitree_g1.py index f319e2c87c..da63687072 100644 --- a/dimos/robot/unitree_webrtc/unitree_g1.py +++ b/dimos/robot/unitree_webrtc/unitree_g1.py @@ -21,6 +21,16 @@ import logging import os import time +from typing import Optional + +from dimos_lcm.foxglove_msgs import SceneUpdate +from geometry_msgs.msg import PoseStamped as ROSPoseStamped +from geometry_msgs.msg import TwistStamped as ROSTwistStamped +from nav_msgs.msg import Odometry as ROSOdometry +from reactivex.disposable import Disposable +from sensor_msgs.msg import Joy as ROSJoy +from sensor_msgs.msg import PointCloud2 as ROSPointCloud2 +from tf2_msgs.msg import TFMessage as ROSTFMessage from dimos import core from dimos.agents2 import Agent @@ -47,7 +57,7 @@ from dimos.msgs.std_msgs.Bool import Bool from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.msgs.vision_msgs import Detection2DArray -from dimos.perception.detection2d.moduleDB import ObjectDBModule +from dimos.perception.detection.moduleDB import ObjectDBModule from dimos.perception.spatial_perception import SpatialMemory from dimos.protocol import pubsub from dimos.protocol.pubsub.lcmpubsub import LCM @@ -62,14 +72,6 @@ from dimos.types.robot_capabilities import RobotCapability from dimos.utils.logging_config import setup_logger from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule -from dimos_lcm.foxglove_msgs import SceneUpdate -from geometry_msgs.msg import TwistStamped as ROSTwistStamped -from nav_msgs.msg import Odometry as ROSOdometry -from reactivex.disposable import Disposable -from sensor_msgs.msg import Joy as ROSJoy -from sensor_msgs.msg import PointCloud2 as ROSPointCloud2 -from tf2_msgs.msg import TFMessage as ROSTFMessage -from typing import Optional logger = setup_logger("dimos.robot.unitree_webrtc.unitree_g1", level=logging.INFO) @@ -417,7 +419,6 @@ def _deploy_ros_bridge(self): "/tf", TFMessage, ROSTFMessage, direction=BridgeDirection.ROS_TO_DIMOS ) - from geometry_msgs.msg import PoseStamped as ROSPoseStamped from std_msgs.msg import Bool as ROSBool from dimos.msgs.std_msgs import Bool diff --git a/dimos/utils/decorators/__init__.py b/dimos/utils/decorators/__init__.py index 22ad478a00..ee17260c20 100644 --- a/dimos/utils/decorators/__init__.py +++ b/dimos/utils/decorators/__init__.py @@ -1,11 +1,12 @@ """Decorators and accumulators for rate limiting and other utilities.""" from .accumulators import Accumulator, LatestAccumulator, RollingAverageAccumulator -from .decorators import limit +from .decorators import limit, retry __all__ = [ "Accumulator", "LatestAccumulator", "RollingAverageAccumulator", "limit", + "retry", ] diff --git a/dimos/utils/decorators/decorators.py b/dimos/utils/decorators/decorators.py index c54e3530e1..067251e5c6 100644 --- a/dimos/utils/decorators/decorators.py +++ b/dimos/utils/decorators/decorators.py @@ -15,7 +15,7 @@ import threading import time from functools import wraps -from typing import Callable, Optional +from typing import Callable, Optional, Type from .accumulators import Accumulator, LatestAccumulator @@ -143,3 +143,59 @@ def getter(self): return getattr(self, attr_name) return getter + + +def retry(max_retries: int = 3, on_exception: Type[Exception] = Exception, delay: float = 0.0): + """ + Decorator that retries a function call if it raises an exception. + + Args: + max_retries: Maximum number of retry attempts (default: 3) + on_exception: Exception type to catch and retry on (default: Exception) + delay: Fixed delay in seconds between retries (default: 0.0) + + Returns: + Decorated function that will retry on failure + + Example: + @retry(max_retries=5, on_exception=ConnectionError, delay=0.5) + def connect_to_server(): + # connection logic that might fail + pass + + @retry() # Use defaults: 3 retries on any Exception, no delay + def risky_operation(): + # might fail occasionally + pass + """ + if max_retries < 0: + raise ValueError("max_retries must be non-negative") + if delay < 0: + raise ValueError("delay must be non-negative") + + def decorator(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + last_exception = None + + for attempt in range(max_retries + 1): + try: + return func(*args, **kwargs) + except on_exception as e: + last_exception = e + if attempt < max_retries: + # Still have retries left + if delay > 0: + time.sleep(delay) + continue + else: + # Out of retries, re-raise the last exception + raise + + # This should never be reached, but just in case + if last_exception: + raise last_exception + + return wrapper + + return decorator diff --git a/dimos/utils/decorators/test_decorators.py b/dimos/utils/decorators/test_decorators.py index 2a9162c762..133fab97c2 100644 --- a/dimos/utils/decorators/test_decorators.py +++ b/dimos/utils/decorators/test_decorators.py @@ -16,7 +16,7 @@ import pytest -from dimos.utils.decorators import LatestAccumulator, RollingAverageAccumulator, limit +from dimos.utils.decorators import LatestAccumulator, RollingAverageAccumulator, limit, retry def test_limit(): @@ -77,3 +77,186 @@ def process(value: float, label: str = ""): # Should see the average of accumulated values assert calls == [(10.0, "first"), (25.0, "third")] # (20+30)/2 = 25 + + +def test_retry_success_after_failures(): + """Test that retry decorator retries on failure and eventually succeeds.""" + attempts = [] + + @retry(max_retries=3) + def flaky_function(fail_times=2): + attempts.append(len(attempts)) + if len(attempts) <= fail_times: + raise ValueError(f"Attempt {len(attempts)} failed") + return "success" + + result = flaky_function() + assert result == "success" + assert len(attempts) == 3 # Failed twice, succeeded on third attempt + + +def test_retry_exhausted(): + """Test that retry decorator raises exception when retries are exhausted.""" + attempts = [] + + @retry(max_retries=2) + def always_fails(): + attempts.append(len(attempts)) + raise RuntimeError(f"Attempt {len(attempts)} failed") + + with pytest.raises(RuntimeError) as exc_info: + always_fails() + + assert "Attempt 3 failed" in str(exc_info.value) + assert len(attempts) == 3 # Initial attempt + 2 retries + + +def test_retry_specific_exception(): + """Test that retry only catches specified exception types.""" + attempts = [] + + @retry(max_retries=3, on_exception=ValueError) + def raises_different_exceptions(): + attempts.append(len(attempts)) + if len(attempts) == 1: + raise ValueError("First attempt") + elif len(attempts) == 2: + raise TypeError("Second attempt - should not be retried") + return "success" + + # Should fail on TypeError (not retried) + with pytest.raises(TypeError) as exc_info: + raises_different_exceptions() + + assert "Second attempt" in str(exc_info.value) + assert len(attempts) == 2 # First attempt with ValueError, second with TypeError + + +def test_retry_no_failures(): + """Test that retry decorator works when function succeeds immediately.""" + attempts = [] + + @retry(max_retries=5) + def always_succeeds(): + attempts.append(len(attempts)) + return "immediate success" + + result = always_succeeds() + assert result == "immediate success" + assert len(attempts) == 1 # Only one attempt needed + + +def test_retry_with_delay(): + """Test that retry decorator applies delay between attempts.""" + attempts = [] + times = [] + + @retry(max_retries=2, delay=0.1) + def delayed_failures(): + times.append(time.time()) + attempts.append(len(attempts)) + if len(attempts) < 2: + raise ValueError(f"Attempt {len(attempts)}") + return "success" + + start = time.time() + result = delayed_failures() + duration = time.time() - start + + assert result == "success" + assert len(attempts) == 2 + assert duration >= 0.1 # At least one delay occurred + + # Check that delays were applied + if len(times) >= 2: + assert times[1] - times[0] >= 0.1 + + +def test_retry_zero_retries(): + """Test retry with max_retries=0 (no retries, just one attempt).""" + attempts = [] + + @retry(max_retries=0) + def single_attempt(): + attempts.append(len(attempts)) + raise ValueError("Failed") + + with pytest.raises(ValueError): + single_attempt() + + assert len(attempts) == 1 # Only the initial attempt + + +def test_retry_invalid_parameters(): + """Test that retry decorator validates parameters.""" + with pytest.raises(ValueError): + + @retry(max_retries=-1) + def invalid_retries(): + pass + + with pytest.raises(ValueError): + + @retry(delay=-0.5) + def invalid_delay(): + pass + + +def test_retry_with_methods(): + """Test that retry decorator works with class methods, instance methods, and static methods.""" + + class TestClass: + def __init__(self): + self.instance_attempts = [] + self.instance_value = 42 + + @retry(max_retries=3) + def instance_method(self, fail_times=2): + """Test retry on instance method.""" + self.instance_attempts.append(len(self.instance_attempts)) + if len(self.instance_attempts) <= fail_times: + raise ValueError(f"Instance attempt {len(self.instance_attempts)} failed") + return f"instance success with value {self.instance_value}" + + @classmethod + @retry(max_retries=2) + def class_method(cls, attempts_list, fail_times=1): + """Test retry on class method.""" + attempts_list.append(len(attempts_list)) + if len(attempts_list) <= fail_times: + raise ValueError(f"Class attempt {len(attempts_list)} failed") + return f"class success from {cls.__name__}" + + @staticmethod + @retry(max_retries=2) + def static_method(attempts_list, fail_times=1): + """Test retry on static method.""" + attempts_list.append(len(attempts_list)) + if len(attempts_list) <= fail_times: + raise ValueError(f"Static attempt {len(attempts_list)} failed") + return "static success" + + # Test instance method + obj = TestClass() + result = obj.instance_method() + assert result == "instance success with value 42" + assert len(obj.instance_attempts) == 3 # Failed twice, succeeded on third + + # Test class method + class_attempts = [] + result = TestClass.class_method(class_attempts) + assert result == "class success from TestClass" + assert len(class_attempts) == 2 # Failed once, succeeded on second + + # Test static method + static_attempts = [] + result = TestClass.static_method(static_attempts) + assert result == "static success" + assert len(static_attempts) == 2 # Failed once, succeeded on second + + # Test that self is properly maintained across retries + obj2 = TestClass() + obj2.instance_value = 100 + result = obj2.instance_method() + assert result == "instance success with value 100" + assert len(obj2.instance_attempts) == 3 diff --git a/dimos/utils/llm_utils.py b/dimos/utils/llm_utils.py new file mode 100644 index 0000000000..05cc44ad24 --- /dev/null +++ b/dimos/utils/llm_utils.py @@ -0,0 +1,75 @@ +# Copyright 2025 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. + +import json +import re +from typing import Union + + +def extract_json(response: str) -> Union[dict, list]: + """Extract JSON from potentially messy LLM response. + + Tries multiple strategies: + 1. Parse the entire response as JSON + 2. Find and parse JSON arrays in the response + 3. Find and parse JSON objects in the response + + Args: + response: Raw text response that may contain JSON + + Returns: + Parsed JSON object (dict or list) + + Raises: + json.JSONDecodeError: If no valid JSON can be extracted + """ + # First try to parse the whole response as JSON + try: + return json.loads(response) + except json.JSONDecodeError: + pass + + # If that fails, try to extract JSON from the messy response + # Look for JSON arrays or objects in the text + + # Pattern to match JSON arrays (including nested arrays/objects) + # This finds the outermost [...] structure + array_pattern = r"\[(?:[^\[\]]*|\[(?:[^\[\]]*|\[[^\[\]]*\])*\])*\]" + + # Pattern to match JSON objects + object_pattern = r"\{(?:[^{}]*|\{(?:[^{}]*|\{[^{}]*\})*\})*\}" + + # Try to find JSON arrays first (most common for detections) + matches = re.findall(array_pattern, response, re.DOTALL) + for match in matches: + try: + parsed = json.loads(match) + # For detection arrays, we expect a list + if isinstance(parsed, list): + return parsed + except json.JSONDecodeError: + continue + + # Try JSON objects if no arrays found + matches = re.findall(object_pattern, response, re.DOTALL) + for match in matches: + try: + return json.loads(match) + except json.JSONDecodeError: + continue + + # If nothing worked, raise an error with the original response + raise json.JSONDecodeError( + f"Could not extract valid JSON from response: {response[:200]}...", response, 0 + ) diff --git a/dimos/utils/test_llm_utils.py b/dimos/utils/test_llm_utils.py new file mode 100644 index 0000000000..4073fd8af2 --- /dev/null +++ b/dimos/utils/test_llm_utils.py @@ -0,0 +1,123 @@ +# Copyright 2025 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. + +"""Tests for LLM utility functions.""" + +import json + +import pytest + +from dimos.utils.llm_utils import extract_json + + +def test_extract_json_clean_response(): + """Test extract_json with clean JSON response.""" + clean_json = '[["object", 1, 2, 3, 4]]' + result = extract_json(clean_json) + assert result == [["object", 1, 2, 3, 4]] + + +def test_extract_json_with_text_before_after(): + """Test extract_json with text before and after JSON.""" + messy = """Here's what I found: + [ + ["person", 10, 20, 30, 40], + ["car", 50, 60, 70, 80] + ] + Hope this helps!""" + result = extract_json(messy) + assert result == [["person", 10, 20, 30, 40], ["car", 50, 60, 70, 80]] + + +def test_extract_json_with_emojis(): + """Test extract_json with emojis and markdown code blocks.""" + messy = """Sure! šŸ˜Š Here are the detections: + + ```json + [["human", 100, 200, 300, 400]] + ``` + + Let me know if you need anything else! šŸ‘""" + result = extract_json(messy) + assert result == [["human", 100, 200, 300, 400]] + + +def test_extract_json_multiple_json_blocks(): + """Test extract_json when there are multiple JSON blocks.""" + messy = """First attempt (wrong format): + {"error": "not what we want"} + + Correct format: + [ + ["cat", 10, 10, 50, 50], + ["dog", 60, 60, 100, 100] + ] + + Another block: {"also": "not needed"}""" + result = extract_json(messy) + # Should return the first valid array + assert result == [["cat", 10, 10, 50, 50], ["dog", 60, 60, 100, 100]] + + +def test_extract_json_object(): + """Test extract_json with JSON object instead of array.""" + response = 'The result is: {"status": "success", "count": 5}' + result = extract_json(response) + assert result == {"status": "success", "count": 5} + + +def test_extract_json_nested_structures(): + """Test extract_json with nested arrays and objects.""" + response = """Processing complete: + [ + ["label1", 1, 2, 3, 4], + {"nested": {"value": 10}}, + ["label2", 5, 6, 7, 8] + ]""" + result = extract_json(response) + assert result[0] == ["label1", 1, 2, 3, 4] + assert result[1] == {"nested": {"value": 10}} + assert result[2] == ["label2", 5, 6, 7, 8] + + +def test_extract_json_invalid(): + """Test extract_json raises error when no valid JSON found.""" + response = "This response has no valid JSON at all!" + with pytest.raises(json.JSONDecodeError) as exc_info: + extract_json(response) + assert "Could not extract valid JSON" in str(exc_info.value) + + +# Test with actual LLM response format +MOCK_LLM_RESPONSE = """ + Yes :) + + [ + ["humans", 76, 368, 219, 580], + ["humans", 354, 372, 512, 525], + ["humans", 409, 370, 615, 748], + ["humans", 628, 350, 762, 528], + ["humans", 785, 323, 960, 650] + ] + + Hope this helps!šŸ˜€šŸ˜Š :)""" + + +def test_extract_json_with_real_llm_response(): + """Test extract_json with actual messy LLM response.""" + result = extract_json(MOCK_LLM_RESPONSE) + assert isinstance(result, list) + assert len(result) == 5 + assert result[0] == ["humans", 76, 368, 219, 580] + assert result[-1] == ["humans", 785, 323, 960, 650] diff --git a/pyproject.toml b/pyproject.toml index 2eab703602..7a71035d27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,8 +80,7 @@ dependencies = [ "transformers[torch]==4.49.0", # Vector Embedding - "sentence_transformers", - + "sentence_transformers", # Perception Dependencies "ultralytics>=8.3.70", @@ -99,7 +98,6 @@ dependencies = [ "googlemaps>=4.10.0", # Inference - "onnx", # Multiprocess @@ -118,7 +116,7 @@ human-cli = "dimos.agents2.cli.human_cli:main" [project.optional-dependencies] manipulation = [ - + # Contact Graspnet Dependencies "h5py>=3.7.0", "pyrender>=0.1.45", @@ -131,15 +129,16 @@ manipulation = [ "tqdm>=4.65.0", "pyyaml>=6.0", "contact-graspnet-pytorch @ git+https://github.com/dimensionalOS/contact_graspnet_pytorch.git", - + # piper arm "piper-sdk", - + # Visualization (Optional) "kaleido>=0.2.1", "plotly>=5.9.0", ] + cpu = [ # CPU inference backends "onnxruntime", @@ -165,6 +164,10 @@ cuda = [ "nltk", "clip @ git+https://github.com/openai/CLIP.git", "detectron2 @ git+https://github.com/facebookresearch/detectron2.git@v0.6", + + # embedding models + "open_clip_torch>=3.0.0", + "torchreid==0.2.5", ] dev = [