Support keypoint on ElasticTransform (2-D)

albumentations/augmentations/geometric/transforms.py

@@ -1,7 +1,7 @@
 import math
 import random
 from enum import Enum
-from typing import Dict, Optional, Sequence, Tuple, Union
+from typing import Dict, Optional, Sequence, Tuple, Union, List
 
 import cv2
 import numpy as np
@@ -203,12 +203,13 @@ def __init__(
         self.same_dxdy = same_dxdy
 
     def apply(self, img, random_state=None, interpolation=cv2.INTER_LINEAR, **params):
+        # For supporting ``apply_to_keypoint`` function, Using ``elastic_transform_v2`` to replace ``elastic_transform``
         return F.elastic_transform(
             img,
             self.alpha,
             self.sigma,
             self.alpha_affine,
-            interpolation,
+            self.interpolation,
             self.border_mode,
             self.value,
             np.random.RandomState(random_state),
@@ -251,6 +252,28 @@ def apply_to_bbox(self, bbox, random_state=None, **params):
         bbox_returned = bbox_from_mask(mask)
         bbox_returned = F.normalize_bbox(bbox_returned, rows, cols)
         return bbox_returned
+
+    def apply_to_keypoint(self, keypoint: KeypointInternalType, random_state=None, half_win=9, **params) -> KeypointInternalType:
+        """
+            Only consider 2-D case.
+        """
+        image = np.zeros([params["cols"], params["rows"]])
+        X, Y = np.meshgrid(np.arange(image.shape[0]), np.arange(image.shape[1]))
+        image[Y, X] = np.exp(- 0.5 / (0.02 ** 2) *
+                             (((X - keypoint[0]) / image.shape[0]) ** 2 + ((Y - keypoint[1]) / image.shape[1]) ** 2))
+        remap_image = F.elastic_transform(
+            image,
+            self.alpha,
+            self.sigma,
+            self.alpha_affine,
+            self.interpolation,
+            self.border_mode,
+            self.mask_value,
+            np.random.RandomState(random_state),
+            self.approximate,
+        )
+        interp_y, interp_x = np.where(remap_image == np.max(remap_image))
+        return (interp_x[0], interp_y[0], 0.0, 0.0)
 
     def get_params(self):
         return {"random_state": random.randint(0, 10000)}

albumentations/augmentations/transforms.py

@@ -222,7 +222,7 @@ def get_transform_init_args_names(self):
 
 
 class ImageCompression(ImageOnlyTransform):
-    """Decreases image quality by Jpeg, WebP compression of an image.
+    """Decrease Jpeg, WebP compression of an image.
 
     Args:
         quality_lower (float): lower bound on the image quality.
@@ -292,7 +292,7 @@ def get_transform_init_args(self):
 
 
 class JpegCompression(ImageCompression):
-    """Decreases image quality by Jpeg compression of an image.
+    """Decrease Jpeg compression of an image.
 
     Args:
         quality_lower (float): lower bound on the jpeg quality. Should be in [0, 100] range
@@ -2391,10 +2391,6 @@ class Spatter(ImageOnlyTransform):
             If tuple of float intensity will be sampled from range `[intensity[0], intensity[1])`. Default: (0.6).
         mode (string, or list of strings): Type of corruption. Currently, supported options are 'rain' and 'mud'.
              If list is provided type of corruption will be sampled list. Default: ("rain").
-        color (list of (r, g, b) or dict or None): Corruption elements color.
-            If list uses provided list as color for specified mode.
-            If dict uses provided color for specified mode. Color for each specified mode should be provided in dict.
-            If None uses default colors (rain: (238, 238, 175), mud: (20, 42, 63)).
         p (float): probability of applying the transform. Default: 0.5.
 
     Targets:
@@ -2416,7 +2412,6 @@ def __init__(
         cutout_threshold: ScaleFloatType = 0.68,
         intensity: ScaleFloatType = 0.6,
         mode: Union[str, Sequence[str]] = "rain",
-        color: Optional[Union[Sequence[int], Dict[str, Sequence[int]]]] = None,
         always_apply: bool = False,
         p: float = 0.5,
     ):
@@ -2427,34 +2422,10 @@ def __init__(
         self.gauss_sigma = to_tuple(gauss_sigma, gauss_sigma)
         self.intensity = to_tuple(intensity, intensity)
         self.cutout_threshold = to_tuple(cutout_threshold, cutout_threshold)
-        self.color = (
-            color
-            if color is not None
-            else {
-                "rain": [238, 238, 175],
-                "mud": [20, 42, 63],
-            }
-        )
         self.mode = mode if isinstance(mode, (list, tuple)) else [mode]
-
-        if len(set(self.mode)) > 1 and not isinstance(self.color, dict):
-            raise ValueError(f"Unsupported color: {self.color}. Please specify color for each mode (use dict for it).")
-
         for i in self.mode:
             if i not in ["rain", "mud"]:
                 raise ValueError(f"Unsupported color mode: {mode}. Transform supports only `rain` and `mud` mods.")
-            if isinstance(self.color, dict):
-                if i not in self.color:
-                    raise ValueError(f"Wrong color definition: {self.color}. Color for mode: {i} not specified.")
-                if len(self.color[i]) != 3:
-                    raise ValueError(
-                        f"Unsupported color: {self.color[i]} for mode {i}. Color should be presented in RGB format."
-                    )
-
-        if isinstance(self.color, (list, tuple)):
-            if len(self.color) != 3:
-                raise ValueError(f"Unsupported color: {self.color}. Color should be presented in RGB format.")
-            self.color = {self.mode[0]: self.color}
 
     def apply(
         self,
@@ -2480,7 +2451,6 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, A
         sigma = random.uniform(self.gauss_sigma[0], self.gauss_sigma[1])
         mode = random.choice(self.mode)
         intensity = random.uniform(self.intensity[0], self.intensity[1])
-        color = np.array(self.color[mode]) / 255.0
 
         liquid_layer = random_utils.normal(size=(h, w), loc=mean, scale=std)
         liquid_layer = gaussian_filter(liquid_layer, sigma=sigma, mode="nearest")
@@ -2501,16 +2471,15 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, A
             m = liquid_layer * dist
             m *= 1 / np.max(m, axis=(0, 1))
 
-            drops = m[:, :, None] * color * intensity
+            drops = m[:, :, None] * np.array([238 / 255.0, 238 / 255.0, 175 / 255.0]) * intensity
             mud = None
             non_mud = None
         else:
             m = np.where(liquid_layer > cutout_threshold, 1, 0)
             m = gaussian_filter(m.astype(np.float32), sigma=sigma, mode="nearest")
             m[m < 1.2 * cutout_threshold] = 0
             m = m[..., np.newaxis]
-
-            mud = m * color
+            mud = m * np.array([20 / 255.0, 42 / 255.0, 63 / 255.0])
             non_mud = 1 - m
             drops = None
 
@@ -2521,5 +2490,5 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, A
             "mode": mode,
         }
 
-    def get_transform_init_args_names(self) -> Tuple[str, str, str, str, str, str, str]:
-        return "mean", "std", "gauss_sigma", "intensity", "cutout_threshold", "mode", "color"
+    def get_transform_init_args_names(self) -> Tuple[str, str, str, str, str, str]:
+        return "mean", "std", "gauss_sigma", "intensity", "cutout_threshold", "mode"