feat(compressors): add OWLv2 bbox crop to HashPipeline and refactor image utilities

- Add crop_batch method to HashPipeline for cropping images using OWLv2 detection boxes
- Integrate crop_batch into pipeline forward pass (extract_hash and extract_features)
- Move extract_masked_region from compressors/proposal/utils.py to utils/image.py
- Add crop_image_by_bbox utility function in utils/image.py
- Update type annotations to use dict[str, Any] instead of bare dict
- Update .justfile to add memory server command
- Update marimo dependency to >=0.22.0
- Update nvidia CUDA package markers for platform compatibility

mini-nav/compressors/pipeline.py

@@ -1,11 +1,12 @@
 """OWLv2 + SAM + DINO + Hash compression pipeline."""
 
-from typing import Optional, Sequence
+from typing import Any, Optional, Sequence
 
 import torch
 import torch.nn.functional as F
 from PIL import Image
 from utils import get_device
+from utils.image import crop_image_by_bbox, extract_masked_region
 
 from .filter import MaskScoringConfig, select_best_mask
 from .model_loader import (
@@ -17,7 +18,6 @@ from .model_loader import (
 )
 from .proposal import (
     detect_objects_batch,
-    extract_masked_region,
     generate_proposals_batch,
 )
 from .proposal.core import DetectionResult
@@ -51,7 +51,7 @@ class HashPipeline:
 
     Pipeline flow:
         Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) ->
-        DINO (features) -> Hash (binary codes)
+        Crop (OWLv2 box) -> DINO (features) -> Hash (binary codes)
 
     Example:
         pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
@@ -137,7 +137,7 @@ class HashPipeline:
         self,
         images: Sequence[Image.Image],
         bboxes_per_image: list[list[list[float]]],
-    ) -> list[list[dict]]:
+    ) -> list[list[dict[str, Any]]]:
         """Segment objects in images using SAM2 with bounding box prompts.
 
         Args:
@@ -161,7 +161,7 @@ class HashPipeline:
     def filter_batch(
         self,
         images: Sequence[Image.Image],
-        masks_per_image: list[list[dict]],
+        masks_per_image: list[list[dict[str, Any]]],
     ) -> list[Image.Image]:
         """Filter masks and extract best masked region for each image.
 
@@ -177,7 +177,8 @@ class HashPipeline:
             return []
 
         filtered_images: list[Image.Image] = []
-        for image, masks in zip(image_list, masks_per_image):
+        for index, image in enumerate(image_list):
+            masks = masks_per_image[index] if index < len(masks_per_image) else []
             if not masks:
                 filtered_images.append(image)
                 continue
@@ -195,6 +196,40 @@ class HashPipeline:
 
         return filtered_images
 
