Add gaussian noise transform #6192

docs/source/transforms.rst

@@ -123,6 +123,7 @@ Transforms on PIL Image and torch.\*Tensor
     Resize
     TenCrop
     GaussianBlur
+    GaussianNoise
     RandomInvert
     RandomPosterize
     RandomSolarize

gallery/plot_transforms.py

@@ -119,6 +119,15 @@ def plot(imgs, with_orig=True, row_title=None, **imshow_kwargs):
 blurred_imgs = [blurrer(orig_img) for _ in range(4)]
 plot(blurred_imgs)
 
+####################################
+# GaussianNoise
+# ~~~~~~~~~~~~~
+# The :class:`~torchvision.transforms.GaussianNoise` transform
+# perturbs the input image with gaussian noise.
+noisy = T.GaussianNoise(mean=0, sigma=(0.1, 2.0))
+noisy_imgs = [noisy(orig_img) for _ in range(2)]
+plot(noisy_imgs)
+
 ####################################
 # RandomPerspective
 # ~~~~~~~~~~~~~~~~~

test/test_transforms.py

@@ -1314,6 +1314,25 @@ def test_gaussian_blur_asserts():
         transforms.GaussianBlur(3, "sigma_string")
 
 
+def test_gaussian_noise():
+    np_img = np.ones((100, 100, 3), dtype=np.uint8) * 255
+    img = F.to_pil_image(np_img, "RGB")
+    out = transforms.GaussianNoise(2.0, (0.1, 2.0))(img)
+    assert isinstance(out, PIL.Image.Image)
+
+    with pytest.raises(TypeError, match="Tensor is not a torch image"):
+        out = transforms.GaussianNoise(2.0, (0.1, 2.0))(torch.ones((4)))
+
+    with pytest.raises(ValueError, match="Mean should be a positive number"):
+        transforms.GaussianNoise(-1)
+
+    with pytest.raises(ValueError, match="If sigma is a single number, it must be positive."):
+        transforms.GaussianNoise(2.0, -1)
+
+    with pytest.raises(ValueError, match="sigma should be a single number or a list/tuple with length 2."):
+        transforms.GaussianNoise(2.0, (1, 2, 3))
+
+
 def test_lambda():
     trans = transforms.Lambda(lambda x: x.add(10))
     x = torch.randn(10)

torchvision/transforms/transforms.py

@@ -46,6 +46,7 @@
     "RandomPerspective",
     "RandomErasing",
     "GaussianBlur",
+    "GaussianNoise",
     "InterpolationMode",
     "RandomInvert",
     "RandomPosterize",
@@ -1837,6 +1838,88 @@ def __repr__(self) -> str:
         return s
 
 
+class GaussianNoise(torch.nn.Module):
+    """Adds Gaussian noise to the image with specified mean and standard deviation.
+    If the image is torch Tensor, it is expected
+    to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions.
+
+    Args:
+        mean (float or sequence): Mean of the sampling gaussian distribution .
+        sigma (float or tuple of float (min, max)): Standard deviation to be used for
+            sampling the gaussian noise. If float, sigma is fixed. If it is tuple
+            of float (min, max), sigma is chosen uniformly at random to lie in the
+            given range.
+
+    Returns:
+        PIL Image or Tensor: Input image perturbed with Gaussian Noise.
+
+    """
+
+    def __init__(self, mean, sigma = (0.1, 2.0)):
+        super().__init__()
+        _log_api_usage_once(self)
+
+        if mean < 0:
+            raise ValueError("Mean should be a positive number")
+
+        if isinstance(sigma, numbers.Number):
+            if sigma <= 0:
+                raise ValueError("If sigma is a single number, it must be positive.")
+            sigma = (sigma, sigma)
+        elif isinstance(sigma, Sequence) and len(sigma) == 2:
+            if not 0.0 < sigma[0] <= sigma[1]:
+                raise ValueError("sigma values should be positive and of the form (min, max).")
+        else:
+            raise ValueError("sigma should be a single number or a list/tuple with length 2.")
+
+        self.mean = mean
+        self.sigma = sigma
+
+    @staticmethod
+    def get_params(sigma_min: float, sigma_max: float) -> float:
+        return torch.empty(1).uniform_(sigma_min, sigma_max).item()
+
+    def forward(self, image: Tensor) -> Tensor:
+        """
+        Args:
+            image (PIL Image or Tensor): image to be perturbed with gaussian noise.
+
+        Returns:
+            PIL Image or Tensor: Image added with gaussian noise.
+        """
+        sigma = self.get_params(self.sigma[0], self.sigma[1])
+        if not isinstance(image, torch.Tensor):
+            if not F._is_pil_image(image):
+                raise TypeError(f"image should be PIL Image or Tensor. Got {type(image)}")
+
+            t_image = F.pil_to_tensor(image)
+        else:
+            t_image = image
+
+        if not t_image.ndim >= 2:
+            raise TypeError("Tensor is not a torch image.")
+
+        dtype = t_image.dtype
+
+        if not t_image.is_floating_point():
+            t_image = t_image.to(torch.float32)
+
+        gaussian_noise = sigma * torch.randn_like(t_image) + self.mean
+        output = t_image + gaussian_noise
+
+        if output.dtype != dtype:
+            output = output.to(dtype)
+
+        if not isinstance(image, torch.Tensor):
+            output = F.to_pil_image(output)
+
+        return output
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}(mean={self.mean}, sigma={self.sigma})"
+        return s
+
+
 def _setup_size(size, error_msg):
     if isinstance(size, numbers.Number):
         return int(size), int(size)