mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
42acb3ee1b
- 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
155 lines
4.2 KiB
Python
155 lines
4.2 KiB
Python
"""SAM mask proposal generation via bounding box prompts."""
|
|
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from transformers import Sam2Model, Sam2Processor
|
|
from utils import get_device
|
|
|
|
|
|
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 _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,
|
|
}
|