diff --git a/mini-nav/utils/model.py b/mini-nav/compressors/model_loader.py similarity index 92% rename from mini-nav/utils/model.py rename to mini-nav/compressors/model_loader.py index e282292..f778905 100644 --- a/mini-nav/utils/model.py +++ b/mini-nav/compressors/model_loader.py @@ -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) diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index 8af8e35..9901e43 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -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, diff --git a/mini-nav/compressors/proposal/__init__.py b/mini-nav/compressors/proposal/__init__.py new file mode 100644 index 0000000..3d35b3b --- /dev/null +++ b/mini-nav/compressors/proposal/__init__.py @@ -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", +] diff --git a/mini-nav/utils/image.py b/mini-nav/compressors/proposal/core.py similarity index 78% rename from mini-nav/utils/image.py rename to mini-nav/compressors/proposal/core.py index d6ed081..a2a9c0f 100644 --- a/mini-nav/utils/image.py +++ b/mini-nav/compressors/proposal/core.py @@ -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)) diff --git a/mini-nav/compressors/proposal/utils.py b/mini-nav/compressors/proposal/utils.py new file mode 100644 index 0000000..2a53e46 --- /dev/null +++ b/mini-nav/compressors/proposal/utils.py @@ -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)) diff --git a/mini-nav/utils/__init__.py b/mini-nav/utils/__init__.py index 792a8de..1591fce 100644 --- a/mini-nav/utils/__init__.py +++ b/mini-nav/utils/__init__.py @@ -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", ] diff --git a/notebooks/verification.py b/notebooks/verification.py index b86f810..740d635 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -33,7 +33,7 @@ def import_packages(): create_habitat_simulator, render_topdown_scene_map, ) - from utils.image import extract_masked_region, segment_image_dataset + from compressors.proposal import extract_masked_region, generate_proposals_batch return ( HabitatSimulatorConfig, @@ -53,7 +53,7 @@ def import_packages(): pl, plt, render_topdown_scene_map, - segment_image_dataset, + generate_proposals_batch, ) @@ -150,7 +150,7 @@ def build_scene_graph_pipeline( room_nodes, sam_max_masks, sam_min_area, - segment_image_dataset, + generate_proposals_batch, sim, views_per_room, ): @@ -191,7 +191,7 @@ def build_scene_graph_pipeline( rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb room_view_images.append(Image.fromarray(rgb3.astype(np.uint8))) - masks_dataset = segment_image_dataset( + masks_dataset = generate_proposals_batch( hash_pipeline.mask_generator, room_view_images, min_area=hash_pipeline.sam_min_mask_area,