feat(compressors): add object scoring and selection for SAM masks

mini-nav/compressors/object_score/features.py

@@ -0,0 +1,211 @@
+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())