diff --git a/albumentations/augmentations/geometric/transforms.py b/albumentations/augmentations/geometric/transforms.py index 866da1992..e148619be 100644 --- a/albumentations/augmentations/geometric/transforms.py +++ b/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)} diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py index e90b3bef8..fd9b039c7 100644 --- a/albumentations/augmentations/transforms.py +++ b/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,7 +2471,7 @@ 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: @@ -2509,8 +2479,7 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, A 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"