detection pipeline rewrite, embedding, vl model standardization, reid system

.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

.envrc.venv

@@ -0,0 +1,2 @@
+source env/bin/activate
+dotenv_if_exists

.gitignore

@@ -48,3 +48,4 @@ yolo11n.pt
 
 # symlink one of .envrc.* if you'd like to use
 .envrc
+.claude

dimos/conftest.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import threading
+
 import pytest
 
 
@@ -24,37 +25,81 @@ 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
     if any(request.node.get_closest_marker(marker) for marker in _skip_for):
         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."
+        )

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

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