diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9335cdab2..d223c9ad1 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py index a78591122..a2c7b45ab 100644 --- a/albumentations/augmentations/functional.py +++ b/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( diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py index 91bbad2d3..fbf57550e 100644 --- a/albumentations/augmentations/transforms.py +++ b/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} diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py index 6ac744b98..797a534fb 100644 --- a/tests/functional/test_functional.py +++ b/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)