PostProcessingVerification.py

from __future__ import annotations

from dataclasses import dataclass
import cv2
import numpy as np


@dataclass
class VerifyParams:
    tooth_close_k: int = 3
    tooth_open_k: int = 0
    fill_holes_flag: bool = True

    min_component_area: int = 10
    min_label_pixels: int = 6

    min_roi_overlap_pixels: int = 6
    min_roi_overlap_frac: float = 0.03

    max_exclusion_overlap_frac: float = 0.45

    bridge_close_k: int = 3
    min_seed_overlap_pixels: int = 8
    min_seed_overlap_frac: float = 0.05
    keep_only_components_touching_original: bool = True

    side_far_border_frac: float = 0.10
    side_far_border_max_width_frac: float = 0.18
    side_far_border_max_roi_frac: float = 0.08

    keep_only_best_component_per_label: bool = True


def _fill_holes(mask: np.ndarray) -> np.ndarray:
    mask_u8 = (mask > 0).astype(np.uint8) * 255
    h, w = mask_u8.shape
    flood = mask_u8.copy()
    flood_mask = np.zeros((h + 2, w + 2), np.uint8)
    cv2.floodFill(flood, flood_mask, (0, 0), 255)
    holes = cv2.bitwise_not(flood)
    out = cv2.bitwise_or(mask_u8, holes)
    return (out > 0).astype(np.uint8)


def _largest_component(mask: np.ndarray) -> np.ndarray:
    mask = (mask > 0).astype(np.uint8)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8)
    if n <= 1:
        return mask
    idx = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
    return (labels == idx).astype(np.uint8)


def _clean_binary_piece(mask: np.ndarray, params: VerifyParams) -> np.ndarray:
    out = (mask > 0).astype(np.uint8)

    # IMPORTANT:
    # apply morphology only on this binary label piece,
    # never on the full integer-labeled mask
    if params.tooth_close_k > 1:
        ker = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE,
            (params.tooth_close_k, params.tooth_close_k),
        )
        out = cv2.morphologyEx(out, cv2.MORPH_CLOSE, ker)

    if params.tooth_open_k > 1:
        ker = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE,
            (params.tooth_open_k, params.tooth_open_k),
        )
        out = cv2.morphologyEx(out, cv2.MORPH_OPEN, ker)

    if params.fill_holes_flag:
        out = _fill_holes(out)

    return (out > 0).astype(np.uint8)


def _seed_overlap_ok(
    piece: np.ndarray,
    original_seed: np.ndarray,
    params: VerifyParams,
) -> bool:
    overlap = int(((piece > 0) & (original_seed > 0)).sum())
    area = int((piece > 0).sum())

    if overlap < params.min_seed_overlap_pixels:
        return False

    if overlap / max(area, 1) < params.min_seed_overlap_frac:
        return False

    return True


def _is_tiny_far_side_artifact(
    piece: np.ndarray,
    roi: np.ndarray,
    exc_overlap: int,
    area: int,
    view: str | None,
    params: VerifyParams,
) -> bool:
    if view not in ("left", "right"):
        return False

    h, w = piece.shape
    ys, xs = np.where(piece > 0)
    if len(xs) == 0:
        return True

    x_min, x_max = int(xs.min()), int(xs.max())
    width_frac = (x_max - x_min + 1) / float(max(w, 1))

    border_w = int(params.side_far_border_frac * w)
    far_side = x_min > (w - border_w) if view == "left" else x_max < border_w

    roi_area = int((roi > 0).sum())
    roi_side_frac = area / max(roi_area, 1)
    exc_frac = exc_overlap / max(area, 1)

    tiny_far_artifact = (
        far_side
        and width_frac < 0.07
        and roi_side_frac < 0.008
        and exc_frac > 0.16
    )

    return tiny_far_artifact


def postprocess_mask(
    raw_mask: np.ndarray,
    roi_mask: np.ndarray,
    exclusion_mask: np.ndarray,
    params: VerifyParams,
    view: str | None = None,
) -> np.ndarray:
    raw = np.asarray(raw_mask).astype(np.uint8)
    roi = (np.asarray(roi_mask) > 0).astype(np.uint8)
    exc = (np.asarray(exclusion_mask) > 0).astype(np.uint8)

    h, w = raw.shape
    final = np.zeros((h, w), dtype=np.uint8)

    labels = [int(v) for v in np.unique(raw) if int(v) != 0]

    for lab in labels:
        original_seed = (raw == lab).astype(np.uint8)
        if int(original_seed.sum()) < params.min_label_pixels:
            continue

        # keep only inside ROI from the start
        piece0 = ((original_seed > 0) & (roi > 0)).astype(np.uint8)
        if piece0.sum() == 0:
            continue

        # clean only this label separately
        piece_clean = _clean_binary_piece(piece0, params)

        # optional: reconnect tiny gaps within the SAME label only
        if params.bridge_close_k > 1:
            ker = cv2.getStructuringElement(
                cv2.MORPH_ELLIPSE,
                (params.bridge_close_k, params.bridge_close_k),
            )
            piece_bridge = cv2.morphologyEx(piece_clean, cv2.MORPH_CLOSE, ker)

            # keep only if still anchored to original seed
            if not params.keep_only_components_touching_original or _seed_overlap_ok(
                piece_bridge, original_seed, params
            ):
                piece_clean = piece_bridge

        n, cc, stats, _ = cv2.connectedComponentsWithStats(piece_clean, 8)

        for i in range(1, n):
            piece = (cc == i).astype(np.uint8)
            area = int(stats[i, cv2.CC_STAT_AREA])

            if area < params.min_component_area:
                continue
            if area < params.min_label_pixels:
                continue

            roi_overlap = int(((piece > 0) & (roi > 0)).sum())
            if roi_overlap < params.min_roi_overlap_pixels:
                continue
            if roi_overlap / max(area, 1) < params.min_roi_overlap_frac:
                continue

            exc_overlap = int(((piece > 0) & (exc > 0)).sum())
            if exc_overlap / max(area, 1) > params.max_exclusion_overlap_frac:
                continue

            if params.keep_only_components_touching_original:
                if not _seed_overlap_ok(piece, original_seed, params):
                    continue

            if _is_tiny_far_side_artifact(piece, roi, exc_overlap, area, view, params):
                continue

            # final write for this label
            final[piece > 0] = lab

    return final.astype(np.uint8)