Source code for draugr.opencv_utilities.color_space.clean

from enum import Enum

import cv2
from sorcery import assigned_names

from draugr.opencv_utilities.namespaces.enums import (
    MorphShapeEnum,
    MorphTypeEnum,
    BorderTypeEnum,
)

__all__ = ["clean_up", "CleanUpMethod"]


[docs]class CleanUpMethod(Enum): close, open, nclose, nopen, erode, dilate, none = assigned_names()
[docs]def clean_up(img, method: CleanUpMethod = CleanUpMethod.open, **kwargs): """ :param img: :type img: :param method: :type method: :param kwargs: :type kwargs: :return: :rtype: """ filter_size = ( max(int(img.shape[1] * 0.01), kwargs.get("min_filter_size", 3)), max(int(img.shape[0] * 0.01), kwargs.get("min_filter_size", 3)), ) morph_shape = MorphShapeEnum.rect.value kernel = cv2.getStructuringElement(morph_shape, filter_size) if method == CleanUpMethod.close: img = cv2.morphologyEx( img, MorphTypeEnum.close.value, kernel, iterations=kwargs.get("iterations", 1), borderType=kwargs.get("borderType", None), # BorderTypeEnum.constant.value, borderValue=kwargs.get("borderValue", 0), # 255 ) elif method == CleanUpMethod.open: img = cv2.morphologyEx( img, MorphTypeEnum.open.value, kernel, iterations=kwargs.get("iterations", 1), borderType=kwargs.get("borderType", None), # BorderTypeEnum.constant.value, borderValue=kwargs.get("borderValue", 0), # 255 ) elif method == CleanUpMethod.nclose: img = cv2.morphologyEx( img, MorphTypeEnum.dilate.value, kernel, None, None, iterations=kwargs.get("iterations", 2), borderType=kwargs.get("borderType", BorderTypeEnum.reflect101.value), ) img = cv2.morphologyEx( img, MorphTypeEnum.erode.value, kernel, None, None, iterations=kwargs.get("iterations", 2), borderType=kwargs.get("borderType", BorderTypeEnum.reflect101.value), ) elif method == CleanUpMethod.erode: img = cv2.dilate( img, kernel, iterations=kwargs.get("iterations", 1), ) elif method == CleanUpMethod.dilate: img = cv2.erode( img, kernel, iterations=kwargs.get("iterations", 1), ) return img