mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
147 lines
4.2 KiB
Python
147 lines
4.2 KiB
Python
from typing import Any
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
def segment_image(
|
|
mask_generator: Any,
|
|
image: Image.Image,
|
|
min_area: int = 32 * 32,
|
|
max_masks: int = 5,
|
|
points_per_batch: int = 64,
|
|
) -> list[dict[str, Any]]:
|
|
"""Segment image using SAM to extract object masks.
|
|
|
|
Args:
|
|
mask_generator: SAM2 mask generator.
|
|
image: PIL Image to segment.
|
|
min_area: Minimum mask area threshold in pixels.
|
|
max_masks: Maximum number of masks to return.
|
|
points_per_batch: Number of prompt points to process in each batch.
|
|
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
|
|
"""
|
|
image_rgb = image.convert("RGB")
|
|
raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
|
|
raw_masks = raw_output.get("masks", raw_output)
|
|
|
|
normalized_masks: list[dict[str, Any]] = []
|
|
|
|
if isinstance(raw_masks, list):
|
|
if raw_masks and isinstance(raw_masks[0], dict):
|
|
normalized_masks = raw_masks
|
|
else:
|
|
for mask_like in raw_masks:
|
|
mask_dict = _to_mask_dict(mask_like)
|
|
if mask_dict is not None:
|
|
normalized_masks.append(mask_dict)
|
|
else:
|
|
mask_array = _to_numpy_mask_array(raw_masks)
|
|
if mask_array is not None:
|
|
if mask_array.ndim == 2:
|
|
mask_array = np.expand_dims(mask_array, axis=0)
|
|
if mask_array.ndim == 3:
|
|
for single_mask in mask_array:
|
|
mask_dict = _to_mask_dict(single_mask)
|
|
if mask_dict is not None:
|
|
normalized_masks.append(mask_dict)
|
|
|
|
if not normalized_masks:
|
|
return []
|
|
|
|
filtered_masks = [m for m in normalized_masks if int(m["area"]) >= min_area]
|
|
|
|
if not filtered_masks:
|
|
return []
|
|
|
|
sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
|
|
return sorted_masks[:max_masks]
|
|
|
|
|
|
def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
|
|
if mask_like is None:
|
|
return None
|
|
if isinstance(mask_like, np.ndarray):
|
|
return mask_like
|
|
|
|
try:
|
|
import torch
|
|
|
|
if isinstance(mask_like, torch.Tensor):
|
|
return mask_like.detach().cpu().numpy()
|
|
except ImportError:
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
|
|
if isinstance(mask_like, dict):
|
|
if "area" in mask_like and "bbox" in mask_like and "segment" in mask_like:
|
|
return mask_like
|
|
|
|
segment = mask_like.get("segment")
|
|
if segment is None and "mask" in mask_like:
|
|
segment = mask_like["mask"]
|
|
if segment is None:
|
|
return None
|
|
|
|
mask_array = _to_numpy_mask_array(segment)
|
|
if mask_array is None:
|
|
return None
|
|
return _build_mask_dict(mask_array)
|
|
|
|
mask_array = _to_numpy_mask_array(mask_like)
|
|
if mask_array is None:
|
|
return None
|
|
return _build_mask_dict(mask_array)
|
|
|
|
|
|
def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
|
|
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 extract_masked_region(
|
|
image: Image.Image,
|
|
mask: np.ndarray,
|
|
) -> Image.Image:
|
|
"""Extract masked region from image.
|
|
|
|
Args:
|
|
image: Original PIL Image.
|
|
mask: Binary mask as numpy array (True = keep).
|
|
|
|
Returns:
|
|
PIL Image with only the masked region visible.
|
|
"""
|
|
image_np = np.array(image.convert("RGB"))
|
|
|
|
# Apply mask
|
|
masked_np = image_np * mask[:, :, np.newaxis]
|
|
|
|
return Image.fromarray(masked_np.astype(np.uint8))
|