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

mini-nav/compressors/object_score/__init__.py

@@ -0,0 +1,14 @@
+from .config import MaskScoringConfig
+from .features import MaskFeatures, compute_mask_features
+from .scorer import score_mask, should_reject_mask
+from .selector import rank_masks, select_best_mask
+
+__all__ = [
+    "MaskFeatures",
+    "MaskScoringConfig",
+    "compute_mask_features",
+    "should_reject_mask",
+    "score_mask",
+    "rank_masks",
+    "select_best_mask",
+]

mini-nav/compressors/object_score/config.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class MaskScoringConfig:
+    min_area_ratio: float = 0.003
+    max_area_ratio: float = 0.70
+    max_aspect_ratio: float = 6.0
+    min_fill_ratio_hard: float = 0.08
+    max_components: int = 6
+    min_largest_component_ratio: float = 0.60
+    reject_edge_touch_count: int = 4
+    reject_large_edge_touch_count: int = 3
+    reject_large_edge_area_ratio: float = 0.25
+
+    area_score_low: float = 0.02
+    area_score_high: float = 0.35
+    fill_score_low: float = 0.25
+    fill_score_high: float = 0.90
+    soft_aspect_ratio: float = 2.5
+    hole_penalty_step: float = 0.05
+    max_hole_penalty_count: int = 5
+
+    weight_area: float = 0.35
+    weight_shape: float = 0.25
+    weight_fragment: float = 0.25
+    weight_boundary: float = 0.15

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

mini-nav/compressors/object_score/scorer.py

@@ -0,0 +1,129 @@
+from typing import Any
+
+from .config import MaskScoringConfig
+from .features import MaskFeatures, compute_mask_features
+
+
+def should_reject_mask(features: MaskFeatures, config: MaskScoringConfig) -> bool:
+    if features.area_ratio < config.min_area_ratio:
+        return True
+    if features.area_ratio > config.max_area_ratio:
+        return True
+
+    aspect = max(features.aspect_ratio, 1.0 / max(features.aspect_ratio, 1e-6))
+    if aspect > config.max_aspect_ratio:
+        return True
+
+    if features.fill_ratio < config.min_fill_ratio_hard:
+        return True
+
+    if (
+        features.num_components > config.max_components
+        and features.largest_component_ratio < config.min_largest_component_ratio
+    ):
+        return True
+
+    if features.touch_edge_count >= config.reject_edge_touch_count:
+        return True
+    if (
+        features.touch_edge_count >= config.reject_large_edge_touch_count
+        and features.area_ratio > config.reject_large_edge_area_ratio
+    ):
+        return True
+
+    return False
+
+
+def score_mask(
+    mask_dict: dict[str, Any],
+    image_shape: tuple[int, int],
+    config: MaskScoringConfig | None = None,
+) -> float:
+    cfg = config or MaskScoringConfig()
+    features = compute_mask_features(mask_dict, image_shape=image_shape)
+    return _score_from_features(features, cfg)
+
+
+def _score_from_features(features: MaskFeatures, config: MaskScoringConfig) -> float:
+    area = _score_area(features.area_ratio, config)
+    shape = _score_shape(features.aspect_ratio, features.fill_ratio, config)
+    fragment = _score_fragment(
+        num_components=features.num_components,
+        largest_component_ratio=features.largest_component_ratio,
+        num_holes=features.num_holes,
+        config=config,
+    )
+    boundary = _score_boundary(features)
+
+    return (
+        config.weight_area * area
+        + config.weight_shape * shape
+        + config.weight_fragment * fragment
+        + config.weight_boundary * boundary
+    )
+
+
+def _score_area(area_ratio: float, config: MaskScoringConfig) -> float:
+    if area_ratio <= 0:
+        return 0.0
+    if area_ratio < config.area_score_low:
+        return area_ratio / config.area_score_low
+    if area_ratio <= config.area_score_high:
+        return 1.0
+    if area_ratio >= config.max_area_ratio:
+        return 0.0
+    width = max(1e-6, config.max_area_ratio - config.area_score_high)
+    return max(0.0, 1.0 - (area_ratio - config.area_score_high) / width)
+
+
+def _score_shape(
+    aspect_ratio: float,
+    fill_ratio: float,
+    config: MaskScoringConfig,
+) -> float:
+    aspect = max(aspect_ratio, 1.0 / max(aspect_ratio, 1e-6))
+    if aspect <= config.soft_aspect_ratio:
+        aspect_score = 1.0
+    elif aspect >= config.max_aspect_ratio:
+        aspect_score = 0.0
+    else:
+        span = max(1e-6, config.max_aspect_ratio - config.soft_aspect_ratio)
+        aspect_score = 1.0 - (aspect - config.soft_aspect_ratio) / span
+
+    if fill_ratio <= 0:
+        fill_score = 0.0
+    elif fill_ratio < config.fill_score_low:
+        fill_score = fill_ratio / config.fill_score_low
+    elif fill_ratio <= config.fill_score_high:
+        fill_score = 1.0
+    elif fill_ratio >= 1.0:
+        fill_score = 0.0
+    else:
+        span = max(1e-6, 1.0 - config.fill_score_high)
+        fill_score = 1.0 - (fill_ratio - config.fill_score_high) / span
+
+    return 0.5 * aspect_score + 0.5 * fill_score
+
+
+def _score_fragment(
+    num_components: int,
+    largest_component_ratio: float,
+    num_holes: int,
+    config: MaskScoringConfig,
+) -> float:
+    comp_score = max(0.0, 1.0 - 0.15 * max(0, num_components - 1))
+    main_score = max(0.0, min(1.0, largest_component_ratio))
+    hole_penalty = (
+        min(max(0, num_holes), config.max_hole_penalty_count) * config.hole_penalty_step
+    )
+    return max(0.0, 0.5 * comp_score + 0.5 * main_score - hole_penalty)
+
+
+def _score_boundary(features: MaskFeatures) -> float:
+    sam_score = 0.5
+    if features.predicted_iou is not None and features.stability_score is not None:
+        sam_score = 0.5 * features.predicted_iou + 0.5 * features.stability_score
+
+    edge_penalty = 0.05 * features.touch_edge_count
+    compact = max(0.0, min(1.0, features.compactness))
+    return max(0.0, min(1.0, 0.7 * sam_score + 0.3 * compact - edge_penalty))

