Fix in GaussNoise (#1801)

* Fix in GaussNoise * Added tests for GaussNoise
albumentations-team · Jun 19, 2024 · a209d6a · a209d6a
1 parent a6ceade
commit a209d6a
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py
@@ -5,7 +5,7 @@
 import cv2
 import numpy as np
 import skimage
-from albucore.functions import add, add_weighted, multiply, multiply_add
+from albucore.functions import add, add_array, add_weighted, multiply, multiply_add
 from albucore.utils import (
     MAX_VALUES_BY_DTYPE,
     clip,
@@ -1536,3 +1536,8 @@ def generate_approx_gaussian_noise(
     # Upsample the noise to the original shape using OpenCV
     result = cv2.resize(low_res_noise, (shape[1], shape[0]), interpolation=cv2.INTER_LINEAR)
     return result.reshape(shape)
+
+
+@clipped
+def add_noise(img: np.ndarray, noise: np.ndarray) -> np.ndarray:
+    return add_array(img, noise)
diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py
@@ -1650,7 +1650,7 @@ def __init__(
         self.noise_scale_factor = noise_scale_factor
 
     def apply(self, img: np.ndarray, gauss: np.ndarray, **params: Any) -> np.ndarray:
-        return albucore.add_array(img, gauss)
+        return fmain.add_noise(img, gauss)
 
     def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, float]:
         image = params["image"]
@@ -1673,7 +1673,7 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, f
             if image.ndim > MONO_CHANNEL_DIMENSIONS:
                 gauss = np.expand_dims(gauss, -1)
 
-        return {"gauss": gauss.astype(image.dtype)}
+        return {"gauss": gauss}
 
     @property
     def targets_as_params(self) -> List[str]:

diff --git a/tests/test_transforms.py b/tests/test_transforms.py
@@ -1748,4 +1748,18 @@ def test_random_fog_initialization(params, expected):
 ])
 def test_random_fog_invalid_input(params):
     with pytest.raises(Exception):
-        img_fog = A.RandomFog(**params)
+        A.RandomFog(**params)
+
+
+@pytest.mark.parametrize("image", IMAGES + [np.full((10, 10), 128, dtype=np.uint8)])
+@pytest.mark.parametrize("mean", (0, 10, -10))
+def test_gauss_noise(mean, image):
+    set_seed(42)
+    aug = A.GaussNoise(p=1, noise_scale_factor=1.0, mean=mean)
+
+    apply_params = aug.get_params_dependent_on_targets(params = {"image":image })
+
+    assert np.abs(mean - apply_params["gauss"].mean()) < 0.5
+    result = A.Compose([aug])(image=image)
+
+    assert not (result["image"] >= image).all()