小目标分割的"框"靠得住吗？YOLO26x 11% 的匹配分割落在检测框外

"""
Pseudocode for:
"On Detection Reliability for Small Object Segmentation"
arXiv:2610.34432

This file implements the core centroid-containment analysis pipeline
described in Sections 2.1 and 2.3 of the paper. It covers:
  - Reference annotation consistency check (Eq. 1–4)
  - Matching of predicted mask centroids to ground-truth mask centroids
  - Classification of matched centroids as INSIDE, OUTSIDE, or NO_MATCH
"""

from dataclasses import dataclass
from enum import Enum
from typing import List, Tuple, Optional

import numpy as np
from scipy.optimize import linear_sum_assignment


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

@dataclass
class BBox:
    """Axis-aligned bounding box in pixel coordinates (x1, y1, x2, y2)."""
    x1: float
    y1: float
    x2: float
    y2: float

    @property
    def center(self) -> Tuple[float, float]:
        return (self.x1 + self.x2) / 2.0, (self.y1 + self.y2) / 2.0

    def contains(self, point: Tuple[float, float]) -> bool:
        """Check whether a point lies inside the box (boundary inclusive)."""
        px, py = point
        return self.x1 <= px <= self.x2 and self.y1 <= py <= self.y2


@dataclass
class MaskObject:
    """A single connected foreground component with its mask centroid."""
    centroid: Tuple[float, float]
    bbox: BBox  # expanded reference box (5 px margin applied for GT)


class ContainmentOutcome(Enum):
    INSIDE = "inside"
    OUTSIDE = "outside"
    NO_MATCH = "no_match"


# ---------------------------------------------------------------------------
# Eq. (1): centroids and reference box construction
# ---------------------------------------------------------------------------

def compute_mask_centroid(mask: np.ndarray) -> Tuple[float, float]:
    """
    Compute the centroid of a binary mask using standard image moments.

    Args:
        mask: 2-D binary array (foreground = 255 or True).

    Returns:
        (cx, cy) centroid coordinates.
    """
    ys, xs = np.nonzero(mask)
    if len(xs) == 0:
        raise ValueError("Empty mask has no centroid.")
    return float(xs.mean()), float(ys.mean())


def expand_bbox(bbox: BBox, margin: int = 5,
                img_shape: Optional[Tuple[int, int]] = None) -> BBox:
    """Expand a bounding box by `margin` pixels on each side (clipped)."""
    x1 = bbox.x1 - margin
    y1 = bbox.y1 - margin
    x2 = bbox.x2 + margin
    y2 = bbox.y2 + margin

    if img_shape is not None:
        h, w = img_shape[:2]
        x1 = max(0.0, x1)
        y1 = max(0.0, y1)
        x2 = min(float(w - 1), x2)
        y2 = min(float(h - 1), y2)

    return BBox(x1, y1, x2, y2)


# ---------------------------------------------------------------------------
# Eq. (2): optimal one-to-one assignment
# ---------------------------------------------------------------------------

def match_centroids(
    source_points: List[Tuple[float, float]],
    target_points: List[Tuple[float, float]],
) -> List[Tuple[int, int]]:
    """
    Solve the one-to-one assignment problem (Eq. 2) using the
    Jonker–Volgenant algorithm via SciPy's linear_sum_assignment.

    Returns:
        List of (source_index, target_index) matched pairs.
    """
    n = len(source_points)
    m = len(target_points)
    if n == 0 or m == 0:
        return []

    cost = np.zeros((n, m), dtype=float)
    for i, (sx, sy) in enumerate(source_points):
        for j, (tx, ty) in enumerate(target_points):
            cost[i, j] = np.hypot(sx - tx, sy - ty)

    row_ind, col_ind = linear_sum_assignment(cost)
    return list(zip(row_ind.tolist(), col_ind.tolist()))


# ---------------------------------------------------------------------------
# Eq. (3) & (4): reference annotation consistency metrics
# ---------------------------------------------------------------------------

def reference_consistency_metrics(
    gt_mask_objects: List[MaskObject],
    gt_bboxes: List[BBox],
    tolerance: float = 5.0,
) -> Tuple[float, float]:
    """
    Compute centre-coordinate agreement (Eq. 3) and mask-centroid
    containment (Eq. 4) for reference annotations.

    Returns:
        (agreement_pct, containment_pct)
    """
    bbox_centers = [b.center for b in gt_bboxes]
    mask_centroids = [obj.centroid for obj in gt_mask_objects]

    matches = match_centroids(mask_centroids, bbox_centers)
    total = len(matches)
    if total == 0:
        return 0.0, 0.0

    agree = 0
    contained = 0
    for mi, bi in matches:
        mx, my = mask_centroids[mi]
        bx, by = bbox_centers[bi]
        if np.hypot(mx - bx, my - by) <= tolerance:
            agree += 1
        if gt_bboxes[bi].contains((mx, my)):
            contained += 1

    return 100.0 * agree / total, 100.0 * contained / total


# ---------------------------------------------------------------------------
# Section 2.3: predicted-mask centroid containment analysis
# ---------------------------------------------------------------------------

