From 5f41cf579437488b2fae0ce3f4f03e5c5330719e Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Fri, 3 Apr 2026 19:23:11 +0800 Subject: [PATCH] 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 --- .justfile | 3 ++ mini-nav/compressors/pipeline.py | 49 +++++++++++++++++--- mini-nav/compressors/proposal/__init__.py | 2 - mini-nav/compressors/proposal/utils.py | 25 ---------- mini-nav/utils/__init__.py | 4 +- mini-nav/utils/image.py | 56 ++++++++++++++++++++++- uv.lock | 16 +++---- 7 files changed, 111 insertions(+), 44 deletions(-) delete mode 100644 mini-nav/compressors/proposal/utils.py diff --git a/.justfile b/.justfile index 66acfb1..38ebe83 100644 --- 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 diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index db099de..e1f71d0 100644 --- 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): diff --git a/mini-nav/compressors/proposal/__init__.py b/mini-nav/compressors/proposal/__init__.py index 64a3d13..28318f3 100644 --- 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", ] diff --git a/mini-nav/compressors/proposal/utils.py b/mini-nav/compressors/proposal/utils.py deleted file mode 100644 index 2a53e46..0000000 --- a/mini-nav/compressors/proposal/utils.py +++ /dev/null @@ -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)) diff --git a/mini-nav/utils/__init__.py b/mini-nav/utils/__init__.py index b0ba98e..967ea25 100644 --- 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", ] diff --git a/mini-nav/utils/image.py b/mini-nav/utils/image.py index 9c9ad7e..133db4f 100644 --- 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. diff --git a/uv.lock b/uv.lock index 1d48431..aa4401c 100644 --- 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" },