cv_ops/morphology.py

"""Morphological image operations."""

from __future__ import annotations

import cv2
import numpy as np

from filters.builtin import ensure_rgb


SHAPES = {"rect": cv2.MORPH_RECT, "ellipse": cv2.MORPH_ELLIPSE, "cross": cv2.MORPH_CROSS}
OPS = {
    "Erosion": cv2.MORPH_ERODE,
    "Dilation": cv2.MORPH_DILATE,
    "Opening": cv2.MORPH_OPEN,
    "Closing": cv2.MORPH_CLOSE,
    "Gradient": cv2.MORPH_GRADIENT,
    "Top-Hat": cv2.MORPH_TOPHAT,
    "Black-Hat": cv2.MORPH_BLACKHAT,
}


def make_kernel(shape: str = "rect", size: int = 5) -> np.ndarray:
    size = max(1, int(size))
    if size % 2 == 0:
        size += 1
    return cv2.getStructuringElement(SHAPES.get(shape, cv2.MORPH_RECT), (size, size))


def threshold_image(image: np.ndarray, method: str = "Otsu", threshold: int = 128) -> np.ndarray:
    gray = cv2.cvtColor(ensure_rgb(image), cv2.COLOR_RGB2GRAY)
    if method == "Otsu":
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    else:
        _, binary = cv2.threshold(gray, int(threshold), 255, cv2.THRESH_BINARY)
    return binary


def apply_morphology(image: np.ndarray, operation: str = "Opening", shape: str = "rect", size: int = 5, iterations: int = 1, threshold_method: str = "Otsu", threshold: int = 128) -> tuple[np.ndarray, np.ndarray]:
    binary = threshold_image(image, threshold_method, threshold)
    kernel = make_kernel(shape, size)
    op = OPS.get(operation, cv2.MORPH_OPEN)
    if op == cv2.MORPH_ERODE:
        result = cv2.erode(binary, kernel, iterations=int(iterations))
    elif op == cv2.MORPH_DILATE:
        result = cv2.dilate(binary, kernel, iterations=int(iterations))
    else:
        result = cv2.morphologyEx(binary, op, kernel, iterations=int(iterations))
    return cv2.cvtColor(result, cv2.COLOR_GRAY2RGB), kernel