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

mini-nav/compressors/__init__.py

@@ -6,6 +6,14 @@ from .common import (
     hash_to_bits,
 )
 from .hash_compressor import HashCompressor, HashLoss, VideoPositiveMask
+from .object_score import (
+    MaskFeatures,
+    MaskScoringConfig,
+    compute_mask_features,
+    rank_masks,
+    score_mask,
+    select_best_mask,
+)
 from .pipeline import HashPipeline, create_pipeline_from_config
 from .train import train
 
@@ -16,6 +24,12 @@ __all__ = [
     "VideoPositiveMask",
     "HashPipeline",
     "create_pipeline_from_config",
+    "MaskFeatures",
+    "MaskScoringConfig",
+    "compute_mask_features",
+    "score_mask",
+    "rank_masks",
+    "select_best_mask",
     "BinarySign",
     "hamming_distance",
     "hamming_similarity",

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)

mini-nav/compressors/pipeline.py

@@ -7,6 +7,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 
+from .object_score import select_best_mask
 from utils import get_device
 from utils.image import extract_masked_region, segment_image, segment_image_dataset
 from utils.model import (
@@ -111,7 +112,10 @@ class HashPipeline(nn.Module):
         if not masks:
             return image
 
-        return extract_masked_region(image, masks[0]["segment"])
+        best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
+        if best_mask is None:
+            return image
+        return extract_masked_region(image, best_mask["segment"])
 
     def _segment_with_sam_dataset(
         self,
@@ -125,10 +129,19 @@ class HashPipeline(nn.Module):
             max_masks=self.sam_max_masks,
             points_per_batch=self.sam_points_per_batch,
         )
-        return [
+        selected_images: list[Image.Image] = []
-            extract_masked_region(image, masks[0]["segment"]) if masks else image
+        for image, masks in zip(image_list, masks_dataset):
-            for image, masks in zip(image_list, masks_dataset)
+            if not masks:
-        ]
+                selected_images.append(image)
+                continue
+
+            best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
+            if best_mask is None:
+                selected_images.append(image)
+                continue
+            selected_images.append(extract_masked_region(image, best_mask["segment"]))
+
+        return selected_images
 
     def _dino_forward(self, image: Image.Image) -> torch.Tensor:
         """Extract DINO tokens from an image.

mini-nav/tests/test_object_score.py

@@ -0,0 +1,113 @@
+import numpy as np
+
+from compressors.object_score import (
+    MaskScoringConfig,
+    compute_mask_features,
+    rank_masks,
+    score_mask,
+    select_best_mask,
+)
+
+
+def _rect_mask(height: int, width: int, x: int, y: int, w: int, h: int) -> np.ndarray:
+    mask = np.zeros((height, width), dtype=bool)
+    mask[y : y + h, x : x + w] = True
+    return mask
+
+
+def test_compute_mask_features_core_metrics() -> None:
+    mask = _rect_mask(height=20, width=20, x=5, y=4, w=6, h=5)
+    mask_dict = {
+        "segment": mask,
+        "area": int(mask.sum()),
+        "bbox": [5, 4, 6, 5],
+        "predicted_iou": 0.8,
+        "stability_score": 0.9,
+    }
+
+    features = compute_mask_features(mask_dict, image_shape=(20, 20))
+
+    assert features.area_ratio == 30 / 400
+    assert features.fill_ratio == 1.0
+    assert features.aspect_ratio == 6 / 5
+    assert features.touch_top is False
+    assert features.touch_left is False
+    assert features.num_components == 1
+    assert features.largest_component_ratio == 1.0
+    assert features.num_holes == 0
+
+
+def test_rank_masks_rejects_extreme_small_and_fragmented_masks() -> None:
+    cfg = MaskScoringConfig(min_area_ratio=0.02)
+    good_mask = _rect_mask(height=30, width=30, x=6, y=6, w=10, h=10)
+
+    fragmented = np.zeros((30, 30), dtype=bool)
+    fragmented[2, 2] = True
+    fragmented[4, 7] = True
+    fragmented[8, 12] = True
+    fragmented[12, 16] = True
+    fragmented[16, 20] = True
+    fragmented[20, 24] = True
+    fragmented[24, 26] = True
+
+    masks = [
+        {"segment": np.zeros((30, 30), dtype=bool), "area": 1, "bbox": [0, 0, 1, 1]},
+        {
+            "segment": fragmented,
+            "area": int(fragmented.sum()),
+            "bbox": [2, 2, 25, 25],
+        },
+        {
+            "segment": good_mask,
+            "area": int(good_mask.sum()),
+            "bbox": [6, 6, 10, 10],
+            "predicted_iou": 0.9,
+            "stability_score": 0.9,
+        },
+    ]
+
+    ranked = rank_masks(masks=masks, image_shape=(30, 30), config=cfg, max_masks=3)
+
+    assert len(ranked) == 1
+    assert ranked[0]["area"] == int(good_mask.sum())
+    assert "mask_score" in ranked[0]
+
+
+def test_score_mask_prefers_stable_reasonable_object() -> None:
+    cfg = MaskScoringConfig()
+
+    candidate = {
+        "segment": _rect_mask(height=100, width=100, x=30, y=20, w=24, h=25),
+        "area": 24 * 25,
+        "bbox": [30, 20, 24, 25],
+        "predicted_iou": 0.92,
+        "stability_score": 0.91,
+    }
+    weak = {
+        "segment": _rect_mask(height=100, width=100, x=0, y=0, w=4, h=60),
+        "area": 4 * 60,
+        "bbox": [0, 0, 4, 60],
+        "predicted_iou": 0.4,
+        "stability_score": 0.3,
+    }
+
+    score_candidate = score_mask(candidate, image_shape=(100, 100), config=cfg)
+    score_weak = score_mask(weak, image_shape=(100, 100), config=cfg)
+
+    assert score_candidate > score_weak
+
+
+def test_select_best_mask_falls_back_to_largest_area_when_all_rejected() -> None:
+    cfg = MaskScoringConfig(min_area_ratio=0.2)
+    tiny = _rect_mask(height=20, width=20, x=1, y=1, w=2, h=2)
+    larger = _rect_mask(height=20, width=20, x=5, y=5, w=4, h=4)
+
+    masks = [
+        {"segment": tiny, "area": int(tiny.sum()), "bbox": [1, 1, 2, 2]},
+        {"segment": larger, "area": int(larger.sum()), "bbox": [5, 5, 4, 4]},
+    ]
+
+    best = select_best_mask(masks=masks, image_shape=(20, 20), config=cfg)
+
+    assert best is not None
+    assert best["area"] == int(larger.sum())