ihc_qc.py

"""QC visualization for spheroid segmentation and DAB quantification."""

from pathlib import Path

import cv2
import numpy as np


DEFAULT_QC_MAX_DIMENSION = 1200
DEFAULT_DAB_VMAX = 0.3


def _resize(
    image: np.ndarray,
    max_dimension: int,
    *,
    nearest: bool = False,
) -> np.ndarray:
    height, width = image.shape[:2]
    scale = min(1.0, max_dimension / max(height, width))
    if scale == 1:
        return image.copy()

    size = (round(width * scale), round(height * scale))
    interpolation = cv2.INTER_NEAREST if nearest else cv2.INTER_AREA
    return cv2.resize(image, size, interpolation=interpolation)


def _draw_spheroids(
    image: np.ndarray,
    labels: np.ndarray,
    boundary_spheroid_ids: set[int],
) -> np.ndarray:
    output = image.copy()

    for spheroid_id in np.unique(labels):
        if spheroid_id == 0:
            continue

        spheroid = labels == spheroid_id
        contours, _ = cv2.findContours(
            spheroid.astype(np.uint8),
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE,
        )
        outline_color = (
            (255, 165, 0)
            if spheroid_id in boundary_spheroid_ids
            else (0, 255, 255)
        )
        cv2.drawContours(output, contours, -1, outline_color, 3)

        y, x = np.nonzero(spheroid)
        center = (round(x.mean()), round(y.mean()))
        text = str(spheroid_id)
        cv2.circle(output, center, 18, (0, 0, 0), cv2.FILLED)
        cv2.putText(
            output,
            text,
            (center[0] - 7 * len(text), center[1] + 7),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.65,
            (255, 255, 255),
            2,
            cv2.LINE_AA,
        )

    return output


def _draw_debris_outlines(
    image: np.ndarray,
    debris_mask: np.ndarray,
) -> np.ndarray:
    output = image.copy()
    contours, _ = cv2.findContours(
        debris_mask.astype(np.uint8),
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE,
    )
    cv2.drawContours(output, contours, -1, (255, 0, 255), 1)
    return output


def _add_title(image: np.ndarray, title: str) -> np.ndarray:
    titled = cv2.copyMakeBorder(
        image,
        54,
        0,
        0,
        0,
        cv2.BORDER_CONSTANT,
        value=(28, 28, 28),
    )
    cv2.putText(
        titled,
        title,
        (18, 36),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        (255, 255, 255),
        2,
        cv2.LINE_AA,
    )
    return titled


def create_qc_image(
    preview: np.ndarray,
    labels: np.ndarray,
    debris_mask: np.ndarray,
    dab: np.ndarray,
    *,
    boundary_spheroid_ids: set[int] | None = None,
    positive_threshold: float | None = None,
    dab_vmax: float = DEFAULT_DAB_VMAX,
    max_dimension: int = DEFAULT_QC_MAX_DIMENSION,
) -> np.ndarray:
    """Return an RGB QC image with segmentation, debris, and DAB panels."""
    if preview.shape[:2] != labels.shape:
        raise ValueError("preview and labels must have the same height and width.")
    if labels.shape != debris_mask.shape or labels.shape != dab.shape:
        raise ValueError("labels, debris_mask, and dab must have the same shape.")
    if max_dimension <= 0:
        raise ValueError("max_dimension must be positive.")
    if positive_threshold is not None and not np.isfinite(positive_threshold):
        raise ValueError("positive_threshold must be finite.")
    if not np.isfinite(dab_vmax) or dab_vmax <= 0:
        raise ValueError("dab_vmax must be finite and positive.")

    preview_small = _resize(preview, max_dimension)
    labels_small = _resize(labels, max_dimension, nearest=True)
    debris_small = _resize(
        (debris_mask > 0).astype(np.uint8),
        max_dimension,
        nearest=True,
    ).astype(bool)
    dab_small = _resize(dab, max_dimension)
    if boundary_spheroid_ids is None:
        boundary_spheroid_ids = {
            int(spheroid_id)
            for spheroid_id in np.unique(
                np.concatenate(
                    (
                        labels[0, :],
                        labels[-1, :],
                        labels[:, 0],
                        labels[:, -1],
                    )
                )
            )
            if spheroid_id != 0
        }

    debris_outline = _draw_debris_outlines(preview_small, debris_small)
    debris_outline = _draw_spheroids(
        debris_outline,
        labels_small,
        boundary_spheroid_ids,
    )
    debris_outline = _add_title(
        debris_outline,
        "cyan: within edge tolerance | orange: exceeds tolerance | magenta: excluded",
    )

    valid = (labels_small > 0) & ~debris_small
    heatmap = np.full((*dab_small.shape, 3), 235, dtype=np.uint8)
    if valid.any():
        scaled_dab = np.clip(dab_small / dab_vmax, 0, 1)
        colors = cv2.applyColorMap(
            (scaled_dab * 255).astype(np.uint8),
            cv2.COLORMAP_INFERNO,
        )
        colors = cv2.cvtColor(colors, cv2.COLOR_BGR2RGB)
        heatmap[valid] = colors[valid]

    heatmap[debris_small & (labels_small > 0)] = (255, 0, 255)
    heatmap = _draw_spheroids(
        heatmap,
        labels_small,
        boundary_spheroid_ids,
    )
    heatmap = _add_title(
        heatmap,
        (
            f"DAB signal (fixed 0-{dab_vmax:g}) | "
            "orange: exceeds edge tolerance"
        ),
    )

    panels = [debris_outline, heatmap]
    if positive_threshold is not None:
        positive = (
            (labels > 0)
            & (debris_mask == 0)
            & (dab >= positive_threshold)
        )
        positive_small = _resize(
            positive.astype(np.uint8),
            max_dimension,
            nearest=True,
        ).astype(bool)

        binary_positive = np.zeros((*labels_small.shape, 3), dtype=np.uint8)
        binary_positive[positive_small] = (255, 255, 255)
        binary_positive = _add_title(
            binary_positive,
            f"white: DAB-positive | threshold >= {positive_threshold:.4g}",
        )
        panels.append(binary_positive)

    return np.concatenate(panels, axis=1)


def save_qc_image(
    output_path: str | Path,
    preview: np.ndarray,
    labels: np.ndarray,
    debris_mask: np.ndarray,
    dab: np.ndarray,
    *,
    boundary_spheroid_ids: set[int] | None = None,
    positive_threshold: float | None = None,
    dab_vmax: float = DEFAULT_DAB_VMAX,
    max_dimension: int = DEFAULT_QC_MAX_DIMENSION,
) -> Path:
    """Create and save the RGB QC image, returning its output path."""
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    qc_image = create_qc_image(
        preview,
        labels,
        debris_mask,
        dab,
        boundary_spheroid_ids=boundary_spheroid_ids,
        positive_threshold=positive_threshold,
        dab_vmax=dab_vmax,
        max_dimension=max_dimension,
    )
    saved = cv2.imwrite(
        str(output_path),
        cv2.cvtColor(qc_image, cv2.COLOR_RGB2BGR),
    )
    if not saved:
        raise OSError(f"Could not write QC image: {output_path}")
    return output_path