diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py index 364f381f8..4363e6e23 100644 --- a/albumentations/augmentations/functional.py +++ b/albumentations/augmentations/functional.py @@ -2662,7 +2662,6 @@ def create_directional_gradient(height: int, width: int, angle: float) -> np.nda @float32_io -@clipped def apply_linear_illumination(img: np.ndarray, intensity: float, angle: float) -> np.ndarray: """Apply directional illumination effect to an image using a linear gradient. @@ -2728,25 +2727,38 @@ def apply_corner_illumination( corner: Literal[0, 1, 2, 3], ) -> np.ndarray: """Apply corner-based illumination effect.""" - result, height, width = prepare_illumination_input(img) + if intensity == 0: + return img.copy() - # Create distance map coordinates - y, x = np.ogrid[:height, :width] + height, width = img.shape[:2] + + # Pre-compute diagonal length once + diagonal_length = math.sqrt(height * height + width * width) - # Adjust coordinates based on corner - if corner == 1: # top-right - x = width - 1 - x - elif corner == 2: # bottom-right - x = width - 1 - x - y = height - 1 - y - elif corner == 3: # bottom-left - y = height - 1 - y + # Create inverted distance map mask directly + # Use uint8 for distanceTransform regardless of input dtype + mask = np.full((height, width), 255, dtype=np.uint8) - # Calculate normalized distance - distance = np.sqrt(x * x + y * y) / np.sqrt(height * height + width * width) - pattern = 1 - distance # Invert so corner is brightest + # Use array indexing instead of conditionals + corners = [(0, 0), (0, width - 1), (height - 1, width - 1), (height - 1, 0)] + mask[corners[corner]] = 0 + + # Calculate distance transform + pattern = cv2.distanceTransform( + mask, + distanceType=cv2.DIST_L2, + maskSize=cv2.DIST_MASK_PRECISE, + dstType=cv2.CV_32F, # Specify float output directly + ) + + # Combine operations to reduce array copies + cv2.multiply(pattern, -intensity / diagonal_length, dst=pattern) + cv2.add(pattern, 1, dst=pattern) + + if img.ndim == NUM_MULTI_CHANNEL_DIMENSIONS: + pattern = cv2.merge([pattern] * img.shape[2]) - return apply_illumination_pattern(result, pattern, intensity) + return multiply_by_array(img, pattern) @clipped diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py index 85814e095..7b397b9a9 100644 --- a/tests/functional/test_functional.py +++ b/tests/functional/test_functional.py @@ -2118,3 +2118,126 @@ def test_gradient_range(): gradient = fmain.create_directional_gradient(10, 10, angle) assert gradient.min() >= 0 - 1e-7 assert gradient.max() <= 1 + 1e-7 + + + +@pytest.mark.parametrize( + ["corner", "intensity", "expected_corner"], + [ + # Test each corner with positive intensity (brightening) + (0, 0.2, (0, 0)), # top-left is brightest + (1, 0.2, (0, 9)), # top-right is brightest + (2, 0.2, (9, 9)), # bottom-right is brightest + (3, 0.2, (9, 0)), # bottom-left is brightest + + # Test with negative intensity (darkening) + (0, -0.2, (9, 9)), # top-left is darkest, opposite corner is brightest + (1, -0.2, (9, 0)), # top-right is darkest, opposite corner is brightest + (2, -0.2, (0, 0)), # bottom-right is darkest, opposite corner is brightest + (3, -0.2, (0, 9)), # bottom-left is darkest, opposite corner is brightest + ], +) +def test_corner_illumination_brightest_point(corner, intensity, expected_corner): + """Test that the illumination pattern has maximum intensity at the correct corner.""" + # Create a constant test image + image = np.full((10, 10), 0.5, dtype=np.float32) + + # Apply corner illumination + result = fmain.apply_corner_illumination(image, intensity, corner) + + # Find the brightest point + actual_corner = np.unravel_index(np.argmax(result), result.shape) + + assert actual_corner == expected_corner + + +@pytest.mark.parametrize( + ["shape", "dtype"], + [ + ((10, 10), np.float32), # grayscale float32 + ((10, 10), np.uint8), # grayscale uint8 + ((10, 10, 3), np.float32), # RGB float32 + ((10, 10, 3), np.uint8), # RGB uint8 + # Removed single channel test case as it's not supported + ], +) +def test_corner_illumination_preserves_shape_and_type(shape, dtype): + """Test that the output maintains the input shape and dtype.""" + # Create test image + image = np.ones(shape, dtype=dtype) + if dtype == np.uint8: + image *= 255 + + # Apply corner illumination + result = fmain.apply_corner_illumination(image, intensity=0.2, corner=0) + + assert result.shape == shape + assert result.dtype == dtype + + +@pytest.mark.parametrize("intensity", [-0.2, 0, 0.2]) +def test_corner_illumination_intensity_range(intensity): + """Test that the output values stay within valid range.""" + # Create test images with extreme values + image_zeros = np.zeros((10, 10), dtype=np.float32) + image_ones = np.ones((10, 10), dtype=np.float32) + + # Apply corner illumination + result_zeros = fmain.apply_corner_illumination(image_zeros, intensity, corner=0) + result_ones = fmain.apply_corner_illumination(image_ones, intensity, corner=0) + + # Check that values stay in valid range + assert np.all(result_zeros >= 0) + assert np.all(result_zeros <= 1) + assert np.all(result_ones >= 0) + assert np.all(result_ones <= 1) + + +def test_corner_illumination_identity_zero_intensity(): + """Test that zero intensity returns the input image unchanged.""" + # Create random test image + image = np.random.rand(10, 10).astype(np.float32) + + # Apply corner illumination with zero intensity + result = fmain.apply_corner_illumination(image, intensity=0, corner=0) + + np.testing.assert_array_almost_equal(result, image, decimal=2) + + +@pytest.mark.parametrize("corner", [0, 1, 2, 3]) +def test_corner_illumination_symmetry(corner): + """Test that the illumination pattern is symmetric around the corner.""" + # Create test image + image = np.ones((11, 11), dtype=np.float32) # Odd dimensions for clear center + + # Apply corner illumination + result = fmain.apply_corner_illumination(image, intensity=0.2, corner=corner) + + # Get distances from corner to test symmetry + if corner == 0: # top-left + d1 = result[0, 1] # one step right + d2 = result[1, 0] # one step down + elif corner == 1: # top-right + d1 = result[0, -2] # one step left + d2 = result[1, -1] # one step down + elif corner == 2: # bottom-right + d1 = result[-1, -2] # one step left + d2 = result[-2, -1] # one step up + else: # bottom-left + d1 = result[-1, 1] # one step right + d2 = result[-2, 0] # one step up + + np.testing.assert_almost_equal(d1, d2) + + +def test_corner_illumination_multichannel_consistency(): + """Test that all channels are modified identically for RGB images.""" + # Create RGB test image + image = np.ones((10, 10, 3), dtype=np.float32) + + # Apply corner illumination + result = fmain.apply_corner_illumination(image, intensity=0.2, corner=0) + + # Check that all channels are identical + np.testing.assert_array_almost_equal(result[..., 0], result[..., 1]) + np.testing.assert_array_almost_equal(result[..., 1], result[..., 2])