Update random rain

.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.9.3
+    rev: v0.9.4
     hooks:
       # Run the linter.
       - id: ruff

albumentations/augmentations/functional.py

@@ -746,45 +746,41 @@ def add_rain(
     drop_color: tuple[int, int, int],
     blur_value: int,
     brightness_coefficient: float,
-    rain_drops: list[tuple[int, int]],
+    rain_drops: np.ndarray,
 ) -> np.ndarray:
-    """Adds rain drops to the image.
+    """Optimized version using OpenCV line drawing."""
+    if not rain_drops.size:
+        return img.copy()
 
-    Args:
-        img (np.ndarray): Input image.
-        slant (int): The angle of the rain drops.
-        drop_length (int): The length of each rain drop.
-        drop_width (int): The width of each rain drop.
-        drop_color (tuple[int, int, int]): The color of the rain drops in RGB format.
-        blur_value (int): The size of the kernel used to blur the image. Rainy views are blurry.
-        brightness_coefficient (float): Coefficient to adjust the brightness of the image. Rainy days are usually shady.
-        rain_drops (list[tuple[int, int]]): A list of tuples where each tuple represents the (x, y)
-            coordinates of the starting point of a rain drop.
+    img = img.copy()
 
-    Returns:
-        np.ndarray: Image with rain effect added.
+    # Pre-allocate rain layer
+    rain_layer = np.zeros_like(img, dtype=np.uint8)
 
-    Reference:
-        https://github.com/UjjwalSaxena/Automold--Road-Augmentation-Library
-    """
-    img = img.copy()
-    for rain_drop_x0, rain_drop_y0 in rain_drops:
-        rain_drop_x1 = rain_drop_x0 + slant
-        rain_drop_y1 = rain_drop_y0 + drop_length
-
-        cv2.line(
-            img,
-            (rain_drop_x0, rain_drop_y0),
-            (rain_drop_x1, rain_drop_y1),
-            drop_color,
-            drop_width,
-        )
+    # Calculate end points correctly
+    end_points = rain_drops + np.array([[slant, drop_length]])  # This creates correct shape
+
+    # Stack arrays properly - both must be same shape arrays
+    lines = np.stack((rain_drops, end_points), axis=1)  # Use tuple and proper axis
+
+    cv2.polylines(
+        rain_layer,
+        lines.astype(np.int32),
+        False,
+        drop_color,
+        drop_width,
+        lineType=cv2.LINE_4,
+    )
+
+    if blur_value > 1:
+        cv2.blur(rain_layer, (blur_value, blur_value), dst=rain_layer)
 
-    img = cv2.blur(img, (blur_value, blur_value))  # rainy view are blurry
-    image_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
-    image_hsv[:, :, 2] *= brightness_coefficient
+    cv2.add(img, rain_layer, dst=img)
 
