Spaces:
Running
Running
| import albumentations as alb | |
| import numpy as np | |
| import cv2 | |
| class Erosion(alb.ImageOnlyTransform): | |
| """ | |
| Apply erosion operation to an image. | |
| Erosion is a morphological operation that shrinks the white regions in a binary image. | |
| Args: | |
| scale (int or tuple/list of int): The scale or range for the size of the erosion kernel. | |
| If an integer is provided, a square kernel of that size will be used. | |
| If a tuple or list is provided, it should contain two integers representing the minimum | |
| and maximum sizes for the erosion kernel. | |
| always_apply (bool, optional): Whether to always apply this transformation. Default is False. | |
| p (float, optional): The probability of applying this transformation. Default is 0.5. | |
| Returns: | |
| numpy.ndarray: The transformed image. | |
| """ | |
| def __init__(self, scale, always_apply=False, p=0.5): | |
| super().__init__(always_apply=always_apply, p=p) | |
| if type(scale) is tuple or type(scale) is list: | |
| assert len(scale) == 2 | |
| self.scale = scale | |
| else: | |
| self.scale = (scale, scale) | |
| def apply(self, img, **params): | |
| kernel = cv2.getStructuringElement( | |
| cv2.MORPH_ELLIPSE, tuple(np.random.randint(self.scale[0], self.scale[1], 2)) | |
| ) | |
| img = cv2.erode(img, kernel, iterations=1) | |
| return img | |
| class Dilation(alb.ImageOnlyTransform): | |
| """ | |
| Apply dilation operation to an image. | |
| Dilation is a morphological operation that expands the white regions in a binary image. | |
| Args: | |
| scale (int or tuple/list of int): The scale or range for the size of the dilation kernel. | |
| If an integer is provided, a square kernel of that size will be used. | |
| If a tuple or list is provided, it should contain two integers representing the minimum | |
| and maximum sizes for the dilation kernel. | |
| always_apply (bool, optional): Whether to always apply this transformation. Default is False. | |
| p (float, optional): The probability of applying this transformation. Default is 0.5. | |
| Returns: | |
| numpy.ndarray: The transformed image. | |
| """ | |
| def __init__(self, scale, always_apply=False, p=0.5): | |
| super().__init__(always_apply=always_apply, p=p) | |
| if type(scale) is tuple or type(scale) is list: | |
| assert len(scale) == 2 | |
| self.scale = scale | |
| else: | |
| self.scale = (scale, scale) | |
| def apply(self, img, **params): | |
| kernel = cv2.getStructuringElement( | |
| cv2.MORPH_ELLIPSE, tuple(np.random.randint(self.scale[0], self.scale[1], 2)) | |
| ) | |
| img = cv2.dilate(img, kernel, iterations=1) | |
| return img | |
| class Bitmap(alb.ImageOnlyTransform): | |
| """ | |
| Apply a bitmap-style transformation to an image. | |
| This transformation replaces all pixel values below a certain threshold with a specified value. | |
| Args: | |
| value (int, optional): The value to replace pixels below the threshold with. Default is 0. | |
| lower (int, optional): The threshold value below which pixels will be replaced. Default is 200. | |
| always_apply (bool, optional): Whether to always apply this transformation. Default is False. | |
| p (float, optional): The probability of applying this transformation. Default is 0.5. | |
| Returns: | |
| numpy.ndarray: The transformed image. | |
| """ | |
| def __init__(self, value=0, lower=200, always_apply=False, p=0.5): | |
| super().__init__(always_apply=always_apply, p=p) | |
| self.lower = lower | |
| self.value = value | |
| def apply(self, img, **params): | |
| img = img.copy() | |
| img[img < self.lower] = self.value | |
| return img | |