From af0531a5eb216c75e0852c6c6e9f17a5a8dd040b Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Thu, 2 Apr 2026 19:41:54 +0800 Subject: [PATCH] feat(compressors/proposal): add OWLv2 object detection pipeline and typed results - switch OWLv2 loader return type to Owlv2ForObjectDetection - add detect_objects and detect_objects_batch with two-stage score filtering - add DetectionResult typed dict and conversion helper for post-processed outputs - export new detection APIs from proposal module --- mini-nav/compressors/model_loader.py | 3 +- mini-nav/compressors/proposal/__init__.py | 11 +- mini-nav/compressors/proposal/core.py | 149 +++++++++++++++++++++- 3 files changed, 157 insertions(+), 6 deletions(-) diff --git a/mini-nav/compressors/model_loader.py b/mini-nav/compressors/model_loader.py index 15aae84..5365031 100644 --- a/mini-nav/compressors/model_loader.py +++ b/mini-nav/compressors/model_loader.py @@ -7,7 +7,6 @@ from transformers import ( AutoImageProcessor, AutoModel, Owlv2ForObjectDetection, - Owlv2Model, Owlv2Processor, Sam2Model, Sam2Processor, @@ -47,7 +46,7 @@ def load_dino_model( def load_owlv2_model( model_name: str = "google/owlv2-base-patch16-ensemble", -) -> tuple[Owlv2Processor, Owlv2Model]: +) -> tuple[Owlv2Processor, Owlv2ForObjectDetection]: device = get_device() processor = Owlv2Processor.from_pretrained(model_name) diff --git a/mini-nav/compressors/proposal/__init__.py b/mini-nav/compressors/proposal/__init__.py index 3d35b3b..64a3d13 100644 --- a/mini-nav/compressors/proposal/__init__.py +++ b/mini-nav/compressors/proposal/__init__.py @@ -1,9 +1,16 @@ -"""Proposal module — SAM mask generation and extraction.""" +"""Proposal module — SAM mask generation and OWLv2 object detection.""" -from .core import generate_proposals, generate_proposals_batch +from .core import ( + detect_objects, + detect_objects_batch, + 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/core.py b/mini-nav/compressors/proposal/core.py index 486f348..d1c1e2f 100644 --- a/mini-nav/compressors/proposal/core.py +++ b/mini-nav/compressors/proposal/core.py @@ -1,13 +1,25 @@ """SAM mask proposal generation via bounding box prompts.""" -from typing import Any, Sequence +from typing import Any, Sequence, TypedDict import numpy as np +import torch from PIL import Image -from transformers import Sam2Model, Sam2Processor +from transformers import ( + Owlv2ForObjectDetection, + Owlv2Processor, + Sam2Model, + Sam2Processor, +) from utils import get_device +class DetectionResult(TypedDict): + bbox: list[float] + score: float + label: str + + def generate_proposals( model: Sam2Model, processor: Sam2Processor, @@ -90,6 +102,103 @@ def generate_proposals_batch( return [_masks_to_proposals(image_masks) for image_masks in all_masks] +def detect_objects( + model: Owlv2ForObjectDetection, + processor: Owlv2Processor, + image: Image.Image, + text_labels: list[str], + score_threshold: float = 0.25, + postprocess_threshold: float = 0.1, +) -> list[DetectionResult]: + """Detect objects in a single image using OWLv2 with text labels. + + Runs OWLv2 object detection on a single image using provided text label + queries. Applies post-process thresholding first, then filters by + score_threshold. + + Args: + model: Owlv2ForObjectDetection instance. + processor: Owlv2Processor instance. + image: PIL Image to detect objects in. + text_labels: List of label groups, e.g., ["cat", "dog"], ["car"]. + score_threshold: Minimum score for final detection results (>=). + postprocess_threshold: Threshold for processor's post-processing (>=). + + Returns: + List of DetectionResult dicts with bbox, score, label. + """ + device = get_device() + inputs = processor(text=text_labels, images=image, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model(**inputs) + target_sizes = [(image.height, image.width)] + result = processor.post_process_grounded_object_detection( + outputs=outputs, + target_sizes=target_sizes, + threshold=postprocess_threshold, + text_labels=[text_labels], + )[0] + return _to_detection_results(result, score_threshold) + + +def detect_objects_batch( + model: Owlv2ForObjectDetection, + processor: Owlv2Processor, + images: Sequence[Image.Image], + text_labels_per_image: list[list[str]], + score_threshold: float = 0.25, + postprocess_threshold: float = 0.1, +) -> list[list[DetectionResult]]: + """Detect objects in a batch of images using OWLv2 with text labels. + + Runs OWLv2 object detection on multiple images using per-image text label + queries. Processes all images in a single batch inference call. + + Args: + model: Owlv2ForObjectDetection instance. + processor: Owlv2Processor instance. + images: Sequence of PIL Images to detect objects in. + text_labels_per_image: Text labels per image, outer list matches images length. + Each inner list is per-image label, e.g., [["cat", "dog"], ["car"]]. + score_threshold: Minimum score for final detection results (>=). + postprocess_threshold: Threshold for processor's post-processing (>=). + + Returns: + List of lists of DetectionResult dicts, one inner list per image. + + Raises: + ValueError: If len(images) != len(text_labels_per_image). + """ + image_list = list(images) + if not image_list: + return [] + + if len(image_list) != len(text_labels_per_image): + raise ValueError( + f"Length mismatch: {len(image_list)} images, " + f"{len(text_labels_per_image)} text label groups" + ) + + device = get_device() + inputs = processor( + text=text_labels_per_image, + images=image_list, + return_tensors="pt", + ).to(device) + with torch.no_grad(): + outputs = model(**inputs) + + target_sizes = [(img.height, img.width) for img in image_list] + results = processor.post_process_grounded_object_detection( + outputs=outputs, + target_sizes=target_sizes, + threshold=postprocess_threshold, + text_labels=text_labels_per_image, + ) + + return [_to_detection_results(result, score_threshold) for result in results] + + def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]: """Convert model output masks to list of mask dicts.""" mask_array = _to_numpy_mask_array(masks) @@ -152,3 +261,39 @@ def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None: "predicted_iou": None, "stability_score": None, } + + +def _to_detection_results( + result: dict[str, Any], score_threshold: float +) -> list[DetectionResult]: + """Convert OWLv2 post-process result to detection results list. + + Extracts boxes, scores, and text labels from result dict, converts tensors + to Python native types, and filters by score threshold. + + Args: + result: OWLv2 post-process result dict with keys 'boxes', 'scores', + 'text_labels'. + score_threshold: Minimum score to include detection (>= threshold). + + Returns: + List of DetectionResult dicts with bbox, score, label. + """ + boxes = result["boxes"] + scores = result["scores"] + text_labels = result["text_labels"] + + detections: list[DetectionResult] = [] + for box, score, label in zip(boxes, scores, text_labels): + score_float = float(score.item()) + if score_float >= score_threshold: + bbox = [float(v) for v in box.tolist()] + detections.append( + { + "bbox": bbox, + "score": score_float, + "label": str(label), + } + ) + + return detections