Mini-Nav/mini-nav/utils/image.py

from typing import Any, Sequence

import numpy as np
from PIL import Image


def segment_image(
    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)
    raw_masks = raw_output.get("masks", 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 = [m for m in normalized_masks if int(m["area"]) >= min_area]

    if not filtered_masks:
        return []

    sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
    return sorted_masks[:max_masks]


def segment_image_dataset(
    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]]]:
    image_list = list(images)
    return [
        segment_image(
            mask_generator,
            image,
            min_area=min_area,
            max_masks=max_masks,
            points_per_batch=points_per_batch,
        )
        for image in image_list
    ]


def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
    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

    return None


def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
    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:
    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,
    }


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))