SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-13 04:25:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(compressors): migrate pipeline to OWLv2-based detection with text labels

Browse Source 
- Replace bbox-prompted segmentation with OWLv2 text-guided object detection
- Refactor HashPipeline from nn.Module to plain class with modular stage methods
- Add detect_batch, segment_batch, filter_batch for explicit pipeline stages
- Rename forward to forward_batch with text_labels API instead of bboxes
- Add mask_scoring_config, score_threshold, postprocess_threshold configuration
- Update model_loader to expose Dinov2Model type annotation
This commit is contained in:
SikongJueluo
2026-04-03 15:49:23 +08:00
parent 4918b654e7
commit 4e16e38f32
3 changed files with 166 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

302
mini-nav/compressors/pipeline.py
View File

@@ -1,25 +1,26 @@
"""SAM + DINO + Hash compression pipeline."""
"""OWLv2 + SAM + DINO + Hash compression pipeline."""
from typing import Optional, Sequence
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from utils import get_device
from .filter import select_best_mask
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,
extract_masked_region,
generate_proposals,
generate_proposals_batch,
)
from .proposal.core import DetectionResult
def create_pipeline_from_config(config) -> "HashPipeline":
@@ -32,50 +33,64 @@ def create_pipeline_from_config(config) -> "HashPipeline":
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,
sam_min_mask_area=config.model.sam_min_mask_area,
sam_max_masks=config.model.sam_max_masks,
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(nn.Module):
"""Pipeline for SAM segmentation + DINO features + Hash compression.
class HashPipeline:
"""Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression.
Pipeline flow:
PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)
Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) ->
DINO (features) -> Hash (binary codes)
Example:
pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
image = Image.open("path/to/image.jpg")
hash_bits = pipeline(image) # Returns [1, 512] binary bits
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",
sam_min_mask_area: int = 100,
sam_max_masks: int = 10,
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()
# SAM2 filter settings.
self.sam_min_mask_area = sam_min_mask_area
self.sam_max_masks = sam_max_masks
# 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
# Load models.
# 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)
self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
# DINO feature dimension based on model size.
# 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.
@@ -90,82 +105,112 @@ class HashPipeline(nn.Module):
"""Number of bits in the hash code."""
return self.hash_compressor.hash_bits
def _segment_with_sam(
self, image: Image.Image, bboxes: list[list[float]]
) -> Image.Image:
"""Segment image with SAM and extract the best object mask.
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:
image: Input PIL Image.
bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
images: Sequence of PIL Images.
text_labels: Text labels for all images (same labels used for each image).
Returns:
Masked image containing only the best object, or original if no masks.
List of lists of DetectionResult dicts, one inner list per image.
"""
masks = generate_proposals(
self.sam_model,
self.sam_processor,
image,
bboxes,
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,
)
masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
if not masks:
return image
best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
if best_mask is None:
return image
return extract_masked_region(image, best_mask["segment"])
def _segment_with_sam_dataset(
def segment_batch(
self,
images: Sequence[Image.Image],
bboxes_per_image: list[list[list[float]]],
) -> list[Image.Image]:
) -> list[list[dict]]:
"""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)
masks_dataset = generate_proposals_batch(
if not image_list:
return []
return generate_proposals_batch(
self.sam_model,
self.sam_processor,
image_list,
bboxes_per_image,
)
masks_dataset = [
_filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
for masks in masks_dataset
]
selected_images: list[Image.Image] = []
for image, masks in zip(image_list, masks_dataset):
if not masks:
selected_images.append(image)
continue
best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
if best_mask is None:
selected_images.append(image)
continue
selected_images.append(extract_masked_region(image, best_mask["segment"]))
return selected_images
def _dino_forward(self, image: Image.Image) -> torch.Tensor:
"""Extract DINO tokens from an image.
def filter_batch(
self,
images: Sequence[Image.Image],
masks_per_image: list[list[dict]],
) -> list[Image.Image]:
"""Filter masks and extract best masked region for each image.
Args:
image: Input PIL Image.
images: Sequence of PIL Images.
masks_per_image: Masks per image from segment_batch.
Returns:
Last hidden state tokens of shape [1, N, dim].
List of PIL Images, one per input image (original if no valid masks).
"""
inputs = self.dino_processor(image, return_tensors="pt").to(self.device)
image_list = list(images)
if not image_list:
return []
with torch.no_grad():
outputs = self.dino(**inputs)
return outputs.last_hidden_state
filtered_images: list[Image.Image] = []
for image, masks in zip(image_list, masks_per_image):
if not masks:
filtered_images.append(image)
continue
def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
inputs = self.dino_processor(images=list(images), return_tensors="pt").to(
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 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
)
@@ -173,113 +218,92 @@ class HashPipeline(nn.Module):
outputs = self.dino(**inputs)
return outputs.last_hidden_state
def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor:
"""Process a single image through the full pipeline.
def compress_batch(self, tokens: torch.Tensor) -> torch.Tensor:
"""Compress DINO tokens to binary hash codes.
Args:
image: Input PIL Image.
bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
tokens: DINO tokens of shape [B, N, dim].
Returns:
Binary hash codes of shape [1, hash_bits] as int32.
Binary hash codes of shape [B, hash_bits] as int32.
"""
image = self._segment_with_sam(image, bboxes)
tokens = self._dino_forward(image)
_, _, bits = self.hash_compressor(tokens)
return bits
def forward_dataset(
def forward_batch(
self,
images: Sequence[Image.Image],
bboxes_per_image: list[list[list[float]]],
text_labels: list[str],
batch_size: int = 32,
apply_sam: bool = True,
) -> 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 len(image_list) == 0:
if not image_list:
return torch.empty(
(0, self.hash_bits), dtype=torch.int32, device=self.device
)
if apply_sam:
processed_images = self._segment_with_sam_dataset(
image_list, bboxes_per_image
)
else:
processed_images = image_list
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)
batch_bits: list[torch.Tensor] = []
for i in range(0, len(processed_images), batch_size):
batch_images = processed_images[i : i + batch_size]
tokens = self._dino_forward_batch(batch_images)
_, _, bits = self.hash_compressor(tokens)
batch_bits.append(bits)
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(batch_bits, dim=0)
def extract_features(
self, image: Image.Image, bboxes: list[list[float]]
) -> torch.Tensor:
"""Extract normalized DINO features from an image.
Args:
image: Input PIL Image.
bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
Returns:
Normalized DINO features of shape [1, dino_dim].
"""
image = self._segment_with_sam(image, bboxes)
tokens = self._dino_forward(image)
features = tokens.mean(dim=1)
return F.normalize(features, dim=-1)
return torch.cat(all_bits, dim=0)
def extract_features_dataset(
self,
images: Sequence[Image.Image],
bboxes_per_image: list[list[list[float]]],
text_labels: list[str],
batch_size: int = 32,
apply_sam: bool = True,
) -> 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 len(image_list) == 0:
if not image_list:
return torch.empty(
(0, self.dino_dim), dtype=torch.float32, device=self.device
)
if apply_sam:
processed_images = self._segment_with_sam_dataset(
image_list, bboxes_per_image
)
else:
processed_images = image_list
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)
all_features: list[torch.Tensor] = []
for i in range(0, len(processed_images), batch_size):
batch_images = processed_images[i : i + batch_size]
tokens = self._dino_forward_batch(batch_images)
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)
def _filter_masks(
masks: list[dict],
min_area: int,
max_masks: int,
) -> list[dict]:
"""Filter masks by area and keep top-N largest."""
filtered = [m for m in masks if int(m["area"]) >= min_area]
if not filtered:
return []
sorted_masks = sorted(filtered, key=lambda m: m["area"], reverse=True)
return sorted_masks[:max_masks]