mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
feat(compressors): refactor pipeline with FramePacket dataclass and unified process_batch
- Add FramePacket dataclass to encapsulate per-image pipeline state - Rename internal methods with underscore prefix convention - Replace separate filter_batch/crop_batch with unified process_batch method - Update notebook to use new HashPipeline API
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
"""OWLv2 + SAM + DINO + Hash compression pipeline."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
import torch
|
||||
@@ -8,7 +9,12 @@ 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 .filter import (
|
||||
MaskScoringConfig,
|
||||
compute_mask_features,
|
||||
score_mask,
|
||||
should_reject_mask,
|
||||
)
|
||||
from .model_loader import (
|
||||
get_dino_dim,
|
||||
load_dino_model,
|
||||
@@ -24,14 +30,6 @@ 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"
|
||||
@@ -46,18 +44,40 @@ def create_pipeline_from_config(config) -> "HashPipeline":
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FramePacket:
|
||||
image: Image.Image
|
||||
boxes_xyxy: list[list[float]] = field(default_factory=list)
|
||||
scores: list[float] = field(default_factory=list)
|
||||
labels: list[str] = field(default_factory=list)
|
||||
masks: list[dict[str, Any]] = field(default_factory=list)
|
||||
selected_idx: int | None = None
|
||||
dropped_indices: list[int] = field(default_factory=list)
|
||||
fallback_reason: str | None = None
|
||||
filtered_image: Image.Image | None = None
|
||||
cropped_image: Image.Image | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineBatchOutput:
|
||||
hash_bits: torch.Tensor
|
||||
cropped_images: list[Image.Image]
|
||||
debug_meta: list[dict[str, Any]]
|
||||
|
||||
|
||||
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)
|
||||
Crop (matched 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]
|
||||
output = pipeline.process_batch(images, text_labels)
|
||||
hash_bits = output.hash_bits # [N, 512]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -73,27 +93,16 @@ class HashPipeline:
|
||||
):
|
||||
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,
|
||||
@@ -102,23 +111,13 @@ class HashPipeline:
|
||||
|
||||
@property
|
||||
def hash_bits(self) -> int:
|
||||
"""Number of bits in the hash code."""
|
||||
return self.hash_compressor.hash_bits
|
||||
|
||||
def detect_batch(
|
||||
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 []
|
||||
@@ -133,20 +132,11 @@ class HashPipeline:
|
||||
postprocess_threshold=self.postprocess_threshold,
|
||||
)
|
||||
|
||||
def segment_batch(
|
||||
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 []
|
||||
@@ -158,77 +148,169 @@ class HashPipeline:
|
||||
bboxes_per_image,
|
||||
)
|
||||
|
||||
def filter_batch(
|
||||
def _build_frame_packets(self, images: Sequence[Image.Image]) -> list[FramePacket]:
|
||||
return [FramePacket(image=image) for image in images]
|
||||
|
||||
def _attach_detections(
|
||||
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:
|
||||
packets: list[FramePacket],
|
||||
text_labels: list[str],
|
||||
) -> list[list[DetectionResult]]:
|
||||
if not packets:
|
||||
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):
|
||||
image_list = [packet.image for packet in packets]
|
||||
detections_per_image = self._detect_batch(image_list, text_labels)
|
||||
for index, packet in enumerate(packets):
|
||||
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"]))
|
||||
packet.boxes_xyxy = [list(det["bbox"]) for det in detections]
|
||||
packet.scores = [float(det["score"]) for det in detections]
|
||||
packet.labels = [str(det["label"]) for det in detections]
|
||||
if not packet.boxes_xyxy:
|
||||
packet.fallback_reason = "no_detection"
|
||||
return detections_per_image
|
||||
|
||||
def _attach_masks(self, packets: list[FramePacket]) -> None:
|
||||
if not packets:
|
||||
return
|
||||
|
||||
image_list = [packet.image for packet in packets]
|
||||
boxes_per_image = [packet.boxes_xyxy for packet in packets]
|
||||
masks_per_image = self._segment_batch(image_list, boxes_per_image)
|
||||
|
||||
for index, packet in enumerate(packets):
|
||||
packet.masks = (
|
||||
masks_per_image[index] if index < len(masks_per_image) else []
|
||||
)
|
||||
if (
|
||||
packet.boxes_xyxy
|
||||
and not packet.masks
|
||||
and packet.fallback_reason is None
|
||||
):
|
||||
packet.fallback_reason = "no_mask"
|
||||
|
||||
def _select_candidates(self, packets: list[FramePacket]) -> None:
|
||||
config = self.mask_scoring_config or MaskScoringConfig()
|
||||
|
||||
for packet in packets:
|
||||
if not packet.masks:
|
||||
packet.selected_idx = None
|
||||
packet.dropped_indices = []
|
||||
if packet.fallback_reason is None:
|
||||
packet.fallback_reason = "no_mask"
|
||||
continue
|
||||
|
||||
cropped_images.append(image)
|
||||
kept: list[tuple[float, int, int]] = []
|
||||
dropped: list[int] = []
|
||||
for index, mask in enumerate(packet.masks):
|
||||
features = compute_mask_features(
|
||||
mask, image_shape=(packet.image.height, packet.image.width)
|
||||
)
|
||||
if should_reject_mask(features, config):
|
||||
dropped.append(index)
|
||||
continue
|
||||
|
||||
return cropped_images
|
||||
mask_score = score_mask(
|
||||
mask,
|
||||
image_shape=(packet.image.height, packet.image.width),
|
||||
config=config,
|
||||
)
|
||||
area = int(mask.get("area", 0))
|
||||
kept.append((float(mask_score), area, index))
|
||||
|
||||
if kept:
|
||||
kept.sort(reverse=True)
|
||||
packet.selected_idx = kept[0][2]
|
||||
packet.dropped_indices = dropped
|
||||
continue
|
||||
|
||||
fallback_index = max(
|
||||
range(len(packet.masks)),
|
||||
key=lambda idx: int(packet.masks[idx].get("area", 0)),
|
||||
)
|
||||
packet.selected_idx = fallback_index
|
||||
packet.dropped_indices = [
|
||||
index for index in range(len(packet.masks)) if index != fallback_index
|
||||
]
|
||||
if packet.fallback_reason is None:
|
||||
packet.fallback_reason = "all_masks_rejected_fallback_area"
|
||||
|
||||
def _render_filtered_images(self, packets: list[FramePacket]) -> None:
|
||||
for packet in packets:
|
||||
if packet.selected_idx is None:
|
||||
packet.filtered_image = packet.image
|
||||
continue
|
||||
|
||||
if packet.selected_idx >= len(packet.masks):
|
||||
packet.filtered_image = packet.image
|
||||
packet.fallback_reason = (
|
||||
packet.fallback_reason or "selected_index_out_of_mask_range"
|
||||
)
|
||||
packet.selected_idx = None
|
||||
continue
|
||||
|
||||
selected_mask = packet.masks[packet.selected_idx]
|
||||
packet.filtered_image = extract_masked_region(
|
||||
packet.image, selected_mask["segment"]
|
||||
)
|
||||
|
||||
def _render_cropped_images(
|
||||
self,
|
||||
packets: list[FramePacket],
|
||||
detections_per_image: list[list[DetectionResult]],
|
||||
) -> None:
|
||||
for index, packet in enumerate(packets):
|
||||
base_image = (
|
||||
packet.filtered_image if packet.filtered_image else packet.image
|
||||
)
|
||||
detections = (
|
||||
detections_per_image[index] if index < len(detections_per_image) else []
|
||||
)
|
||||
|
||||
if packet.selected_idx is None:
|
||||
packet.cropped_image = base_image
|
||||
if packet.fallback_reason is None:
|
||||
packet.fallback_reason = "no_selected_candidate"
|
||||
continue
|
||||
|
||||
if packet.selected_idx >= len(detections):
|
||||
packet.cropped_image = base_image
|
||||
packet.fallback_reason = (
|
||||
packet.fallback_reason or "selected_index_out_of_detection_range"
|
||||
)
|
||||
continue
|
||||
|
||||
selected_detection = detections[packet.selected_idx]
|
||||
cropped = crop_image_by_bbox(base_image, selected_detection["bbox"])
|
||||
packet.cropped_image = cropped
|
||||
if cropped.size == base_image.size:
|
||||
packet.fallback_reason = (
|
||||
packet.fallback_reason or "invalid_or_full_bbox"
|
||||
)
|
||||
|
||||
def _build_debug_meta(
|
||||
self,
|
||||
packets: list[FramePacket],
|
||||
return_debug_details: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
debug_meta: list[dict[str, Any]] = []
|
||||
for packet in packets:
|
||||
item: dict[str, Any] = {
|
||||
"selected_idx": packet.selected_idx,
|
||||
"dropped_indices": list(packet.dropped_indices),
|
||||
"fallback_reason": packet.fallback_reason,
|
||||
"num_boxes": len(packet.boxes_xyxy),
|
||||
"num_masks": len(packet.masks),
|
||||
}
|
||||
if return_debug_details:
|
||||
item["boxes_xyxy"] = [list(box) for box in packet.boxes_xyxy]
|
||||
item["scores"] = [float(score) for score in packet.scores]
|
||||
item["labels"] = [str(label) for label in packet.labels]
|
||||
item["masks"] = packet.masks
|
||||
debug_meta.append(item)
|
||||
return debug_meta
|
||||
|
||||
def extract_dino_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
|
||||
"""Extract DINO tokens from a batch of images.
|
||||
@@ -265,6 +347,65 @@ class HashPipeline:
|
||||
_, _, bits = self.hash_compressor(tokens)
|
||||
return bits
|
||||
|
||||
def process_batch(
|
||||
self,
|
||||
images: Sequence[Image.Image],
|
||||
text_labels: list[str],
|
||||
batch_size: int = 32,
|
||||
return_debug_details: bool = False,
|
||||
) -> PipelineBatchOutput:
|
||||
"""Run full pipeline and return cropped images + hashes + debug metadata.
|
||||
|
||||
Args:
|
||||
images: Sequence of PIL Images.
|
||||
text_labels: Text labels for detection (same for all images).
|
||||
batch_size: Batch size for DINO feature extraction.
|
||||
return_debug_details: Include boxes/scores/labels/masks in debug output.
|
||||
|
||||
Returns:
|
||||
PipelineBatchOutput with final cropped images, binary hash bits,
|
||||
and per-image debug metadata.
|
||||
"""
|
||||
if batch_size <= 0:
|
||||
raise ValueError("batch_size must be greater than 0")
|
||||
|
||||
image_list = list(images)
|
||||
if not image_list:
|
||||
return PipelineBatchOutput(
|
||||
hash_bits=torch.empty(
|
||||
(0, self.hash_bits), dtype=torch.int32, device=self.device
|
||||
),
|
||||
cropped_images=[],
|
||||
debug_meta=[],
|
||||
)
|
||||
|
||||
packets = self._build_frame_packets(image_list)
|
||||
detections_per_image = self._attach_detections(packets, text_labels)
|
||||
self._attach_masks(packets)
|
||||
self._select_candidates(packets)
|
||||
self._render_filtered_images(packets)
|
||||
self._render_cropped_images(packets, detections_per_image)
|
||||
|
||||
cropped_images = [
|
||||
packet.cropped_image if packet.cropped_image is not None else packet.image
|
||||
for packet in packets
|
||||
]
|
||||
|
||||
all_bits: list[torch.Tensor] = []
|
||||
for index in range(0, len(cropped_images), batch_size):
|
||||
sub_batch = cropped_images[index : index + batch_size]
|
||||
tokens = self.extract_dino_batch(sub_batch)
|
||||
bits = self.compress_batch(tokens)
|
||||
all_bits.append(bits)
|
||||
|
||||
hash_bits = torch.cat(all_bits, dim=0)
|
||||
debug_meta = self._build_debug_meta(packets, return_debug_details)
|
||||
return PipelineBatchOutput(
|
||||
hash_bits=hash_bits,
|
||||
cropped_images=cropped_images,
|
||||
debug_meta=debug_meta,
|
||||
)
|
||||
|
||||
def forward_batch(
|
||||
self,
|
||||
images: Sequence[Image.Image],
|
||||
@@ -281,29 +422,11 @@ class HashPipeline:
|
||||
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)
|
||||
return self.process_batch(
|
||||
images=images,
|
||||
text_labels=text_labels,
|
||||
batch_size=batch_size,
|
||||
).hash_bits
|
||||
|
||||
def extract_features_dataset(
|
||||
self,
|
||||
@@ -330,15 +453,15 @@ class HashPipeline:
|
||||
(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)
|
||||
processed = self.process_batch(
|
||||
images=image_list,
|
||||
text_labels=text_labels,
|
||||
batch_size=batch_size,
|
||||
).cropped_images
|
||||
|
||||
all_features: list[torch.Tensor] = []
|
||||
for i in range(0, len(processed), batch_size):
|
||||
sub_batch = processed[i : i + batch_size]
|
||||
for index in range(0, len(processed), batch_size):
|
||||
sub_batch = processed[index : index + batch_size]
|
||||
tokens = self.extract_dino_batch(sub_batch)
|
||||
features = tokens.mean(dim=1)
|
||||
all_features.append(F.normalize(features, dim=-1))
|
||||
|
||||
Reference in New Issue
Block a user