mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
af0531a5eb
- 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
300 lines
8.9 KiB
Python
300 lines
8.9 KiB
Python
"""SAM mask proposal generation via bounding box prompts."""
|
|
|
|
from typing import Any, Sequence, TypedDict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
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,
|
|
image: Image.Image,
|
|
bboxes: list[list[float]],
|
|
) -> list[dict[str, Any]]:
|
|
"""Segment regions in image using SAM2 with bounding box prompts.
|
|
|
|
Args:
|
|
model: Sam2Model instance.
|
|
processor: Sam2Processor instance.
|
|
image: PIL Image to segment.
|
|
bboxes: Bounding boxes as [[x1, y1, x2, y2], ...].
|
|
|
|
Returns:
|
|
List of mask dictionaries with keys:
|
|
- segment: Binary mask (numpy array)
|
|
- area: Mask area in pixels
|
|
- bbox: Bounding box [x, y, width, height]
|
|
- predicted_iou: Model's confidence in the mask
|
|
- stability_score: Stability score for the mask
|
|
"""
|
|
if not bboxes:
|
|
return []
|
|
|
|
device = get_device()
|
|
image_rgb = image.convert("RGB")
|
|
input_boxes = [bboxes]
|
|
inputs = processor(
|
|
images=image_rgb,
|
|
input_boxes=input_boxes,
|
|
return_tensors="pt",
|
|
).to(device)
|
|
|
|
outputs = model(**inputs, multimask_output=False)
|
|
masks = processor.post_process_masks(
|
|
outputs.pred_masks.cpu(),
|
|
inputs["original_sizes"],
|
|
)[0]
|
|
|
|
return _masks_to_proposals(masks)
|
|
|
|
|
|
def generate_proposals_batch(
|
|
model: Sam2Model,
|
|
processor: Sam2Processor,
|
|
images: Sequence[Image.Image],
|
|
bboxes_per_image: list[list[list[float]]],
|
|
) -> list[list[dict[str, Any]]]:
|
|
"""Segment a batch of images using SAM2 with bounding box prompts.
|
|
|
|
Args:
|
|
model: Sam2Model instance.
|
|
processor: Sam2Processor instance.
|
|
images: Sequence of PIL Images to segment.
|
|
bboxes_per_image: Bounding boxes per image, outer list matches images length.
|
|
|
|
Returns:
|
|
List of lists of mask dictionaries, one inner list per image.
|
|
"""
|
|
image_list = list(images)
|
|
if not image_list:
|
|
return []
|
|
|
|
device = get_device()
|
|
image_rgb_list = [img.convert("RGB") for img in image_list]
|
|
|
|
inputs = processor(
|
|
images=image_rgb_list,
|
|
input_boxes=bboxes_per_image,
|
|
return_tensors="pt",
|
|
).to(device)
|
|
|
|
outputs = model(**inputs, multimask_output=False)
|
|
all_masks = processor.post_process_masks(
|
|
outputs.pred_masks.cpu(),
|
|
inputs["original_sizes"],
|
|
)
|
|
|
|
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)
|
|
if mask_array is None:
|
|
return []
|
|
|
|
# Ensure 3D: [num_masks, H, W]
|
|
if mask_array.ndim == 2:
|
|
mask_array = np.expand_dims(mask_array, axis=0)
|
|
|
|
if mask_array.ndim != 3:
|
|
return []
|
|
|
|
# Remove batch dim if present: [1, num_masks, H, W] → [num_masks, H, W]
|
|
if mask_array.ndim == 3 and mask_array.shape[0] == 1:
|
|
mask_array = mask_array[0]
|
|
|
|
proposals: list[dict[str, Any]] = []
|
|
for single_mask in mask_array:
|
|
mask_dict = _build_mask_dict(single_mask)
|
|
if mask_dict is not None:
|
|
proposals.append(mask_dict)
|
|
|
|
return proposals
|
|
|
|
|
|
def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
|
|
"""Convert mask-like object to numpy array."""
|
|
if mask_like is None:
|
|
return None
|
|
if isinstance(mask_like, np.ndarray):
|
|
return mask_like
|
|
|
|
import torch
|
|
|
|
if isinstance(mask_like, torch.Tensor):
|
|
return mask_like.detach().cpu().numpy()
|
|
|
|
return None
|
|
|
|
|
|
def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
|
|
"""Build a mask dictionary from a 2D boolean numpy array."""
|
|
if mask_array.ndim != 2:
|
|
return None
|
|
segment = mask_array.astype(bool)
|
|
area = int(segment.sum())
|
|
if area <= 0:
|
|
return None
|
|
|
|
ys, xs = np.where(segment)
|
|
min_y, max_y = int(ys.min()), int(ys.max())
|
|
min_x, max_x = int(xs.min()), int(xs.max())
|
|
bbox = [min_x, min_y, max_x - min_x + 1, max_y - min_y + 1]
|
|
|
|
return {
|
|
"segment": segment,
|
|
"area": area,
|
|
"bbox": bbox,
|
|
"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
|