| """Tests for physics-aware augmentations.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import pytest |
|
|
| from dipauglib.transforms.physics import ( |
| CastShadow, |
| ColourFade, |
| ColourTempShift, |
| DefocusBlur, |
| DustOverlay, |
| IlluminationGradient, |
| MotionBlur, |
| SensorNoise, |
| ) |
|
|
|
|
| TRANSFORMS = [ |
| IlluminationGradient, |
| CastShadow, |
| MotionBlur, |
| DefocusBlur, |
| ColourTempShift, |
| ColourFade, |
| DustOverlay, |
| SensorNoise, |
| ] |
|
|
|
|
| @pytest.mark.parametrize("transform_cls", TRANSFORMS) |
| def test_transform_preserves_shape_and_mask(transform_cls): |
| image = np.full((32, 32, 3), 128, dtype=np.uint8) |
| mask = np.zeros((32, 32), dtype=np.uint8) |
| mask[8:20, 10:22] = 1 |
| transform = transform_cls(intensity=0.5, p=1.0) |
| output = transform(image=image, mask=mask) |
| assert output["image"].shape == image.shape |
| assert output["mask"].shape == mask.shape |
| assert np.array_equal(output["mask"], mask) |
|
|
|
|
| @pytest.mark.parametrize("transform_cls", TRANSFORMS) |
| def test_transform_handles_all_black_image(transform_cls): |
| image = np.zeros((24, 24, 3), dtype=np.uint8) |
| mask = np.zeros((24, 24), dtype=np.uint8) |
| transform = transform_cls(intensity=1.0, p=1.0) |
| output = transform(image=image, mask=mask) |
| assert output["image"].shape == image.shape |
| assert output["mask"].shape == mask.shape |
|
|
|
|
| @pytest.mark.parametrize("transform_cls", TRANSFORMS) |
| def test_transform_handles_single_pixel_image(transform_cls): |
| image = np.array([[[255, 0, 0]]], dtype=np.uint8) |
| mask = np.array([[1]], dtype=np.uint8) |
| transform = transform_cls(intensity=0.2, p=1.0) |
| output = transform(image=image, mask=mask) |
| assert output["image"].shape == image.shape |
| assert output["mask"].shape == mask.shape |
|
|
