Mini-Nav/mini-nav/compressors/object_score/features.py

from dataclasses import dataclass
from math import pi
from typing import Any

import numpy as np


@dataclass(frozen=True)
class MaskFeatures:
    area_ratio: float
    fill_ratio: float
    aspect_ratio: float
    touch_top: bool
    touch_bottom: bool
    touch_left: bool
    touch_right: bool
    touch_edge_count: int
    num_components: int
    largest_component_ratio: float
    num_holes: int
    perimeter: float
    compactness: float
    center: tuple[float, float]
    predicted_iou: float | None
    stability_score: float | None


def compute_mask_features(
    mask_dict: dict[str, Any],
    image_shape: tuple[int, int],
) -> MaskFeatures:
    height, width = image_shape
    if height <= 0 or width <= 0:
        raise ValueError("image_shape must be positive")

    segment_raw = np.asarray(mask_dict["segment"])
    if segment_raw.ndim != 2:
        raise ValueError("mask segment must be 2D")
    segment = segment_raw.astype(bool)

    area = int(mask_dict.get("area", int(segment.sum())))
    bbox = _get_bbox(mask_dict, segment)
    _, _, bbox_w, bbox_h = bbox

    image_area = height * width
    bbox_area = max(1, bbox_w * bbox_h)
    area_ratio = float(area) / float(image_area)
    fill_ratio = float(area) / float(bbox_area)
    aspect_ratio = float(bbox_w) / float(max(1, bbox_h))

    touch_top = bool(segment[0, :].any())
    touch_bottom = bool(segment[-1, :].any())
    touch_left = bool(segment[:, 0].any())
    touch_right = bool(segment[:, -1].any())
    touch_edge_count = (
        int(touch_top) + int(touch_bottom) + int(touch_left) + int(touch_right)
    )

    num_components, largest_component_ratio = _component_stats(segment, area)
    num_holes = _count_holes(segment)
    perimeter = _estimate_perimeter(segment)
    compactness = 0.0
    if perimeter > 0.0:
        compactness = (4.0 * pi * float(area)) / (perimeter * perimeter)
        compactness = max(0.0, min(1.0, compactness))

    ys, xs = np.where(segment)
    center = (float(xs.mean()), float(ys.mean())) if len(xs) > 0 else (0.0, 0.0)

    predicted_iou = _safe_float(mask_dict.get("predicted_iou"))
    stability_score = _safe_float(mask_dict.get("stability_score"))

    return MaskFeatures(
        area_ratio=area_ratio,
        fill_ratio=fill_ratio,
        aspect_ratio=aspect_ratio,
        touch_top=touch_top,
        touch_bottom=touch_bottom,
        touch_left=touch_left,
        touch_right=touch_right,
        touch_edge_count=touch_edge_count,
        num_components=num_components,
        largest_component_ratio=largest_component_ratio,
        num_holes=num_holes,
        perimeter=perimeter,
        compactness=compactness,
        center=center,
        predicted_iou=predicted_iou,
        stability_score=stability_score,
    )


def _safe_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _get_bbox(
    mask_dict: dict[str, Any], segment: np.ndarray
) -> tuple[int, int, int, int]:
    bbox_raw = mask_dict.get("bbox")
    if isinstance(bbox_raw, (list, tuple)) and len(bbox_raw) == 4:
        x, y, w, h = (int(v) for v in bbox_raw)
        if w > 0 and h > 0:
            return x, y, w, h

    ys, xs = np.where(segment)
    if len(xs) == 0:
        return 0, 0, 1, 1
    min_y, max_y = int(ys.min()), int(ys.max())
    min_x, max_x = int(xs.min()), int(xs.max())
    return min_x, min_y, max_x - min_x + 1, max_y - min_y + 1


def _component_stats(segment: np.ndarray, area: int) -> tuple[int, float]:
    visited = np.zeros_like(segment, dtype=bool)
    height, width = segment.shape
    component_areas: list[int] = []

    for y in range(height):
        for x in range(width):
            if not segment[y, x] or visited[y, x]:
                continue
            stack = [(y, x)]
            visited[y, x] = True
            comp_area = 0

            while stack:
                cy, cx = stack.pop()
                comp_area += 1
                neighbors = ((cy - 1, cx), (cy + 1, cx), (cy, cx - 1), (cy, cx + 1))
                for ny, nx in neighbors:
                    if ny < 0 or nx < 0 or ny >= height or nx >= width:
                        continue
                    if visited[ny, nx] or not segment[ny, nx]:
                        continue
                    visited[ny, nx] = True
                    stack.append((ny, nx))

            component_areas.append(comp_area)

    if not component_areas:
        return 0, 0.0
    largest = max(component_areas)
    largest_ratio = float(largest) / float(max(1, area))
    return len(component_areas), largest_ratio


def _count_holes(segment: np.ndarray) -> int:
    height, width = segment.shape
    inverted = ~segment
    visited = np.zeros_like(inverted, dtype=bool)

    border_stack: list[tuple[int, int]] = []
    for x in range(width):
        border_stack.append((0, x))
        border_stack.append((height - 1, x))
    for y in range(height):
        border_stack.append((y, 0))
        border_stack.append((y, width - 1))

    while border_stack:
        y, x = border_stack.pop()
        if y < 0 or x < 0 or y >= height or x >= width:
            continue
        if visited[y, x] or not inverted[y, x]:
            continue
        visited[y, x] = True
        border_stack.extend(((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)))

    holes = 0
    for y in range(height):
        for x in range(width):
            if visited[y, x] or not inverted[y, x]:
                continue
            holes += 1
            stack = [(y, x)]
            visited[y, x] = True
            while stack:
                cy, cx = stack.pop()
                neighbors = ((cy - 1, cx), (cy + 1, cx), (cy, cx - 1), (cy, cx + 1))
                for ny, nx in neighbors:
                    if ny < 0 or nx < 0 or ny >= height or nx >= width:
                        continue
                    if visited[ny, nx] or not inverted[ny, nx]:
                        continue
                    visited[ny, nx] = True
                    stack.append((ny, nx))

    return holes


def _estimate_perimeter(segment: np.ndarray) -> float:
    segment_int = segment.astype(np.int32)
    padded = np.pad(segment_int, ((1, 1), (1, 1)), mode="constant", constant_values=0)
    up = padded[:-2, 1:-1]
    down = padded[2:, 1:-1]
    left = padded[1:-1, :-2]
    right = padded[1:-1, 2:]

    edges = (
        (segment_int == 1) & (up == 0)
        | (segment_int == 1) & (down == 0)
        | (segment_int == 1) & (left == 0)
        | (segment_int == 1) & (right == 0)
    )
    return float(edges.sum())