mini-nav/compressors/object_score/selector.py

@@ -0,0 +1,55 @@
+from typing import Any
+
+from .config import MaskScoringConfig
+from .features import compute_mask_features
+from .scorer import score_mask, should_reject_mask
+
+
+def rank_masks(
+    masks: list[dict[str, Any]],
+    image_shape: tuple[int, int],
+    config: MaskScoringConfig | None = None,
+    max_masks: int | None = None,
+) -> list[dict[str, Any]]:
+    cfg = config or MaskScoringConfig()
+    kept: list[dict[str, Any]] = []
+
+    for mask in masks:
+        features = compute_mask_features(mask, image_shape=image_shape)
+        if should_reject_mask(features, cfg):
+            continue
+        score = score_mask(mask, image_shape=image_shape, config=cfg)
+        enriched = dict(mask)
+        enriched["mask_score"] = float(score)
+        kept.append(enriched)
+
+    kept.sort(
+        key=lambda m: (
+            float(m.get("mask_score", 0.0)),
+            int(m.get("area", 0)),
+        ),
+        reverse=True,
+    )
+
+    if max_masks is None:
+        return kept
+    if max_masks <= 0:
+        return []
+    return kept[:max_masks]
+
+
+def select_best_mask(
+    masks: list[dict[str, Any]],
+    image_shape: tuple[int, int],
+    config: MaskScoringConfig | None = None,
+) -> dict[str, Any] | None:
+    if not masks:
+        return None
+
+    ranked = rank_masks(
+        masks=masks, image_shape=image_shape, config=config, max_masks=1
+    )
+    if ranked:
+        return ranked[0]
+
+    return max(masks, key=lambda m: int(m.get("area", 0)), default=None)