SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(compressors): reorganize SAM utilities into proposal module

Browse Source
This commit is contained in:
SikongJueluo
2026-03-30 20:09:12 +08:00
parent f421b0c56b
commit 26b00e556a
7 changed files with 89 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
mini-nav/utils/model.py → mini-nav/compressors/model_loader.py
View File

@@ -1,12 +1,12 @@
"""Model loading utilities for DINO, SAM2 and HashCompressor."""
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
import torch
from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
from .common import get_device
from utils import get_device
if TYPE_CHECKING:
from compressors.hash_compressor import HashCompressor
@@ -26,7 +26,7 @@ def load_sam_model(
def load_dino_model(
model_name: str = "facebook/dinov2-large",
) -> tuple[AutoImageProcessor, AutoModel]:
) -> tuple[Any, Any]:
device = get_device()
processor = AutoImageProcessor.from_pretrained(model_name)

12
mini-nav/compressors/pipeline.py
View File

@@ -8,9 +8,13 @@ import torch.nn.functional as F
from PIL import Image
from .object_score import select_best_mask
from .proposal import (
extract_masked_region,
generate_proposals,
generate_proposals_batch,
)
from utils import get_device
from utils.image import extract_masked_region, segment_image, segment_image_dataset
from utils.model import (
from .model_loader import (
get_dino_dim,
load_dino_model,
load_hash_compressor,
@@ -101,7 +105,7 @@ class HashPipeline(nn.Module):
Returns:
Masked image containing only the largest object, or original if no masks.
"""
masks = segment_image(
masks = generate_proposals(
self.mask_generator,
image,
min_area=self.sam_min_mask_area,
@@ -122,7 +126,7 @@ class HashPipeline(nn.Module):
images: Sequence[Image.Image],
) -> list[Image.Image]:
image_list = list(images)
masks_dataset = segment_image_dataset(
masks_dataset = generate_proposals_batch(
self.mask_generator,
image_list,
min_area=self.sam_min_mask_area,

10
mini-nav/compressors/proposal/__init__.py Normal file
View File

@@ -0,0 +1,10 @@
"""Proposal module — SAM mask generation and extraction."""
from .core import generate_proposals, generate_proposals_batch
from .utils import extract_masked_region
__all__ = [
"generate_proposals",
"generate_proposals_batch",
"extract_masked_region",
]

88
mini-nav/utils/image.py → mini-nav/compressors/proposal/core.py
View File

@@ -1,10 +1,13 @@
"""SAM mask proposal generation."""
from typing import Any, Sequence
import torch
import numpy as np
from PIL import Image
def segment_image(
def generate_proposals(
mask_generator: Any,
image: Image.Image,
min_area: int = 32 * 32,
@@ -19,6 +22,7 @@ def segment_image(
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)
@@ -34,37 +38,44 @@ def segment_image(
)
def segment_image_dataset(
def generate_proposals_batch(
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]]]:
"""Segment a batch of images using SAM.
Args:
mask_generator: SAM2 mask generator.
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.
Returns:
List of lists of mask dictionaries, one inner list per image.
"""
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
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
]
return [
_normalize_and_filter_masks(
@@ -77,6 +88,7 @@ def segment_image_dataset(
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
@@ -95,6 +107,7 @@ def _normalize_and_filter_masks(
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)
@@ -135,23 +148,20 @@ def _normalize_and_filter_masks(
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
try:
import torch
if isinstance(mask_like, torch.Tensor):
return mask_like.detach().cpu().numpy()
except ImportError:
pass
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
@@ -174,6 +184,7 @@ def _to_mask_dict(mask_like: Any) -> dict[str, Any] | 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)
@@ -193,24 +204,3 @@ def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
"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))

25
mini-nav/compressors/proposal/utils.py Normal file
View File

@@ -0,0 +1,25 @@
"""Mask extraction utilities."""
import numpy as np
from PIL import Image
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))

9
mini-nav/utils/__init__.py
View File

@@ -4,8 +4,6 @@ from .feature_extractor import (
extract_single_image_feature,
infer_vector_dim,
)
from .image import extract_masked_region, segment_image, segment_image_dataset
from .model import get_dino_dim, load_dino_model, load_hash_compressor, load_sam_model
__all__ = [
"get_device",
@@ -13,11 +11,4 @@ __all__ = [
"infer_vector_dim",
"extract_single_image_feature",
"extract_batch_features",
"segment_image",
"segment_image_dataset",
"extract_masked_region",
"load_dino_model",
"load_sam_model",
"get_dino_dim",
"load_hash_compressor",
]