-    return cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
+    if brightness_coefficient != 1.0:
+        cv2.multiply(img, brightness_coefficient, dst=img, dtype=cv2.CV_8U)
+
+    return img
 
 
 def get_fog_particle_radiuses(

albumentations/augmentations/transforms.py

@@ -796,7 +796,7 @@ def apply(
         img: np.ndarray,
         slant: int,
         drop_length: int,
-        rain_drops: list[tuple[int, int]],
+        rain_drops: np.ndarray,
         **params: Any,
     ) -> np.ndarray:
         non_rgb_error(img)
@@ -817,31 +817,36 @@ def get_params_dependent_on_data(
         params: dict[str, Any],
         data: dict[str, Any],
     ) -> dict[str, Any]:
-        slant = int(self.py_random.uniform(*self.slant_range))
-
         height, width = params["shape"][:2]
-        area = height * width
 
+        # Simpler calculations, directly following Kornia
         if self.rain_type == "drizzle":
-            num_drops = area // 770
-            drop_length = 10
+            num_drops = height // 4
         elif self.rain_type == "heavy":
-            num_drops = width * height // 600
-            drop_length = 30
+            num_drops = height
         elif self.rain_type == "torrential":
-            num_drops = area // 500
-            drop_length = 60
+            num_drops = height * 2
         else:
-            drop_length = self.drop_length
-            num_drops = area // 600
-
-        rain_drops = []
-
-        for _ in range(num_drops):  # If You want heavy rain, try increasing this
-            x = self.py_random.randint(slant, width) if slant < 0 else self.py_random.randint(0, max(width - slant, 0))
-            y = self.py_random.randint(0, max(height - drop_length, 0))
-
-            rain_drops.append((x, y))
+            num_drops = height // 3
+
+        # Fixed proportion for drop length (like Kornia)
+        drop_length = max(1, height // 8)
+
+        # Simplified slant calculation
+        slant = self.random_generator.integers(-width // 50, width // 50)
+
+        # Single random call for all coordinates
+        if num_drops > 0:
+            # Generate all coordinates in one call
+            coords = self.random_generator.integers(
+                low=[0, 0],
+                high=[width, height - drop_length],
+                size=(num_drops, 2),
+                dtype=np.int32,
+            )
+            rain_drops = coords
+        else:
+            rain_drops = np.empty((0, 2), dtype=np.int32)
 
         return {"drop_length": drop_length, "slant": slant, "rain_drops": rain_drops}
 

tests/functional/test_functional.py

@@ -2326,3 +2326,134 @@ def test_gaussian_illumination_sigma(sigma, expected_pattern):
 
     if expected_pattern == "narrow":
         assert diff > wide_diff  # Narrow should have steeper falloff than wide
+
+
+
+@pytest.mark.parametrize(
+    ["img", "slant", "drop_length", "drop_width", "drop_color", "blur_value", "brightness_coefficient", "rain_drops", "expected_shape"],
+    [
+        # Test basic functionality with small image
+        (
+            np.zeros((10, 10, 3), dtype=np.uint8),
+            5,
+            3,
+            1,
+            (200, 200, 200),
+            3,
+            0.7,
+            np.array([(2, 2)]),
+            (10, 10, 3),
+        ),
+        # Test with no rain drops
+        (
+            np.zeros((20, 20, 3), dtype=np.uint8),
+            5,
+            3,
+            1,
+            (200, 200, 200),
+            3,
+            0.7,
+            np.array([]).reshape(0, 2),
+            (20, 20, 3),
+        ),
+        # Test with multiple rain drops
+        (
+            np.zeros((30, 30, 3), dtype=np.uint8),
+            -5,
+            5,
+            2,
+            (255, 255, 255),
+            5,
+            0.8,
+            np.array([(5, 5), (10, 10), (15, 15)]),
+            (30, 30, 3),
+        ),
+    ]
+)
+def test_add_rain_shape_and_type(
+    img, slant, drop_length, drop_width, drop_color, blur_value, brightness_coefficient, rain_drops, expected_shape
+):
+    result = fmain.add_rain(
+        img, slant, drop_length, drop_width, drop_color, blur_value, brightness_coefficient, rain_drops
+    )
+    assert result.shape == expected_shape
+    assert result.dtype == np.uint8
+
+
+@pytest.mark.parametrize("brightness_coefficient", [0.5, 0.7, 1.0])
+def test_add_rain_brightness(brightness_coefficient):
+    """Test that brightness coefficient correctly affects image brightness"""
+    img = np.full((20, 20, 3), 100, dtype=np.uint8)
+    rain_drops = np.array([(5, 5)])
+
+    result = fmain.add_rain(
+        img=img,
+        slant=5,
+        drop_length=3,
+        drop_width=1,
+        drop_color=(200, 200, 200),
+        blur_value=3,
+        brightness_coefficient=brightness_coefficient,
+        rain_drops=rain_drops,
+    )
+
+    # Convert to HSV to check brightness
+    original_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    result_hsv = cv2.cvtColor(result, cv2.COLOR_RGB2HSV)
+
+    if brightness_coefficient < 1.0:
+        # For darkening coefficients, brightness should decrease
+        assert np.mean(result_hsv[:, :, 2]) < np.mean(original_hsv[:, :, 2])
+        np.testing.assert_allclose(
+            np.mean(result_hsv[:, :, 2]) / np.mean(original_hsv[:, :, 2]),
+            brightness_coefficient,
+            rtol=0.1  # Allow 10% tolerance due to rounding and blur effects
+        )
+    else:
+        # For brightness_coefficient = 1.0, brightness might slightly increase
+        # due to bright rain drops and blur, but shouldn't change dramatically
+        np.testing.assert_allclose(
+            np.mean(result_hsv[:, :, 2]) / np.mean(original_hsv[:, :, 2]),
+            1.0,
+            rtol=0.1  # Allow 10% tolerance
+        )
+
+
+def test_add_rain_drops_visibility():
+    """Test that rain drops are actually visible in the output"""
+    img = np.zeros((20, 20, 3), dtype=np.uint8)
+    rain_drops = np.array([(5, 5)])
+    drop_color = (255, 255, 255)
+
+    result = fmain.add_rain(
+        img=img,
+        slant=0,
+        drop_length=5,
+        drop_width=1,
+        drop_color=drop_color,
+        blur_value=1,  # Minimal blur to check drop visibility
+        brightness_coefficient=1.0,  # No brightness change
+        rain_drops=rain_drops,
+    )
+
+    # Check if any pixels have the rain drop color
+    assert np.any(result > 0)
+
+
+def test_add_rain_preserves_input():
+    """Test that the function doesn't modify the input image"""
+    img = np.zeros((10, 10, 3), dtype=np.uint8)
+    img_copy = img.copy()
+
+    fmain.add_rain(
+        img=img,
+        slant=5,
+        drop_length=3,
+        drop_width=1,
+        drop_color=(200, 200, 200),
+        blur_value=3,
+        brightness_coefficient=0.7,
+        rain_drops=np.array([(5, 5)]),
+    )
+
+    np.testing.assert_array_equal(img, img_copy)