From a809803979e7e0a01f48c36a5eb964702c9448f2 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Mon, 30 Mar 2026 16:31:06 +0800 Subject: [PATCH] feat(compressors): add object scoring and selection for SAM masks --- mini-nav/compressors/__init__.py | 14 ++ mini-nav/compressors/object_score/__init__.py | 14 ++ mini-nav/compressors/object_score/config.py | 27 +++ mini-nav/compressors/object_score/features.py | 211 ++++++++++++++++++ mini-nav/compressors/object_score/scorer.py | 129 +++++++++++ mini-nav/compressors/object_score/selector.py | 55 +++++ mini-nav/compressors/pipeline.py | 23 +- mini-nav/tests/test_object_score.py | 113 ++++++++++ 8 files changed, 581 insertions(+), 5 deletions(-) create mode 100644 mini-nav/compressors/object_score/__init__.py create mode 100644 mini-nav/compressors/object_score/config.py create mode 100644 mini-nav/compressors/object_score/features.py create mode 100644 mini-nav/compressors/object_score/scorer.py create mode 100644 mini-nav/compressors/object_score/selector.py create mode 100644 mini-nav/tests/test_object_score.py diff --git a/mini-nav/compressors/__init__.py b/mini-nav/compressors/__init__.py index fef16a9..3e55583 100644 --- a/mini-nav/compressors/__init__.py +++ b/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", diff --git a/mini-nav/compressors/object_score/__init__.py b/mini-nav/compressors/object_score/__init__.py new file mode 100644 index 0000000..8ebcddd --- /dev/null +++ b/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", +] diff --git a/mini-nav/compressors/object_score/config.py b/mini-nav/compressors/object_score/config.py new file mode 100644 index 0000000..f5f3e19 --- /dev/null +++ b/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 diff --git a/mini-nav/compressors/object_score/features.py b/mini-nav/compressors/object_score/features.py new file mode 100644 index 0000000..1bc81db --- /dev/null +++ b/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()) diff --git a/mini-nav/compressors/object_score/scorer.py b/mini-nav/compressors/object_score/scorer.py new file mode 100644 index 0000000..1b57b1f --- /dev/null +++ b/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)) diff --git a/mini-nav/compressors/object_score/selector.py b/mini-nav/compressors/object_score/selector.py new file mode 100644 index 0000000..453a4a9 --- /dev/null +++ b/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) diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index a53eec9..8af8e35 100644 --- a/mini-nav/compressors/pipeline.py +++ b/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 [ - extract_masked_region(image, masks[0]["segment"]) if masks else image - for image, masks in zip(image_list, masks_dataset) - ] + selected_images: list[Image.Image] = [] + 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. diff --git a/mini-nav/tests/test_object_score.py b/mini-nav/tests/test_object_score.py new file mode 100644 index 0000000..1c9121e --- /dev/null +++ b/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())