danvi/dim1486/generalize dan perception stack

dimos/memory/backend.py

@@ -207,13 +207,16 @@ def _vector_search(self, query: StreamQuery) -> Iterator[Observation[T]]:
         ranked: list[Observation[T]] = []
         for obs_id, sim in hits:
             match = obs_by_id.get(obs_id)
-            if match is not None:
-                ranked.append(
-                    match.derive(data=match.data, embedding=query.search_vec, similarity=sim)
-                )
+            if match is None:
+                continue
+            if not all(f.matches(match) for f in query.filters):
+                continue
+            ranked.append(
+                match.derive(data=match.data, embedding=query.search_vec, similarity=sim)
+            )
 
-        # Apply remaining query ops (skip vector search)
-        rest = replace(query, search_vec=None, search_k=None)
+        # Apply remaining query ops (filters already applied; skip vector search)
+        rest = replace(query, search_vec=None, search_k=None, filters=())
         yield from rest.apply(iter(ranked))
 
     def _iterate_live(

dimos/memory/stream.py

@@ -568,8 +568,15 @@ def drain(self) -> int:
         return n
 
     def drain_thread(self) -> DisposableBase:
-        """Drain this stream on the dimos thread pool; returns a disposable."""
-        return self.subscribe(lambda _: None)
+        """Drain this stream on the dimos thread pool; returns a disposable.
+
+        A drain has no consumer to surface an iteration error, so it logs -
+        a dead pipeline must not die silently.
+        """
+        return self.subscribe(
+            lambda _: None,
+            on_error=lambda e: logger.error("drain_thread() pipeline died: %s", e, exc_info=e),
+        )
 
     def observable(self) -> reactivex.Observable[O]:
         """Convert this stream to an RxPY Observable.

dimos/memory/utils/sqlite.py

@@ -27,6 +27,9 @@ def open_sqlite_connection(path: str | Path) -> sqlite3.Connection:
     conn = sqlite3.connect(path, check_same_thread=False)
     conn.execute("PRAGMA journal_mode=WAL")
     conn.execute("PRAGMA synchronous=NORMAL")
+    # WAL has one writer at a time; concurrent writer threads (a recorder
+    # feed plus a live embed save) can starve past the 5s default and die.
+    conn.execute("PRAGMA busy_timeout=30000")
     conn.enable_load_extension(True)
     sqlite_vec.load(conn)
     conn.enable_load_extension(False)

dimos/perception/detection/detectors/owlv2.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+from collections import OrderedDict
 from functools import cached_property
 
 import numpy as np
@@ -27,9 +28,14 @@
 from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox
 from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D
 
+# ~3.8 MB of GPU memory per cached frame
+_FEATURE_CACHE_MAX = 128
+
 
 class Owlv2Config(HuggingFaceModelConfig):
     model_name: str = "google/owlv2-base-patch16-ensemble"
+    # float16 runs the forward under autocast at roughly half the latency;
+    # scores jitter by a few thousandths, so threshold-edge boxes may flip.
     dtype: torch.dtype = torch.float32
 
 
@@ -45,6 +51,22 @@ class Owlv2Detector(HuggingFaceModel):
 
     config: Owlv2Config
 
+    def __init__(self, **kwargs: object) -> None:
+        super().__init__(**kwargs)
+        # image-tower forwards run so far; instrumentation reads deltas
+        self.forwards = 0
+        # (frame ts, label, floor) to that label's floor-filtered
+        # (boxes_kx4, scores_k); filled and read by localize
+        self.score_cache: OrderedDict[tuple[float, str, float], tuple[np.ndarray, np.ndarray]] = (
+            OrderedDict()
+        )
+        # frame ts to text-independent tower outputs; any query set scores
+        # against a cached frame without rerunning the tower
+        self._features: OrderedDict[
+            float, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+        ] = OrderedDict()
+        self._text_embeds: dict[str, torch.Tensor] = {}
+
     @cached_property
     def _model(self):  # type: ignore[no-untyped-def]
         from transformers import Owlv2ForObjectDetection
@@ -62,6 +84,13 @@ def _processor(self):  # type: ignore[no-untyped-def]
 
         return Owlv2Processor.from_pretrained(self.config.model_name)
 
+    def _autocast(self) -> torch.autocast:
+        return torch.autocast(
+            device_type="cuda",
+            dtype=self.config.dtype,
+            enabled=self.config.dtype is not torch.float32 and "cuda" in str(self.config.device),
+        )
+
     def query_detections(
         self,
         image: Image,
@@ -74,38 +103,55 @@ def query_detections(
         ``confidence`` is the calibrated per-box score. ``class_id`` indexes
         into ``queries``.
         """
-        pil = PILImage.fromarray(image.to_rgb().data)
-        with torch.inference_mode():
-            inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to(
-                self.config.device
-            )
+        return self.query_detections_batch([image], queries, threshold)[0]
+
+    def query_detections_batch(
+        self,
+        images: list[Image],
+        queries: list[str],
+        threshold: float = 0.1,
+    ) -> list[ImageDetections2D]:
+        """``query_detections`` over several images in one forward pass.
+
+        Per-call preprocessing, text encoding and kernel launches amortize
+        across the batch, which is what makes many-frame sweeps affordable;
+        results are per-image, in input order.
+        """
+        self.forwards += len(images)
+        pils = [PILImage.fromarray(image.to_rgb().data) for image in images]
+        with torch.inference_mode(), self._autocast():
+            inputs = self._processor(
+                text=[queries] * len(pils), images=pils, return_tensors="pt"
+            ).to(self.config.device)
             outputs = self._model(**inputs)
             results = self._processor.post_process_grounded_object_detection(
                 outputs=outputs,
-                target_sizes=torch.tensor([(pil.height, pil.width)]),
+                target_sizes=torch.tensor([(pil.height, pil.width) for pil in pils]),
                 threshold=threshold,
-            )[0]
-
-        detections: list[Detection2DBBox] = []
-        w, h = float(pil.width), float(pil.height)
-        for box, score, label in zip(
-            results["boxes"], results["scores"], results["labels"], strict=False
-        ):
-            x1, y1, x2, y2 = (float(v) for v in box)
-            bbox = (max(0.0, x1), max(0.0, y1), min(w, x2), min(h, y2))
-            det = Detection2DBBox(
-                bbox=bbox,
-                track_id=-1,
-                class_id=int(label),
-                confidence=float(score),
-                name=queries[int(label)],
-                ts=image.ts,
-                image=image,
             )
-            if det.is_valid():
-                detections.append(det)
 
-        return ImageDetections2D(image=image, detections=detections)
+        batch: list[ImageDetections2D] = []
+        for image, pil, result in zip(images, pils, results, strict=True):
+            detections: list[Detection2DBBox] = []
+            w, h = float(pil.width), float(pil.height)
+            for box, score, label in zip(
+                result["boxes"], result["scores"], result["labels"], strict=False
+            ):
+                x1, y1, x2, y2 = (float(v) for v in box)
+                bbox = (max(0.0, x1), max(0.0, y1), min(w, x2), min(h, y2))
+                det = Detection2DBBox(
+                    bbox=bbox,
+                    track_id=-1,
+                    class_id=int(label),
+                    confidence=float(score),
+                    name=queries[int(label)],
+                    ts=image.ts,
+                    image=image,
+                )
+                if det.is_valid():
+                    detections.append(det)
+            batch.append(ImageDetections2D(image=image, detections=detections))
+        return batch
 
     def query_score_rows(
         self,
@@ -122,8 +168,9 @@ def query_score_rows(
         queries and refuse. Returns pixel ``(x1, y1, x2, y2)`` boxes and
         their score rows.
         """
+        self.forwards += 1
         pil = PILImage.fromarray(image.to_rgb().data)
-        with torch.inference_mode():
+        with torch.inference_mode(), self._autocast():
             inputs = self._processor(text=[queries], images=pil, return_tensors="pt").to(
                 self.config.device
             )
@@ -143,7 +190,98 @@ def query_score_rows(
         boxes[:, 1::2] = boxes[:, 1::2].clip(0.0, float(pil.height))
         return boxes, kept
 
+    def _frame_features(
+        self, image: Image
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Text-independent tower outputs for one frame, cached by frame ts.
+
+        Returns unit-normalized per-box class embeddings, the class head's
+        logit shift and scale, and clipped pixel ``(x1, y1, x2, y2)`` boxes.
+        A miss runs the image tower once; a hit runs nothing.
+        """
+        key = image.ts
+        if key in self._features:
+            self._features.move_to_end(key)
+            return self._features[key]
+
+        from transformers.image_transforms import center_to_corners_format
+
+        self.forwards += 1
+        pil = PILImage.fromarray(image.to_rgb().data)
+        model = self._model
+        head = model.class_head
+        with torch.inference_mode(), self._autocast():
+            pixel_values = self._processor(images=pil, return_tensors="pt").pixel_values.to(
+                self.config.device
+            )
+            feature_map = model.image_embedder(pixel_values)[0]
+            batch, height, width, dim = feature_map.shape
+            image_feats = feature_map.reshape(batch, height * width, dim)
+            pred_boxes = model.box_predictor(image_feats, feature_map)
+            class_embeds = head.dense0(image_feats)
+            class_embeds = class_embeds / (
+                torch.linalg.norm(class_embeds, dim=-1, keepdim=True) + 1e-6
+            )
+            logit_shift = head.logit_shift(image_feats)
+            logit_scale = head.elu(head.logit_scale(image_feats)) + 1
+
+            # the same conversion post_process_grounded_object_detection runs:
+            # corners in model dtype, scaled to the padded square in float32
+            boxes = center_to_corners_format(pred_boxes)[0].float() * float(
+                max(pil.width, pil.height)
+            )
+            boxes[:, 0::2] = boxes[:, 0::2].clip(0.0, float(pil.width))
+            boxes[:, 1::2] = boxes[:, 1::2].clip(0.0, float(pil.height))
+
+        entry = (class_embeds[0], logit_shift[0], logit_scale[0], boxes)
+        self._features[key] = entry
+        if len(self._features) > _FEATURE_CACHE_MAX:
+            self._features.popitem(last=False)
+        return entry
+
+    def _query_embeds(self, queries: list[str]) -> torch.Tensor:
+        """Unit-normalized text embeddings, one row per query, cached per string."""
+        missing = [q for q in queries if q not in self._text_embeds]
+        if missing:
+            with torch.inference_mode(), self._autocast():
+                inputs = self._processor(text=[missing], return_tensors="pt").to(self.config.device)
+                embeds = self._model.owlv2.get_text_features(**inputs)
+                embeds = embeds / torch.linalg.norm(embeds, ord=2, dim=-1, keepdim=True)
+            for query, embed in zip(missing, embeds, strict=True):
+                self._text_embeds[query] = embed
+        return torch.stack([self._text_embeds[q] for q in queries])
+
+    def query_score_rows_batch(
+        self,
+        images: list[Image],
+        queries: list[str],
+        threshold: float = 0.1,
+    ) -> list[tuple[np.ndarray, np.ndarray]]:
+        """``query_score_rows`` over several images: one ``(boxes, rows)`` each.
+
+        The image tower runs only for frames missing from the feature cache;
+        scoring any query set against a cached frame is a matmul against its
+        stored class embeddings, so repeated windows and new labels on seen
+        frames cost no model forwards.
+        """
+        with torch.inference_mode():
+            query_embeds = self._query_embeds(queries)
+            query_embeds = query_embeds / (
+                torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6
+            )
+            out: list[tuple[np.ndarray, np.ndarray]] = []
+            for image in images:
+                class_embeds, logit_shift, logit_scale, boxes_px = self._frame_features(image)
+                logits = (class_embeds @ query_embeds.T + logit_shift) * logit_scale
+                scores = torch.sigmoid(logits.to(torch.float32))
+                keep = scores.max(dim=-1).values > threshold
+                out.append((boxes_px[keep].cpu().numpy(), scores[keep].cpu().numpy()))
+        return out
+
     def stop(self) -> None:
+        self.score_cache.clear()
+        self._features.clear()
+        self._text_embeds.clear()
         if "_processor" in self.__dict__:
             del self.__dict__["_processor"]
         super().stop()