mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
94ed05a039
- Add Owlv2ForObjectDetection and Owlv2Processor imports to model_loader - Refactor load_dino_model to return tuple of processor and model - Rewrite generate_proposals_batch to group images by bbox count for efficient batching - Add _normalize_single_bbox_list helper for bbox normalization - Update verification.py to use new pipeline architecture with detect/segment/filter/crop steps
345 lines
10 KiB
Python
345 lines
10 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
|
|
"""
|
|
normalized_bboxes = _normalize_single_bbox_list(bboxes)
|
|
if not normalized_bboxes:
|
|
return []
|
|
|
|
device = get_device()
|
|
image_rgb = image.convert("RGB")
|
|
input_boxes = [normalized_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 []
|
|
|
|
if len(image_list) != len(bboxes_per_image):
|
|
raise ValueError(
|
|
f"Length mismatch: {len(image_list)} images, {len(bboxes_per_image)} bbox groups"
|
|
)
|
|
|
|
normalized_bboxes_per_image = [
|
|
_normalize_single_bbox_list(bboxes) for bboxes in bboxes_per_image
|
|
]
|
|
|
|
proposals_per_image: list[list[dict[str, Any]]] = [[] for _ in image_list]
|
|
valid_indices = [
|
|
index for index, bboxes in enumerate(normalized_bboxes_per_image) if bboxes
|
|
]
|
|
if not valid_indices:
|
|
return proposals_per_image
|
|
|
|
device = get_device()
|
|
grouped_indices: dict[int, list[int]] = {}
|
|
for image_index in valid_indices:
|
|
box_count = len(normalized_bboxes_per_image[image_index])
|
|
grouped_indices.setdefault(box_count, []).append(image_index)
|
|
|
|
for grouped_image_indices in grouped_indices.values():
|
|
image_rgb_list = [
|
|
image_list[image_index].convert("RGB")
|
|
for image_index in grouped_image_indices
|
|
]
|
|
batched_input_boxes = [
|
|
normalized_bboxes_per_image[image_index]
|
|
for image_index in grouped_image_indices
|
|
]
|
|
|
|
inputs = processor(
|
|
images=image_rgb_list,
|
|
input_boxes=batched_input_boxes,
|
|
return_tensors="pt",
|
|
).to(device)
|
|
|
|
outputs = model(**inputs, multimask_output=False)
|
|
all_masks = processor.post_process_masks(
|
|
outputs.pred_masks.cpu(),
|
|
inputs["original_sizes"],
|
|
)
|
|
|
|
batched_proposals = [
|
|
_masks_to_proposals(image_masks) for image_masks in all_masks
|
|
]
|
|
for output_index, image_index in enumerate(grouped_image_indices):
|
|
proposals_per_image[image_index] = batched_proposals[output_index]
|
|
|
|
return proposals_per_image
|
|
|
|
|
|
def _normalize_single_bbox_list(bboxes: Sequence[Sequence[float]]) -> list[list[float]]:
|
|
normalized: list[list[float]] = []
|
|
for bbox in bboxes:
|
|
bbox_list = list(bbox)
|
|
if len(bbox_list) != 4:
|
|
continue
|
|
normalized.append([float(value) for value in bbox_list])
|
|
|
|
return normalized
|
|
|
|
|
|
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 []
|
|
|
|
if mask_array.ndim < 2:
|
|
return []
|
|
|
|
if mask_array.ndim == 2:
|
|
mask_array = np.expand_dims(mask_array, axis=0)
|
|
else:
|
|
height, width = mask_array.shape[-2], mask_array.shape[-1]
|
|
mask_array = mask_array.reshape(-1, height, width)
|
|
|
|
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
|