mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
5f41cf5794
- 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
347 lines
11 KiB
Python
347 lines
11 KiB
Python
"""OWLv2 + SAM + DINO + Hash compression pipeline."""
|
|
|
|
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 (
|
|
get_dino_dim,
|
|
load_dino_model,
|
|
load_hash_compressor,
|
|
load_owlv2_model,
|
|
load_sam_model,
|
|
)
|
|
from .proposal import (
|
|
detect_objects_batch,
|
|
generate_proposals_batch,
|
|
)
|
|
from .proposal.core import DetectionResult
|
|
|
|
|
|
def create_pipeline_from_config(config) -> "HashPipeline":
|
|
"""Create HashPipeline from a config object.
|
|
|
|
Args:
|
|
config: Configuration object with model settings.
|
|
|
|
Returns:
|
|
Initialized HashPipeline.
|
|
"""
|
|
return HashPipeline(
|
|
owlv2_model=getattr(
|
|
config.model, "owlv2_model", "google/owlv2-base-patch16-ensemble"
|
|
),
|
|
dino_model=config.model.dino_model,
|
|
sam_model=config.model.sam_model,
|
|
hash_bits=config.model.compression_dim,
|
|
compressor_path=config.model.compressor_path,
|
|
mask_scoring_config=getattr(config.model, "mask_scoring_config", None),
|
|
score_threshold=getattr(config.model, "score_threshold", 0.25),
|
|
postprocess_threshold=getattr(config.model, "postprocess_threshold", 0.1),
|
|
)
|
|
|
|
|
|
class HashPipeline:
|
|
"""Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression.
|
|
|
|
Pipeline flow:
|
|
Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) ->
|
|
Crop (OWLv2 box) -> DINO (features) -> Hash (binary codes)
|
|
|
|
Example:
|
|
pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
|
|
images = [Image.open("path/to/image.jpg")]
|
|
text_labels = ["object"]
|
|
hash_bits = pipeline.forward_batch(images, text_labels) # Returns [N, 512]
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
owlv2_model: str = "google/owlv2-base-patch16-ensemble",
|
|
dino_model: str = "facebook/dinov2-large",
|
|
sam_model: str = "facebook/sam2.1-hiera-large",
|
|
hash_bits: int = 512,
|
|
compressor_path: Optional[str] = None,
|
|
mask_scoring_config: Optional["MaskScoringConfig"] = None,
|
|
score_threshold: float = 0.25,
|
|
postprocess_threshold: float = 0.1,
|
|
):
|
|
super().__init__()
|
|
|
|
# Device for model placement.
|
|
self.device = get_device()
|
|
|
|
# OWLv2 detection settings.
|
|
self.owlv2_processor, self.owlv2_model = load_owlv2_model(
|
|
model_name=owlv2_model
|
|
)
|
|
self.score_threshold = score_threshold
|
|
self.postprocess_threshold = postprocess_threshold
|
|
|
|
# Mask scoring config for filter step.
|
|
self.mask_scoring_config = mask_scoring_config
|
|
|
|
# SAM2 model for segmentation.
|
|
self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model)
|
|
|
|
# DINO model for feature extraction.
|
|
self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
|
|
self.dino_dim = get_dino_dim(dino_model)
|
|
|
|
# Hash compressor for binarizing DINO features.
|
|
self.hash_compressor = load_hash_compressor(
|
|
input_dim=self.dino_dim,
|
|
hash_bits=hash_bits,
|
|
compressor_path=compressor_path,
|
|
)
|
|
|
|
@property
|
|
def hash_bits(self) -> int:
|
|
"""Number of bits in the hash code."""
|
|
return self.hash_compressor.hash_bits
|
|
|
|
def detect_batch(
|
|
self,
|
|
images: Sequence[Image.Image],
|
|
text_labels: list[str],
|
|
) -> list[list[DetectionResult]]:
|
|
"""Detect objects in a batch of images using OWLv2.
|
|
|
|
Args:
|
|
images: Sequence of PIL Images.
|
|
text_labels: Text labels for all images (same labels used for each image).
|
|
|
|
Returns:
|
|
List of lists of DetectionResult dicts, one inner list per image.
|
|
"""
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return []
|
|
|
|
text_labels_per_image = [text_labels] * len(image_list)
|
|
return detect_objects_batch(
|
|
self.owlv2_model,
|
|
self.owlv2_processor,
|
|
image_list,
|
|
text_labels_per_image,
|
|
score_threshold=self.score_threshold,
|
|
postprocess_threshold=self.postprocess_threshold,
|
|
)
|
|
|
|
def segment_batch(
|
|
self,
|
|
images: Sequence[Image.Image],
|
|
bboxes_per_image: list[list[list[float]]],
|
|
) -> list[list[dict[str, Any]]]:
|
|
"""Segment objects in images using SAM2 with bounding box prompts.
|
|
|
|
Args:
|
|
images: Sequence of PIL Images.
|
|
bboxes_per_image: Bounding boxes per image as [[[x1,y1,x2,y2], ...], ...].
|
|
|
|
Returns:
|
|
List of lists of mask dictionaries, one inner list per image.
|
|
"""
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return []
|
|
|
|
return generate_proposals_batch(
|
|
self.sam_model,
|
|
self.sam_processor,
|
|
image_list,
|
|
bboxes_per_image,
|
|
)
|
|
|
|
def filter_batch(
|
|
self,
|
|
images: Sequence[Image.Image],
|
|
masks_per_image: list[list[dict[str, Any]]],
|
|
) -> list[Image.Image]:
|
|
"""Filter masks and extract best masked region for each image.
|
|
|
|
Args:
|
|
images: Sequence of PIL Images.
|
|
masks_per_image: Masks per image from segment_batch.
|
|
|
|
Returns:
|
|
List of PIL Images, one per input image (original if no valid masks).
|
|
"""
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return []
|
|
|
|
filtered_images: list[Image.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
|
|
|
|
best_mask = select_best_mask(
|
|
masks,
|
|
image_shape=(image.height, image.width),
|
|
config=self.mask_scoring_config,
|
|
)
|
|
if best_mask is None:
|
|
filtered_images.append(image)
|
|
continue
|
|
|
|
filtered_images.append(extract_masked_region(image, best_mask["segment"]))
|
|
|
|
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.
|
|
|
|
Args:
|
|
images: Sequence of PIL Images.
|
|
|
|
Returns:
|
|
Last hidden state tokens of shape [B, N, dim].
|
|
"""
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return torch.empty(
|
|
(0, 1, self.dino_dim), dtype=torch.float32, device=self.device
|
|
)
|
|
|
|
inputs = self.dino_processor(images=image_list, return_tensors="pt").to(
|
|
self.device
|
|
)
|
|
|
|
with torch.no_grad():
|
|
outputs = self.dino(**inputs)
|
|
return outputs.last_hidden_state
|
|
|
|
def compress_batch(self, tokens: torch.Tensor) -> torch.Tensor:
|
|
"""Compress DINO tokens to binary hash codes.
|
|
|
|
Args:
|
|
tokens: DINO tokens of shape [B, N, dim].
|
|
|
|
Returns:
|
|
Binary hash codes of shape [B, hash_bits] as int32.
|
|
"""
|
|
_, _, bits = self.hash_compressor(tokens)
|
|
return bits
|
|
|
|
def forward_batch(
|
|
self,
|
|
images: Sequence[Image.Image],
|
|
text_labels: list[str],
|
|
batch_size: int = 32,
|
|
) -> torch.Tensor:
|
|
"""Process a batch of images through the full pipeline.
|
|
|
|
Args:
|
|
images: Sequence of PIL Images.
|
|
text_labels: Text labels for detection (same for all images).
|
|
batch_size: Batch size for DINO feature extraction.
|
|
|
|
Returns:
|
|
Binary hash codes of shape [N, hash_bits] as int32.
|
|
"""
|
|
if batch_size <= 0:
|
|
raise ValueError("batch_size must be greater than 0")
|
|
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return torch.empty(
|
|
(0, self.hash_bits), dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
detections = self.detect_batch(image_list, text_labels)
|
|
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):
|
|
sub_batch = processed[i : i + batch_size]
|
|
tokens = self.extract_dino_batch(sub_batch)
|
|
bits = self.compress_batch(tokens)
|
|
all_bits.append(bits)
|
|
|
|
return torch.cat(all_bits, dim=0)
|
|
|
|
def extract_features_dataset(
|
|
self,
|
|
images: Sequence[Image.Image],
|
|
text_labels: list[str],
|
|
batch_size: int = 32,
|
|
) -> torch.Tensor:
|
|
"""Extract normalized DINO features from a batch of images.
|
|
|
|
Args:
|
|
images: Sequence of PIL Images.
|
|
text_labels: Text labels for detection (same for all images).
|
|
batch_size: Batch size for DINO feature extraction.
|
|
|
|
Returns:
|
|
Normalized DINO features of shape [N, dino_dim].
|
|
"""
|
|
if batch_size <= 0:
|
|
raise ValueError("batch_size must be greater than 0")
|
|
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return torch.empty(
|
|
(0, self.dino_dim), dtype=torch.float32, device=self.device
|
|
)
|
|
|
|
detections = self.detect_batch(image_list, text_labels)
|
|
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):
|
|
sub_batch = processed[i : i + batch_size]
|
|
tokens = self.extract_dino_batch(sub_batch)
|
|
features = tokens.mean(dim=1)
|
|
all_features.append(F.normalize(features, dim=-1))
|
|
|
|
return torch.cat(all_features, dim=0)
|