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
This commit is contained in:
2026-04-03 19:23:11 +08:00
parent 4e16e38f32
commit 5f41cf5794
7 changed files with 111 additions and 44 deletions

View File

@@ -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):

View File

@@ -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",
]

View File

@@ -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))