def classify_predicted_centroids(
    pred_mask_centroids: List[Tuple[float, float]],
    gt_mask_centroids: List[Tuple[float, float]],
    pred_bboxes: List[BBox],
) -> Tuple[List[ContainmentOutcome], dict]:
    """
    Core analysis: match predicted mask centroids to GT mask centroids,
    then classify each matched centroid as INSIDE at least one predicted
    bounding box, OUTSIDE all predicted boxes, or NO_MATCH.

    Returns:
        outcomes: per matched prediction, the containment outcome.
        stats:    summary counts and percentages.
    """
    matches = match_centroids(pred_mask_centroids, gt_mask_centroids)

    outcomes: List[ContainmentOutcome] = []
    inside_count = 0
    outside_count = 0

    matched_gt_indices = set()
    for pi, gi in matches:
        matched_gt_indices.add(gi)
        cx, cy = pred_mask_centroids[pi]

        if any(bbox.contains((cx, cy)) for bbox in pred_bboxes):
            outcomes.append(ContainmentOutcome.INSIDE)
            inside_count += 1
        else:
            outcomes.append(ContainmentOutcome.OUTSIDE)
            outside_count += 1

    no_match_count = len(gt_mask_centroids) - len(matched_gt_indices)
    total_gt = len(gt_mask_centroids)

    stats = {
        "total_gt_objects": total_gt,
        "matched_predictions": len(matches),
        "inside": inside_count,
        "outside": outside_count,
        "no_match": no_match_count,
        "inside_pct": 100.0 * inside_count / total_gt if total_gt else 0.0,
        "outside_pct": 100.0 * outside_count / total_gt if total_gt else 0.0,
        "no_match_pct": 100.0 * no_match_count / total_gt if total_gt else 0.0,
    }
    return outcomes, stats


# ---------------------------------------------------------------------------
# End-to-end demonstration (matches the paper's reporting style)
# ---------------------------------------------------------------------------

def run_centroid_containment_analysis(
    gt_mask_objects: List[MaskObject],
    gt_bboxes: List[BBox],
    pred_mask_centroids: List[Tuple[float, float]],
    pred_bboxes: List[BBox],
    tolerance: float = 5.0,
) -> None:
    """
    Full pipeline:
      1. Reference consistency check (Eq. 3–4)
      2. Predicted centroid containment classification (Section 2.3)
      3. Print the same style of summary reported in the paper.
    """
    print("=" * 60)
    print("1. REFERENCE ANNOTATION CONSISTENCY")
    print("=" * 60)

    agree_pct, contain_pct = reference_consistency_metrics(
        gt_mask_objects, gt_bboxes, tolerance
    )
    print(f"Centre-coordinate agreement (≤{tolerance:.0f} px): {agree_pct:.2f}%")
    print(f"Mask-centroid containment in reference boxes: {contain_pct:.2f}%")

    print("\n" + "=" * 60)
    print("2. PREDICTED-MASK CENTROID CONTAINMENT")
    print("=" * 60)

    gt_mask_centroids = [obj.centroid for obj in gt_mask_objects]
    outcomes, stats = classify_predicted_centroids(
        pred_mask_centroids, gt_mask_centroids, pred_bboxes
    )

    print(f"Total GT objects          : {stats['total_gt_objects']}")
    print(f"Matched predictions       : {stats['matched_predictions']}")
    print(f"  INSIDE  at least one box: {stats['inside']:3d}  "
          f"({stats['inside_pct']:.2f}%)")
    print(f"  OUTSIDE all boxes       : {stats['outside']:3d}  "
          f"({stats['outside_pct']:.2f}%)")
    print(f"  NO_MATCH                : {stats['no_match']:3d}  "
          f"({stats['no_match_pct']:.2f}%)")

    print("\nPer-matched-prediction outcomes:")
    for i, outcome in enumerate(outcomes):
        print(f"  Pred {i:03d}: {outcome.value}")


# ---------------------------------------------------------------------------
# Minimal toy example (for quick sanity check)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Toy GT: two small objects
    gt_masks = [
        MaskObject(centroid=(10.0, 10.0), bbox=BBox(5, 5, 15, 15)),
        MaskObject(centroid=(40.0, 40.0), bbox=BBox(35, 35, 45, 45)),
    ]
    gt_boxes = [obj.bbox for obj in gt_masks]

    # Simulated predictions:
    # - pred 0: centroid inside its predicted box (INSIDE)
    # - pred 1: centroid outside its predicted box (OUTSIDE)
    # - GT object 1 receives no matched prediction (NO_MATCH)
    pred_centroids = [(10.5, 10.5), (39.0, 39.0)]
    pred_bboxes = [
        BBox(8, 8, 13, 13),      # contains (10.5, 10.5) → INSIDE
        BBox(50, 50, 60, 60),     # does NOT contain (39, 39) → OUTSIDE
    ]

    run_centroid_containment_analysis(
        gt_mask_objects=gt_masks,
        gt_bboxes=gt_boxes,
        pred_mask_centroids=pred_centroids,
        pred_bboxes=pred_bboxes,
        tolerance=5.0,
    )