【Opencv】形态学运算

以前学习opencv-python的时候，为了锻炼一下python的code习惯，重新实现封装了一下腐蚀、膨胀、开运算、闭运算、黑帽运算等形态学计算，学习分享一下，哈哈。

import cv2
from typing import Optional, Tuple


def _morphology_base(operation,
                     src,
                     kernel,
                     anchor: Optional[Tuple] = (-1, -1),
                     iterations: Optional[int] = 1,
                     borderType: Optional[str] = 'CONSTANT'):
    _border_type_list = ['CONSTANT',
                         'DEFAULT',
                         'ISOLATED',
                         'REFLECT',
                         'REFLECT101',
                         'REFLECT_101,'
                         'REPLICATE',
                         'TRANSPARENT',
                         'WRAP']

    if borderType not in _border_type_list:
        raise ValueError(f'The boederType value is incorrect! Please select in {_border_type_list}')
    if operation not in ['erode', 'dilate']:
        raise ValueError('The operate value is incorrect! Please select in "erode", "dilate"')

    if borderType == 'CONSTANT':
        border_type = cv2.BORDER_CONSTANT
    elif borderType == 'DEFAULT':
        border_type = cv2.BORDER_DEFAULT
    elif borderType == 'ISOLATED':
        border_type = cv2.BORDER_ISOLATED
    elif borderType == 'REFLECT':
        border_type = cv2.BORDER_REFLECT
    elif borderType == 'REFLECT101':
        border_type = cv2.BORDER_REFLECT101
    elif borderType == 'REFLECT_101':
        border_type = cv2.BORDER_REFLECT_101
    elif borderType == 'REPLICATE':
        border_type = cv2.BORDER_REPLICATE
    elif borderType == 'TRANSPARENT':
        border_type = cv2.BORDER_TRANSPARENT
    elif borderType == 'WRAP':
        border_type = cv2.BORDER_WRAP
    else:
        border_type = cv2.BORDER_CONSTANT

    if operation == 'erode':
        return cv2.erode(src=src, kernel=kernel, anchor=anchor, iterations=iterations, borderType=border_type)
    else:
        return cv2.dilate(src=src, kernel=kernel, anchor=anchor, iterations=iterations, borderType=border_type)


def _erode(src,
           kernel,
           anchor: Optional[Tuple] = (-1, -1),
           iterations: Optional[int] = 1,
           borderType: Optional[str] = 'CONSTANT'
           ):
    return _morphology_base(operation='erode', src=src, kernel=kernel,
                            anchor=anchor, iterations=iterations, borderType=borderType)


def _dilate(src,
            kernel,
            anchor: Optional[Tuple] = (-1, -1),
            iterations: Optional[int] = 1,
            borderType: Optional[str] = 'CONSTANT'
            ):
    return _morphology_base(operation='dilate', src=src, kernel=kernel,
                            anchor=anchor, iterations=iterations, borderType=borderType)


def _open(src,
          kernel,
          anchor: Optional[Tuple] = (-1, -1),
          iterations: Optional[int] = 1,
          borderType: Optional[str] = 'CONSTANT'
          ):
    erosion = _erode(src, kernel, anchor, iterations, borderType)
    open_result = _dilate(erosion, kernel, anchor, iterations, borderType)
    return open_result


def _close(src,
           kernel,
           anchor: Optional[Tuple] = (-1, -1),
           iterations: Optional[int] = 1,
           borderType: Optional[str] = 'CONSTANT'
           ):
    dilation = _dilate(src, kernel, anchor, iterations, borderType)
    close_result = _erode(dilation, kernel, anchor, iterations, borderType)
    return close_result


def _gradient(src,
              kernel,
              anchor: Optional[Tuple] = (-1, -1),
              iterations: Optional[int] = 1,
              borderType: Optional[str] = 'CONSTANT'
              ):
    erosion = _erode(src, kernel, anchor, iterations, borderType)
    dilation = _dilate(src, kernel, anchor, iterations, borderType)
    gradient_result = cv2.addWeighted(dilation, 1, erosion, -1, gamma=0)
    return gradient_result


def _tophat(src,
            kernel,
            anchor: Optional[Tuple] = (-1, -1),
            iterations: Optional[int] = 1,
            borderType: Optional[str] = 'CONSTANT'
            ):
    open_result = _open(src, kernel, anchor, iterations, borderType)
    top_hat_result = cv2.addWeighted(src, 1, open_result, -1, gamma=0)
    return top_hat_result


def _blackhat(src,
              kernel,
              anchor: Optional[Tuple] = (-1, -1),
              iterations: Optional[int] = 1,
              borderType: Optional[str] = 'CONSTANT'
              ):
    close_result = _close(src, kernel, anchor, iterations, borderType)
    black_hat_result = cv2.addWeighted(close_result, 1, src, -1, gamma=0)
    return black_hat_result


def morphology_operation(operation,
                         src,
                         kernel,
                         anchor: Optional[Tuple] = (-1, -1),
                         iterations: Optional[int] = 1,
                         borderType: Optional[str] = 'CONSTANT'):
    """
     形态学运算，支持腐蚀，膨胀，开运算，闭运算，形态学梯度运算，顶帽运算以及黑帽运算。

    :param operation: 形态学运算，支持腐蚀，膨胀，开运算，闭运算，形态学梯度运算，顶帽运算以及黑帽运算。
    :param src: 待操作的原始图像，通道数可为任意数，但深度必须属于【CV_8U CV_16U CV_16S CV_32F CV 64F】之一。
    :param kernel: 操作所采用的结构体，可自定义生成，也可使用cv2.getStructuringElement()生成。
    :param anchor: element结构中的锚点位置，默认核的中心。
    :param iterations: 操作的迭代次数，默认为1.
    :param borderType: 边界样式，默认为'CONSTANT'。
    :return: 操作后目标图像
    """
    _operation_list = ['erode', 'dilate', 'open', 'close', 'gradient', 'tophat', 'blackhat']

    if operation not in _operation_list:
        raise ValueError(f'The operation value is incorrect! Please select in {_operation_list}')

    if operation == 'erode':
        return _erode(src, kernel, anchor, iterations, borderType)

    elif operation == 'dilate':
        return _dilate(src, kernel, anchor, iterations, borderType)

    elif operation == 'open':
        return _open(src, kernel, anchor, iterations, borderType)

    elif operation == 'close':
        return _close(src, kernel, anchor, iterations, borderType)

    elif operation == 'gradient':
        return _gradient(src, kernel, anchor, iterations, borderType)

    elif operation == 'tophat':
        return _tophat(src, kernel, anchor, iterations, borderType)

    elif operation == 'blackhat':
        return _blackhat(src, kernel, anchor, iterations, borderType)
    else:
        pass