Mini-Nav/mini-nav/compressors/proposal/core.py

"""SAM mask proposal generation."""

from typing import Any, Sequence
import torch

import numpy as np
from PIL import Image


def generate_proposals(
    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 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]
    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(
            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:
    """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
    )

    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:
    """Convert mask-like object to numpy array."""
    if mask_like is None:
        return None
    if isinstance(mask_like, np.ndarray):
        return mask_like

    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:
        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,
    }