Add erosion

.pre-commit-config.yaml

@@ -44,7 +44,7 @@ repos:
       types: [python]
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.3.4
+    rev: v0.3.5
     hooks:
       # Run the linter.
       - id: ruff

README.md

@@ -164,6 +164,7 @@ Pixel-level transforms will change just an input image and will leave any additi
 - [Downscale](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Downscale)
 - [Emboss](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Emboss)
 - [Equalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Equalize)
+- [Erosion](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Erosion)
 - [FDA](https://albumentations.ai/docs/api_reference/augmentations/domain_adaptation/#albumentations.augmentations.domain_adaptation.FDA)
 - [FancyPCA](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.FancyPCA)
 - [FromFloat](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.FromFloat)

albumentations/augmentations/functional.py

@@ -69,6 +69,7 @@
     "MAX_VALUES_BY_DTYPE",
     "split_uniform_grid",
     "chromatic_aberration",
+    "erode",
 ]
 
 TWO = 2
@@ -1498,3 +1499,8 @@ def _distort_channel(
         interpolation=interpolation,
         borderMode=cv2.BORDER_REPLICATE,
     )
+
+
+@preserve_shape
+def erode(img: np.ndarray, kernel: np.ndarray) -> np.ndarray:
+    return cv2.erode(img, kernel, iterations=1)

albumentations/augmentations/transforms.py

@@ -14,7 +14,7 @@
 
 from albumentations import random_utils
 from albumentations.augmentations.blur.functional import blur
-from albumentations.augmentations.functional import split_uniform_grid
+from albumentations.augmentations.functional import erode, split_uniform_grid
 from albumentations.augmentations.utils import (
     get_num_channels,
     is_grayscale_image,
@@ -79,6 +79,7 @@
     "PixelDropout",
     "Spatter",
     "ChromaticAberration",
+    "Erosion",
 ]
 
 MAX_JPEG_QUALITY = 100
@@ -2866,3 +2867,60 @@ def _unmatch_sign(a: float, b: float) -> float:
 
     def get_transform_init_args_names(self) -> Tuple[str, str, str, str]:
         return "primary_distortion_limit", "secondary_distortion_limit", "mode", "interpolation"
+
+
+class Erosion(ImageOnlyTransform):
+    """Apply erosion operation to an image.
+
+    Erosion is a morphological operation that shrinks the white regions in a binary image,
+    or the foreground objects in case of a grayscale image. It is particularly useful in removing
+    small white noise, detaching objects, and emphasizing the structure of larger objects. This operation
+    is especially beneficial when working with scans of texts, as it can help clean up the background and
+    make the text more legible by reducing the size of the white spaces around the letters.
+
+
+    Args:
+        scale (int or tuple/list of int): Defines the size of the structuring element (kernel) used for erosion.
+            If an integer is provided, a square kernel of that size will be used.
+            If a tuple or list is provided, it should specify the range of sizes from which the kernel size
+            will be randomly selected for each application.
+        always_apply (bool, optional): Whether to always apply this transformation or not. Default is False.
+        p (float, optional): The probability with which the erosion operation will be applied. Default is 0.5.
+
+    Targets:
+        image
+
+    Image types:
+        uint8, float32
+
+    Reference:
+        https://github.com/facebookresearch/nougat
+
+    Example:
+        >>> import albumentations as A
+        >>> transform = A.Compose([A.Erosion(scale=(2, 3), p=0.5)])
+        >>> image = transform(image=image)["image"]
+    """
+
+    def __init__(
+        self,
+        scale: ScaleIntType = (2, 3),
+        always_apply: bool = False,
+        p: float = 0.5,
+    ):
+        super().__init__(always_apply, p)
+        self.scale = to_tuple(scale, scale)
+
+    def apply(self, img: np.ndarray, kernel: Tuple[int, int], **params: Any) -> np.ndarray:
+        return erode(img, kernel)
+
+    def get_params(self) -> Dict[str, float]:
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, tuple(random_utils.randint(self.scale[0], self.scale[1], 2))
+        )
+        return {
+            "kernel": kernel,
+        }
+
+    def get_transform_init_args_names(self) -> Tuple[str, ...]:
+        return ("scale",)

tests/test_serialization.py

@@ -450,6 +450,7 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
             mask_fill_value=20,
         )
     ],
+    [A.Erosion, {}]
 ]
 
 AUGMENTATION_CLS_EXCEPT = {