postprocessing/types.py

"""Public data types for model-independent contour post-processing."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np


PixelSpacing = Tuple[float, float]


@dataclass(frozen=True)
class AdaptiveContourConfig:
    """Configuration for sparse, editable closed contours.

    Distances are expressed in millimetres. ``pixel_spacing`` supplied to the
    public API is ordered as ``(row_mm, column_mm)`` while all returned points
    are ordered as ``[x, y]`` in input-image pixel coordinates.
    """

    min_control_points: int = 8
    max_control_points: int = 10
    tolerance_mm: float = 1.5
    min_spacing_mm: float = 4.0
    tension: float = 0.95
    samples_per_segment: int = 20
    name: str = "adaptive"

    def validate(self) -> None:
        if self.min_control_points < 3:
            raise ValueError("min_control_points must be at least 3")
        if self.max_control_points < self.min_control_points:
            raise ValueError("max_control_points must be >= min_control_points")
        if self.tolerance_mm < 0 or self.min_spacing_mm < 0:
            raise ValueError("contour distances must be non-negative")
        if self.tension < 0:
            raise ValueError("tension must be non-negative")
        if self.samples_per_segment < 1:
            raise ValueError("samples_per_segment must be positive")


@dataclass(frozen=True)
class SplineContourConfig:
    """Configuration for historical B-spline smoothing and simplification.

    The dense mask boundary is first smoothed to ``n_points`` samples.  Those
    samples are reduced to the first requested control-point count whose
    zero-smoothing reconstruction reaches ``control_point_iou_threshold``;
    otherwise the best candidate is retained.
    """

    smoothing: float = 40.0
    n_points: int = 100
    control_point_counts: Tuple[int, ...] = (6, 8, 10)
    control_point_iou_threshold: float = 0.985
    reconstruction_smoothing: float = 0.0
    simplification_method: str = "curvature"
    name: str = "periodic_bspline"

    def validate(self) -> None:
        if self.smoothing < 0:
            raise ValueError("smoothing must be non-negative")
        if self.n_points < 4:
            raise ValueError("n_points must be at least 4")
        if not self.control_point_counts:
            raise ValueError("control_point_counts must not be empty")
        if any(count < 3 for count in self.control_point_counts):
            raise ValueError("all control-point counts must be at least 3")
        if tuple(sorted(set(self.control_point_counts))) != self.control_point_counts:
            raise ValueError("control_point_counts must be unique and increasing")
        if not 0.0 <= self.control_point_iou_threshold <= 1.0:
            raise ValueError("control_point_iou_threshold must be between 0 and 1")
        if self.reconstruction_smoothing < 0:
            raise ValueError("reconstruction_smoothing must be non-negative")
        if self.simplification_method != "curvature":
            raise ValueError("only curvature simplification is supported")


ContourConfig = Union[AdaptiveContourConfig, SplineContourConfig]


@dataclass(frozen=True)
class ContourQualityMetrics:
    """Agreement between the dense mask boundary and rendered smooth curve."""

    control_point_count: int
    boundary_mean_mm: float
    boundary_p95_mm: float
    boundary_max_mm: float
    mask_iou: float
    source_area_px: int
    rendered_area_px: int
    area_change_pct: float
    roughness_ratio: float

    def as_dict(self) -> dict[str, float | int]:
        return {
            "control_point_count": self.control_point_count,
            "boundary_mean_mm": self.boundary_mean_mm,
            "boundary_p95_mm": self.boundary_p95_mm,
            "boundary_max_mm": self.boundary_max_mm,
            "mask_iou": self.mask_iou,
            "source_area_px": self.source_area_px,
            "rendered_area_px": self.rendered_area_px,
            "area_change_pct": self.area_change_pct,
            "roughness_ratio": self.roughness_ratio,
        }


@dataclass(frozen=True)
class MaskToContourResult:
    """Dense, sparse, and rendered representations of one mask boundary."""

    dense_contour: np.ndarray
    control_points: np.ndarray
    smooth_contour: np.ndarray
    rendered_mask: np.ndarray
    metrics: ContourQualityMetrics
    pixel_spacing: PixelSpacing
    preset_name: str


@dataclass(frozen=True)
class MyocardiumContourResult:
    """Paired inner/endocardial and outer/epicardial ring boundaries."""

    endocardium: MaskToContourResult
    epicardium: MaskToContourResult
    rendered_myocardium_mask: np.ndarray
    mask_iou: float
    area_change_pct: float


def validate_pixel_spacing(pixel_spacing: Optional[PixelSpacing]) -> PixelSpacing:
    spacing = (1.0, 1.0) if pixel_spacing is None else tuple(float(v) for v in pixel_spacing)
    if len(spacing) != 2 or not all(np.isfinite(v) and v > 0 for v in spacing):
        raise ValueError("pixel_spacing must contain two positive finite values")
    return spacing[0], spacing[1]