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
2026-07-12 20:15:31 +08:00 · 2026-04-03 19:23:11 +08:00
parent 4e16e38f32
commit 5f41cf5794
7 changed files with 111 additions and 44 deletions
--- a/.justfile
+++ b/.justfile
@@ -54,3 +54,6 @@ add-dev +packages:
 remove-dev +packages:
    uv remove {{ packages }} --group dev --no-sync
    just sync-pkgs
+
+memory:
+    MCP_ALLOW_ANONYMOUS_ACCESS=true memory server --http
--- a/mini-nav/compressors/pipeline.py
+++ b/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):
--- a/mini-nav/compressors/proposal/init.py
+++ b/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",
 ]
--- a/mini-nav/compressors/proposal/utils.py
+++ b/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))
--- a/mini-nav/utils/init.py
+++ b/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",
 ]
--- a/mini-nav/utils/image.py
+++ b/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.
--- a/uv.lock
+++ b/uv.lock
@@ -889,7 +889,7 @@ name = "cuda-bindings"
 version = "12.9.4"
 source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/" }
 dependencies = [
-    { name = "cuda-pathfinder", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "cuda-pathfinder", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
 ]
 wheels = [
    { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/7a/d8/b546104b8da3f562c1ff8ab36d130c8fe1dd6a045ced80b4f6ad74f7d4e1/cuda_bindings-12.9.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d3c842c2a4303b2a580fe955018e31aea30278be19795ae05226235268032e5", size = 12148218, upload-time = "2025-10-21T14:51:28.855Z" },
@@ -2397,7 +2397,7 @@ requires-dist = [
    { name = "habitat-lab", specifier = ">=0.3.320250127" },
    { name = "httpx", extras = ["socks"], specifier = ">=0.28.1" },
    { name = "lancedb", specifier = ">=0.30.1" },
-    { name = "marimo", extras = ["mcp"], specifier = ">=0.21.1" },
+    { name = "marimo", extras = ["mcp"], specifier = ">=0.22.0" },
    { name = "matplotlib", specifier = ">=3.10.8" },
    { name = "polars", extras = ["database", "numpy", "pandas", "pydantic"], specifier = ">=1.37.1" },
    { name = "pyarrow", specifier = ">=23.0.0" },
@@ -2968,7 +2968,7 @@ name = "nvidia-cudnn-cu12"
 version = "9.10.2.21"
 source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
 ]
 wheels = [
    { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" },
@@ -2979,7 +2979,7 @@ name = "nvidia-cufft-cu12"
 version = "11.3.3.83"
 source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
 ]
 wheels = [
    { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" },
@@ -3006,9 +3006,9 @@ name = "nvidia-cusolver-cu12"
 version = "11.7.3.90"
 source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
-    { name = "nvidia-cusparse-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
-    { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
+    { name = "nvidia-cusparse-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
+    { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
 ]
 wheels = [
    { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" },
@@ -3019,7 +3019,7 @@ name = "nvidia-cusparse-cu12"
 version = "12.5.8.93"
 source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')" },
 ]
 wheels = [
    { url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" },