+    def crop_batch(
+        self,
+        images: Sequence[Image.Image],
+        masks_per_image: list[list[dict[str, Any]]],
+        detections_per_image: list[list[DetectionResult]],
+    ) -> list[Image.Image]:
+        """Crop filtered images using OWLv2 detection boxes.
+
+        Args:
+            images: Sequence of PIL Images after filter_batch.
+            masks_per_image: Masks per image from segment_batch.
+            detections_per_image: Detection results per image from detect_batch.
+
+        Returns:
+            List of cropped PIL Images. Returns original image when no detection exists.
+        """
+        image_list = list(images)
+        if not image_list:
+            return []
+
+        cropped_images: list[Image.Image] = []
+        for index, image in enumerate(image_list):
+            detections = (
+                detections_per_image[index] if index < len(detections_per_image) else []
+            )
+            if detections:
+                best_detection = max(detections, key=lambda d: d["score"])
+                cropped_images.append(crop_image_by_bbox(image, best_detection["bbox"]))
+                continue
+
+            cropped_images.append(image)
+
+        return cropped_images
+
     def extract_dino_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
         """Extract DINO tokens from a batch of images.
 
@@ -259,6 +294,7 @@ class HashPipeline:
         bboxes = [[d["bbox"] for d in dets] for dets in detections]
         masks = self.segment_batch(image_list, bboxes)
         processed = self.filter_batch(image_list, masks)
+        processed = self.crop_batch(processed, masks, detections)
 
         all_bits: list[torch.Tensor] = []
         for i in range(0, len(processed), batch_size):
@@ -298,6 +334,7 @@ class HashPipeline:
         bboxes = [[d["bbox"] for d in dets] for dets in detections]
         masks = self.segment_batch(image_list, bboxes)
         processed = self.filter_batch(image_list, masks)
+        processed = self.crop_batch(processed, masks, detections)
 
         all_features: list[torch.Tensor] = []
         for i in range(0, len(processed), batch_size):

mini-nav/compressors/proposal/__init__.py

@@ -6,12 +6,10 @@ from .core import (
     generate_proposals,
     generate_proposals_batch,
 )
-from .utils import extract_masked_region
 
 __all__ = [
     "detect_objects",
     "detect_objects_batch",
     "generate_proposals",
     "generate_proposals_batch",
-    "extract_masked_region",
 ]

mini-nav/compressors/proposal/utils.py

@@ -1,25 +0,0 @@
-"""Mask extraction utilities."""
-
-import numpy as np
-from PIL import Image
-
-
-def extract_masked_region(
-    image: Image.Image,
-    mask: np.ndarray,
-) -> Image.Image:
-    """Extract masked region from image.
-
-    Args:
-        image: Original PIL Image.
-        mask: Binary mask as numpy array (True = keep).
-
-    Returns:
-        PIL Image with only the masked region visible.
-    """
-    image_np = np.array(image.convert("RGB"))
-
-    # Apply mask.
-    masked_np = image_np * mask[:, :, np.newaxis]
-
-    return Image.fromarray(masked_np.astype(np.uint8))

mini-nav/utils/__init__.py

@@ -4,7 +4,7 @@ from .feature_extractor import (
     extract_single_image_feature,
     infer_vector_dim,
 )
-from .image import numpy_to_pil
+from .image import crop_image_by_bbox, extract_masked_region, numpy_to_pil
 
 __all__ = [
     "get_device",
@@ -13,4 +13,6 @@ __all__ = [
     "extract_single_image_feature",
     "extract_batch_features",
     "numpy_to_pil",
+    "extract_masked_region",
+    "crop_image_by_bbox",
 ]

mini-nav/utils/image.py

@@ -1,12 +1,64 @@
-"""Image conversion utilities."""
+"""Image utilities for conversion, masking, and cropping."""
 
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import numpy as np
+from numpy.typing import NDArray
 from PIL import Image
 
 
-def numpy_to_pil(rgb: np.ndarray) -> Image.Image:
+def extract_masked_region(
+    image: Image.Image,
+    mask: NDArray[np.bool_],
+) -> Image.Image:
+    """Extract masked region from image.
+
+    Args:
+        image: Original PIL Image.
+        mask: Binary mask as numpy array (True = keep).
+
+    Returns:
+        PIL Image with only the masked region visible.
+    """
+    image_np = np.array(image.convert("RGB"))
+    masked_np = image_np * mask[:, :, np.newaxis]
+    return Image.fromarray(masked_np.astype(np.uint8))
+
+
+def crop_image_by_bbox(
+    image: Image.Image,
+    bbox: Sequence[float],
+) -> Image.Image:
+    """Crop an image by bounding box [x1, y1, x2, y2].
+
+    Args:
+        image: Source PIL image.
+        bbox: OWLv2-style box coordinates [x1, y1, x2, y2].
+
+    Returns:
+        Cropped PIL image. Returns the original image when bbox is invalid.
+    """
+    if len(bbox) != 4:
+        return image
+
+    x1, y1, x2, y2 = tuple(float(v) for v in bbox)
+    if not np.isfinite([x1, y1, x2, y2]).all():
+        return image
+
+    left = max(0, int(np.floor(x1)))
+    top = max(0, int(np.floor(y1)))
+    right = min(image.width, int(np.ceil(x2)))
+    bottom = min(image.height, int(np.ceil(y2)))
+
+    if right <= left or bottom <= top:
+        return image
+
+    return image.crop((left, top, right, bottom))
+
+
+def numpy_to_pil(rgb: NDArray[np.uint8]) -> Image.Image:
     """Convert an RGB numpy array to a PIL Image.
 
     Handles arrays with 4 channels (RGBA) by dropping the alpha channel.