Update elastic

.pre-commit-config.yaml

@@ -51,7 +51,7 @@ repos:
         files: setup.py
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.7.1
+    rev: v0.7.2
     hooks:
       # Run the linter.
       - id: ruff

albumentations/augmentations/geometric/functional.py

@@ -1488,18 +1488,37 @@
     same_dxdy: bool,
     kernel_size: tuple[int, int],
     random_generator: np.random.Generator,
+    noise_distribution: Literal["gaussian", "uniform"],
 ) -> tuple[np.ndarray, np.ndarray]:
-    """Generate displacement fields for elastic transform."""
-    dx = random_generator.standard_normal(size=image_shape[:2]).astype(np.float32) * 2 - 1
-    cv2.GaussianBlur(dx, kernel_size, sigma, dst=dx)
-    dx *= alpha
+    """Generate displacement fields for elastic transform.
 
-    if same_dxdy:
-        dy = dx
-    else:
-        dy = random_generator.standard_normal(size=image_shape[:2]).astype(np.float32) * 2 - 1
-        cv2.GaussianBlur(dy, kernel_size, sigma, dst=dy)
-        dy *= alpha
+    Args:
+        image_shape: Shape of the image (height, width)
+        alpha: Scaling factor for displacement
+        sigma: Standard deviation for Gaussian blur
+        same_dxdy: Whether to use same displacement field for both directions
+        kernel_size: Size of Gaussian blur kernel
+        random_generator: NumPy random number generator
+        noise_distribution: Type of noise distribution to use ("gaussian" or "uniform")
+
+    Returns:
+        tuple: (dx, dy) displacement fields
+    """
+
+    def generate_noise_field() -> np.ndarray:
+        if noise_distribution == "gaussian":
+            field = random_generator.standard_normal(size=image_shape[:2]).astype(np.float32) * 2 - 1
+            cv2.GaussianBlur(field, kernel_size, sigma, dst=field)
+        elif noise_distribution == "uniform":
+            field = random_generator.uniform(low=-1, high=1, size=image_shape[:2]).astype(np.float32)
+            cv2.GaussianBlur(field, kernel_size, sigma, dst=field)
+        return field * alpha
+
+    # Generate first displacement field
+    dx = generate_noise_field()
+
+    # Generate or copy second displacement field
+    dy = dx if same_dxdy else generate_noise_field()
 
     return dx, dy
 

albumentations/augmentations/geometric/transforms.py

@@ -186,6 +186,11 @@ class ElasticTransform(BaseDistortion):
         mask_interpolation (int): Flag that is used to specify the interpolation algorithm for mask.
             Should be one of: cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4.
             Default: cv2.INTER_NEAREST.
+        noise_distribution (Literal["gaussian", "uniform"]): Distribution used to generate the displacement fields.
+            "gaussian" generates fields using normal distribution (more natural deformations).
+            "uniform" generates fields using uniform distribution (more mechanical deformations).
+            Default: "gaussian".
+
         p (float): Probability of applying the transform. Default: 0.5
 
     Targets:
@@ -220,6 +225,7 @@ class InitSchema(BaseDistortion.InitSchema):
         sigma: Annotated[float, Field(ge=1)]
         approximate: bool
         same_dxdy: bool
+        noise_distribution: Literal["gaussian", "uniform"]
 
     def __init__(
         self,
@@ -233,6 +239,7 @@ def __init__(
         approximate: bool = False,
         same_dxdy: bool = False,
         mask_interpolation: int = cv2.INTER_NEAREST,
+        noise_distribution: Literal["gaussian", "uniform"] = "gaussian",
         p: float = 0.5,
     ):
         super().__init__(
@@ -248,6 +255,7 @@ def __init__(
         self.sigma = sigma
         self.approximate = approximate
         self.same_dxdy = same_dxdy
+        self.noise_distribution = noise_distribution
 
     def get_params_dependent_on_data(self, params: dict[str, Any], data: dict[str, Any]) -> dict[str, Any]:
         height, width = params["shape"][:2]
@@ -261,6 +269,7 @@ def get_params_dependent_on_data(self, params: dict[str, Any], data: dict[str, A
             same_dxdy=self.same_dxdy,
             kernel_size=kernel_size,
             random_generator=self.random_generator,
+            noise_distribution=self.noise_distribution,
         )
 
         x, y = np.meshgrid(np.arange(width), np.arange(height))
@@ -270,7 +279,14 @@ def get_params_dependent_on_data(self, params: dict[str, Any], data: dict[str, A
         return {"map_x": map_x, "map_y": map_y}
 
     def get_transform_init_args_names(self) -> tuple[str, ...]:
-        return (*super().get_transform_init_args_names(), "alpha", "sigma", "approximate", "same_dxdy")
+        return (
+            *super().get_transform_init_args_names(),
+            "alpha",
+            "sigma",
+            "approximate",
+            "same_dxdy",
+            "noise_distribution",
+        )
 
 
 class Perspective(DualTransform):