mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
refactor(compressors): reorganize SAM utilities into proposal module
This commit is contained in:
@@ -1,216 +0,0 @@
|
||||
from typing import Any, Sequence
|
||||
|
||||
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)
|
||||
return _normalize_and_filter_masks(
|
||||
raw_output, min_area=min_area, max_masks=max_masks
|
||||
)
|
||||
|
||||
|
||||
def segment_image_dataset(
|
||||
mask_generator: Any,
|
||||
images: Sequence[Image.Image],
|
||||
min_area: int = 32 * 32,
|
||||
max_masks: int = 5,
|
||||
points_per_batch: int = 64,
|
||||
) -> list[list[dict[str, Any]]]:
|
||||
image_list = list(images)
|
||||
if not image_list:
|
||||
return []
|
||||
|
||||
image_rgb_list = [image.convert("RGB") for image in image_list]
|
||||
try:
|
||||
raw_batch_output = mask_generator(
|
||||
image_rgb_list,
|
||||
points_per_batch=points_per_batch,
|
||||
)
|
||||
batch_items = _split_batch_output(
|
||||
raw_batch_output, expected_size=len(image_list)
|
||||
)
|
||||
if batch_items is not None:
|
||||
return [
|
||||
_normalize_and_filter_masks(
|
||||
batch_item,
|
||||
min_area=min_area,
|
||||
max_masks=max_masks,
|
||||
)
|
||||
for batch_item in batch_items
|
||||
]
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
return [
|
||||
_normalize_and_filter_masks(
|
||||
mask_generator(image_rgb, points_per_batch=points_per_batch),
|
||||
min_area=min_area,
|
||||
max_masks=max_masks,
|
||||
)
|
||||
for image_rgb in image_rgb_list
|
||||
]
|
||||
|
||||
|
||||
def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None:
|
||||
if isinstance(raw_output, list):
|
||||
if len(raw_output) == expected_size:
|
||||
return raw_output
|
||||
return None
|
||||
|
||||
if isinstance(raw_output, dict):
|
||||
raw_masks = raw_output.get("masks", raw_output)
|
||||
if isinstance(raw_masks, list) and len(raw_masks) == expected_size:
|
||||
return raw_masks
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_and_filter_masks(
|
||||
raw_output: Any,
|
||||
min_area: int,
|
||||
max_masks: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
raw_masks = (
|
||||
raw_output.get("masks", raw_output)
|
||||
if isinstance(raw_output, dict)
|
||||
else 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 = [
|
||||
mask for mask in normalized_masks if int(mask["area"]) >= min_area
|
||||
]
|
||||
if not filtered_masks:
|
||||
return []
|
||||
|
||||
sorted_masks = sorted(filtered_masks, key=lambda mask: mask["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))
|
||||
Reference in New Issue
Block a user