mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
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:
@@ -7,7 +7,6 @@ from transformers import (
|
|||||||
AutoImageProcessor,
|
AutoImageProcessor,
|
||||||
AutoModel,
|
AutoModel,
|
||||||
Owlv2ForObjectDetection,
|
Owlv2ForObjectDetection,
|
||||||
Owlv2Model,
|
|
||||||
Owlv2Processor,
|
Owlv2Processor,
|
||||||
Sam2Model,
|
Sam2Model,
|
||||||
Sam2Processor,
|
Sam2Processor,
|
||||||
@@ -47,7 +46,7 @@ def load_dino_model(
|
|||||||
|
|
||||||
def load_owlv2_model(
|
def load_owlv2_model(
|
||||||
model_name: str = "google/owlv2-base-patch16-ensemble",
|
model_name: str = "google/owlv2-base-patch16-ensemble",
|
||||||
) -> tuple[Owlv2Processor, Owlv2Model]:
|
) -> tuple[Owlv2Processor, Owlv2ForObjectDetection]:
|
||||||
device = get_device()
|
device = get_device()
|
||||||
|
|
||||||
processor = Owlv2Processor.from_pretrained(model_name)
|
processor = Owlv2Processor.from_pretrained(model_name)
|
||||||
|
|||||||
@@ -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
|
from .utils import extract_masked_region
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
|
"detect_objects",
|
||||||
|
"detect_objects_batch",
|
||||||
"generate_proposals",
|
"generate_proposals",
|
||||||
"generate_proposals_batch",
|
"generate_proposals_batch",
|
||||||
"extract_masked_region",
|
"extract_masked_region",
|
||||||
|
|||||||
@@ -1,13 +1,25 @@
|
|||||||
"""SAM mask proposal generation via bounding box prompts."""
|
"""SAM mask proposal generation via bounding box prompts."""
|
||||||
|
|
||||||
from typing import Any, Sequence
|
from typing import Any, Sequence, TypedDict
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
from transformers import Sam2Model, Sam2Processor
|
from transformers import (
|
||||||
|
Owlv2ForObjectDetection,
|
||||||
|
Owlv2Processor,
|
||||||
|
Sam2Model,
|
||||||
|
Sam2Processor,
|
||||||
|
)
|
||||||
from utils import get_device
|
from utils import get_device
|
||||||
|
|
||||||
|
|
||||||
|
class DetectionResult(TypedDict):
|
||||||
|
bbox: list[float]
|
||||||
|
score: float
|
||||||
|
label: str
|
||||||
|
|
||||||
|
|
||||||
def generate_proposals(
|
def generate_proposals(
|
||||||
model: Sam2Model,
|
model: Sam2Model,
|
||||||
processor: Sam2Processor,
|
processor: Sam2Processor,
|
||||||
@@ -90,6 +102,103 @@ def generate_proposals_batch(
|
|||||||
return [_masks_to_proposals(image_masks) for image_masks in all_masks]
|
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]]:
|
def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]:
|
||||||
"""Convert model output masks to list of mask dicts."""
|
"""Convert model output masks to list of mask dicts."""
|
||||||
mask_array = _to_numpy_mask_array(masks)
|
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,
|
"predicted_iou": None,
|
||||||
"stability_score": 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
|
||||||
|
|||||||
Reference in New Issue
Block a user