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
This commit is contained in:
2026-04-02 19:41:54 +08:00
parent 42acb3ee1b
commit af0531a5eb
3 changed files with 157 additions and 6 deletions

View File

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

View File

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

View File

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