mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
refactor(compressors): switch SAM from automatic mask generation to bbox-prompted segmentation
- Replace SAM2AutomaticMaskGenerator pipeline with Sam2Processor+Sam2Model - Freeze SAM model parameters at load time, removing torch.no_grad() at call sites - Rewrite proposal/core.py to use bbox prompts instead of automatic point sampling - Add bboxes parameter to all HashPipeline public methods (forward, forward_dataset, extract_features, extract_features_dataset) - Extract mask filtering logic (_filter_masks) from proposal into pipeline - Rename object_score/ to filter/ - Add load_owlv2_model to model_loader - Rename notebooks/test.py to habitat_sim_setup.py
This commit is contained in:
@@ -1,27 +1,26 @@
|
||||
"""SAM mask proposal generation."""
|
||||
"""SAM mask proposal generation via bounding box prompts."""
|
||||
|
||||
from typing import Any, Sequence
|
||||
import torch
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from transformers import Sam2Model, Sam2Processor
|
||||
from utils import get_device
|
||||
|
||||
|
||||
def generate_proposals(
|
||||
mask_generator: Any,
|
||||
model: Sam2Model,
|
||||
processor: Sam2Processor,
|
||||
image: Image.Image,
|
||||
min_area: int = 32 * 32,
|
||||
max_masks: int = 5,
|
||||
points_per_batch: int = 64,
|
||||
bboxes: list[list[float]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Segment image using SAM to extract object masks.
|
||||
"""Segment regions in image using SAM2 with bounding box prompts.
|
||||
|
||||
Args:
|
||||
mask_generator: SAM2 mask generator.
|
||||
model: Sam2Model instance.
|
||||
processor: Sam2Processor instance.
|
||||
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.
|
||||
bboxes: Bounding boxes as [[x1, y1, x2, y2], ...].
|
||||
|
||||
Returns:
|
||||
List of mask dictionaries with keys:
|
||||
@@ -31,28 +30,40 @@ def generate_proposals(
|
||||
- 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")
|
||||
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
|
||||
)
|
||||
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(
|
||||
mask_generator: Any,
|
||||
model: Sam2Model,
|
||||
processor: Sam2Processor,
|
||||
images: Sequence[Image.Image],
|
||||
min_area: int = 32 * 32,
|
||||
max_masks: int = 5,
|
||||
points_per_batch: int = 64,
|
||||
bboxes_per_image: list[list[list[float]]],
|
||||
) -> list[list[dict[str, Any]]]:
|
||||
"""Segment a batch of images using SAM.
|
||||
"""Segment a batch of images using SAM2 with bounding box prompts.
|
||||
|
||||
Args:
|
||||
mask_generator: SAM2 mask generator.
|
||||
model: Sam2Model instance.
|
||||
processor: Sam2Processor instance.
|
||||
images: Sequence of PIL Images to segment.
|
||||
min_area: Minimum mask area threshold in pixels.
|
||||
max_masks: Maximum number of masks to return per image.
|
||||
points_per_batch: Number of prompt points to process in each batch.
|
||||
bboxes_per_image: Bounding boxes per image, outer list matches images length.
|
||||
|
||||
Returns:
|
||||
List of lists of mask dictionaries, one inner list per image.
|
||||
@@ -61,90 +72,48 @@ def generate_proposals_batch(
|
||||
if not image_list:
|
||||
return []
|
||||
|
||||
image_rgb_list = [image.convert("RGB") for image in image_list]
|
||||
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
|
||||
]
|
||||
device = get_device()
|
||||
image_rgb_list = [img.convert("RGB") for img in image_list]
|
||||
|
||||
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
|
||||
]
|
||||
inputs = processor(
|
||||
images=image_rgb_list,
|
||||
input_boxes=bboxes_per_image,
|
||||
return_tensors="pt",
|
||||
).to(device)
|
||||
|
||||
|
||||
def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None:
|
||||
"""Attempt to split raw batch output into per-image results."""
|
||||
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]]:
|
||||
"""Normalize raw SAM output into mask dicts and filter by area/count."""
|
||||
raw_masks = (
|
||||
raw_output.get("masks", raw_output)
|
||||
if isinstance(raw_output, dict)
|
||||
else raw_output
|
||||
outputs = model(**inputs, multimask_output=False)
|
||||
all_masks = processor.post_process_masks(
|
||||
outputs.pred_masks.cpu(),
|
||||
inputs["original_sizes"],
|
||||
)
|
||||
|
||||
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)
|
||||
return [_masks_to_proposals(image_masks) for image_masks in all_masks]
|
||||
|
||||
if not normalized_masks:
|
||||
|
||||
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 []
|
||||
|
||||
filtered_masks = [
|
||||
mask for mask in normalized_masks if int(mask["area"]) >= min_area
|
||||
]
|
||||
if not filtered_masks:
|
||||
# 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 []
|
||||
|
||||
sorted_masks = sorted(filtered_masks, key=lambda mask: mask["area"], reverse=True)
|
||||
return sorted_masks[:max_masks]
|
||||
# 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:
|
||||
@@ -154,35 +123,14 @@ def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | 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 _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
|
||||
"""Convert a single mask-like object to a standardized mask dict."""
|
||||
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:
|
||||
"""Build a mask dictionary from a 2D boolean numpy array."""
|
||||
if mask_array.ndim != 2:
|
||||
|
||||
Reference in New Issue
Block a user