Newer
Older
RobotCar / src / pi / vision / morphology.py
"""
morphology
二値画像の形態学的処理ユーティリティモジュール
"""

import cv2
import numpy as np


def apply_iso_closing(
    binary: np.ndarray, size: int,
) -> np.ndarray:
    """等方クロージングで穴を埋める

    Args:
        binary: 二値画像
        size: カーネルサイズ

    Returns:
        クロージング後の二値画像
    """
    if size < 3:
        return binary
    k = size | 1
    kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (k, k),
    )
    return cv2.morphologyEx(
        binary, cv2.MORPH_CLOSE, kernel,
    )


def apply_staged_closing(
    binary: np.ndarray,
    small_size: int,
    min_area: int,
    large_size: int,
) -> np.ndarray:
    """段階クロージング: 小穴埋め → 孤立除去 → 大穴埋め

    Args:
        binary: 二値画像
        small_size: 第1段クロージングのカーネルサイズ
        min_area: 孤立領域除去の最小面積(0 で無効)
        large_size: 第2段クロージングのカーネルサイズ(0 で無効)

    Returns:
        処理後の二値画像
    """
    result = apply_iso_closing(binary, small_size)

    if min_area > 0:
        contours, _ = cv2.findContours(
            result, cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE,
        )
        mask = np.zeros_like(result)
        for cnt in contours:
            if cv2.contourArea(cnt) >= min_area:
                cv2.drawContours(
                    mask, [cnt], -1, 255, -1,
                )
        result = mask

    result = apply_iso_closing(result, large_size)

    return result


def apply_width_filter(
    binary: np.ndarray,
    width_near: int,
    width_far: int,
    tolerance: float,
) -> np.ndarray:
    """透視補正付き幅フィルタで広がりすぎた行を除外する

    Args:
        binary: 二値画像
        width_near: 画像下端での期待線幅(px)
        width_far: 画像上端での期待線幅(px)
        tolerance: 上限倍率

    Returns:
        幅フィルタ適用後の二値画像
    """
    result = binary.copy()
    h = binary.shape[0]
    denom = max(h - 1, 1)

    for y_local in range(h):
        xs = np.where(binary[y_local] > 0)[0]
        if len(xs) == 0:
            continue
        t = (h - 1 - y_local) / denom
        expected = float(width_far) + (
            float(width_near) - float(width_far)
        ) * t
        max_w = expected * tolerance
        actual_w = int(xs[-1]) - int(xs[0]) + 1
        if actual_w > max_w:
            result[y_local] = 0

    return result


def apply_dist_mask(
    binary: np.ndarray, thresh: float,
) -> np.ndarray:
    """距離変換で中心部のみを残す

    Args:
        binary: 二値画像
        thresh: 距離の閾値(ピクセル)

    Returns:
        中心部のみの二値画像
    """
    if thresh <= 0:
        return binary
    dist = cv2.distanceTransform(
        binary, cv2.DIST_L2, 5,
    )
    _, mask = cv2.threshold(
        dist, thresh, 255, cv2.THRESH_BINARY,
    )
    return mask.astype(np.uint8)