diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index 6d9f9732552..69a2dec48c9 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -4142,9 +4142,21 @@ def adapter(_, input, __): check_transform(transforms.GaussianNoise(), make_input(dtype=torch.uint8), check_sample_input=adapter) + def test_pil_image(self): + image = make_image_pil() + out = F.gaussian_noise(image) + assert isinstance(out, PIL.Image.Image) + assert out.mode == image.mode + assert out.size == image.size + + out = F.gaussian_noise(image, mean=100, clip=True) + assert isinstance(out, PIL.Image.Image) + assert out.mode == image.mode + assert out.size == image.size + assert F.pil_to_tensor(out).min() >= 0 + assert F.pil_to_tensor(out).max() <= 255 + def test_bad_input(self): - with pytest.raises(ValueError, match="Gaussian Noise is not implemented for PIL images."): - F.gaussian_noise(make_image_pil()) with pytest.raises(ValueError, match="Input tensor is expected to be in uint8 or float dtype"): F.gaussian_noise(make_image(dtype=torch.int32)) with pytest.raises(ValueError, match="sigma shouldn't be negative"): diff --git a/torchvision/transforms/v2/_misc.py b/torchvision/transforms/v2/_misc.py index 305149c87b1..32229ee0a64 100644 --- a/torchvision/transforms/v2/_misc.py +++ b/torchvision/transforms/v2/_misc.py @@ -223,8 +223,8 @@ class GaussianNoise(Transform): noise added to each image will be different. The input tensor is also expected to be of float dtype in ``[0, 1]``, - or of ``uint8`` dtype in ``[0, 255]``. This transform does not support PIL - images. + or of ``uint8`` dtype in ``[0, 255]``. PIL images are supported and + processed through the ``uint8`` path via ``pil_to_tensor``. Regardless of the dtype used, the parameters of the function use the same scale, so a ``mean`` parameter of 0.5 will result in an average value diff --git a/torchvision/transforms/v2/functional/_misc.py b/torchvision/transforms/v2/functional/_misc.py index 55d085a2e24..cb9d7bf8701 100644 --- a/torchvision/transforms/v2/functional/_misc.py +++ b/torchvision/transforms/v2/functional/_misc.py @@ -226,9 +226,11 @@ def gaussian_noise_video(video: torch.Tensor, mean: float = 0.0, sigma: float = @_register_kernel_internal(gaussian_noise, PIL.Image.Image) def _gaussian_noise_pil( - video: torch.Tensor, mean: float = 0.0, sigma: float = 0.1, clip: bool = True + image: PIL.Image.Image, mean: float = 0.0, sigma: float = 0.1, clip: bool = True ) -> PIL.Image.Image: - raise ValueError("Gaussian Noise is not implemented for PIL images.") + t_img = pil_to_tensor(image) + output = gaussian_noise_image(t_img, mean=mean, sigma=sigma, clip=clip) + return to_pil_image(output, mode=image.mode) def to_dtype(inpt: torch.Tensor, dtype: torch.dtype = torch.float, scale: bool = False) -> torch.Tensor: