postprocessing/api.py

"""Public mask-to-contour post-processing functions."""

from __future__ import annotations

import numpy as np

from .geometry import (
    largest_external_contour,
    myocardium_ring_boundaries,
    physical_to_xy,
    render_periodic_bspline,
    render_tension_curve,
    select_adaptive_control_indices,
    select_sparse_spline_controls,
    xy_to_physical,
)
from .metrics import contour_quality_metrics
from .presets import LAX_EDITABLE_CONTOUR_V1
from .types import (
    AdaptiveContourConfig,
    ContourConfig,
    MaskToContourResult,
    MyocardiumContourResult,
    PixelSpacing,
    SplineContourConfig,
    validate_pixel_spacing,
)


def _convert_dense_contour(
    dense_xy: np.ndarray,
    shape: tuple[int, int],
    pixel_spacing: PixelSpacing,
    config: ContourConfig,
) -> MaskToContourResult:
    physical_dense = xy_to_physical(dense_xy, pixel_spacing)
    if isinstance(config, AdaptiveContourConfig):
        indices = select_adaptive_control_indices(physical_dense, config)
        physical_control = physical_dense[indices]
        physical_smooth = render_tension_curve(
            physical_control,
            tension=config.tension,
            samples_per_segment=config.samples_per_segment,
        )
        control_xy = dense_xy[indices].astype(np.float64)
    elif isinstance(config, SplineContourConfig):
        config.validate()
        physical_bspline = render_periodic_bspline(physical_dense, config)
        physical_control, physical_smooth = select_sparse_spline_controls(
            physical_bspline, config
        )
        control_xy = physical_to_xy(physical_control, pixel_spacing)
    else:  # pragma: no cover - protected by the public type and explicit error
        raise TypeError(f"unsupported contour config: {type(config).__name__}")
    smooth_xy = physical_to_xy(physical_smooth, pixel_spacing)
    metrics, rendered_mask = contour_quality_metrics(
        dense_xy,
        smooth_xy,
        shape,
        pixel_spacing,
        control_point_count=len(control_xy),
    )
    return MaskToContourResult(
        dense_contour=dense_xy.astype(np.float64),
        control_points=control_xy,
        smooth_contour=smooth_xy,
        rendered_mask=rendered_mask,
        metrics=metrics,
        pixel_spacing=pixel_spacing,
        preset_name=config.name,
    )


def mask_to_smooth_contour(
    mask: np.ndarray,
    pixel_spacing: PixelSpacing | None = None,
    config: ContourConfig = LAX_EDITABLE_CONTOUR_V1,
) -> MaskToContourResult:
    """Convert the largest mask component to a smooth closed contour."""

    spacing = validate_pixel_spacing(pixel_spacing)
    dense_xy = largest_external_contour(mask)
    return _convert_dense_contour(dense_xy, np.asarray(mask).shape, spacing, config)


def myocardium_mask_to_smooth_contours(
    mask: np.ndarray,
    pixel_spacing: PixelSpacing | None = None,
    config: ContourConfig = LAX_EDITABLE_CONTOUR_V1,
) -> MyocardiumContourResult:
    """Convert a myocardium ring mask into paired endo/epi smooth contours."""

    spacing = validate_pixel_spacing(pixel_spacing)
    array = np.asarray(mask)
    inner, outer = myocardium_ring_boundaries(array)
    endocardium = _convert_dense_contour(inner, array.shape, spacing, config)
    epicardium = _convert_dense_contour(outer, array.shape, spacing, config)
    rendered = np.logical_and(
        epicardium.rendered_mask > 0, endocardium.rendered_mask == 0
    ).astype(np.uint8)
    source = array > 0
    intersection = int(np.logical_and(source, rendered > 0).sum())
    union = int(np.logical_or(source, rendered > 0).sum())
    source_area = int(source.sum())
    rendered_area = int(rendered.sum())
    return MyocardiumContourResult(
        endocardium=endocardium,
        epicardium=epicardium,
        rendered_myocardium_mask=rendered,
        mask_iou=float(intersection / union) if union else 0.0,
        area_change_pct=(
            float(100.0 * (rendered_area - source_area) / source_area)
            if source_area
            else float("inf")
        ),
    )


def cavity_myocardium_masks_to_smooth_contours(
    cavity_mask: np.ndarray,
    myocardium_mask: np.ndarray,
    pixel_spacing: PixelSpacing | None = None,
    config: ContourConfig = LAX_EDITABLE_CONTOUR_V1,
) -> MyocardiumContourResult:
    """Convert cavity and myocardium predictions to endo/epi contours.

    The endocardium is the cavity's largest external boundary.  The
    epicardium is the largest external boundary of ``cavity | myocardium``.
    This supports model outputs where the myocardium class is either a ring or
    an overlapping/filled epicardial region.
    """

    spacing = validate_pixel_spacing(pixel_spacing)
    cavity = np.asarray(cavity_mask)
    myocardium = np.asarray(myocardium_mask)
    if cavity.ndim != 2 or myocardium.ndim != 2:
        raise ValueError("cavity_mask and myocardium_mask must both be 2D")
    if cavity.shape != myocardium.shape:
        raise ValueError(
            "cavity_mask and myocardium_mask must have identical shapes"
        )
    cavity_binary = cavity > 0
    myocardium_binary = myocardium > 0
    outer_binary = np.logical_or(cavity_binary, myocardium_binary)
    endocardium = _convert_dense_contour(
        largest_external_contour(cavity_binary), cavity.shape, spacing, config
    )
    epicardium = _convert_dense_contour(
        largest_external_contour(outer_binary), cavity.shape, spacing, config
    )
    rendered = np.logical_and(
        epicardium.rendered_mask > 0, endocardium.rendered_mask == 0
    ).astype(np.uint8)
    source = np.logical_and(outer_binary, np.logical_not(cavity_binary))
    intersection = int(np.logical_and(source, rendered > 0).sum())
    union = int(np.logical_or(source, rendered > 0).sum())
    source_area = int(source.sum())
    rendered_area = int(rendered.sum())
    return MyocardiumContourResult(
        endocardium=endocardium,
        epicardium=epicardium,
        rendered_myocardium_mask=rendered,
        mask_iou=float(intersection / union) if union else 0.0,
        area_change_pct=(
            float(100.0 * (rendered_area - source_area) / source_area)
            if source_area
            else float("inf")
        ),
    )