add torchscriptable adjust_gamma transform

#1375

Author: nairbv

test/test_functional_tensor.py

@@ -23,6 +23,8 @@ def _create_data(self, height=3, width=3, channels=3):
 
     def compareTensorToPIL(self, tensor, pil_image, msg=None):
         pil_tensor = torch.as_tensor(np.array(pil_image).transpose((2, 0, 1)))
+        if msg is None:
+            msg = "tensor:\n{} \ndid not equal PIL tensor:\n{}".format(tensor, pil_tensor)
         self.assertTrue(tensor.equal(pil_tensor), msg)
 
     def approxEqualTensorToPIL(self, tensor, pil_image, tol=1e-5, msg=None):
@@ -293,6 +295,33 @@ def test_pad(self):
         with self.assertRaises(ValueError, msg="Padding can not be negative for symmetric padding_mode"):
             F_t.pad(tensor, (-2, -3), padding_mode="symmetric")
 
+    def test_adjust_gamma(self):
+        script_fn = torch.jit.script(F_t.adjust_gamma)
+        tensor, pil_img = self._create_data(26, 36)
+
+        for dt in [torch.float64, torch.float32, None]:
+
+            if dt is not None:
+                tensor = F.convert_image_dtype(tensor, dt)
+
+            gammas = [0.8, 1.0, 1.2]
+            gains = [0.7, 1.0, 1.3]
+            for gamma, gain in zip(gammas, gains):
+
+                adjusted_tensor = F_t.adjust_gamma(tensor, gamma, gain)
+                adjusted_pil = F_pil.adjust_gamma(pil_img, gamma, gain)
+                scripted_result = script_fn(tensor, gamma, gain)
+                self.assertEqual(adjusted_tensor.dtype, scripted_result.dtype)
+                self.assertEqual(adjusted_tensor.size()[1:], adjusted_pil.size[::-1])
+
+                rbg_tensor = adjusted_tensor
+                if adjusted_tensor.dtype != torch.uint8:
+                    rbg_tensor = F.convert_image_dtype(adjusted_tensor, torch.uint8)
+
+                self.compareTensorToPIL(rbg_tensor, adjusted_pil)
+
+                self.assertTrue(adjusted_tensor.equal(scripted_result))
+
     def test_resize(self):
         script_fn = torch.jit.script(F_t.resize)
         tensor, pil_img = self._create_data(26, 36)

torchvision/transforms/functional.py

@@ -161,8 +161,8 @@ def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) -
             msg = f"The cast from {image.dtype} to {dtype} cannot be performed safely."
             raise RuntimeError(msg)
 
-        eps = 1e-3
-        return image.mul(torch.iinfo(dtype).max + 1 - eps).to(dtype)
+        max = torch.iinfo(dtype).max
+        return image.mul(torch.iinfo(dtype).max).clamp(0, max).to(dtype)
     else:
         # int to float
         if dtype.is_floating_point:
@@ -760,7 +760,7 @@ def adjust_hue(img: Tensor, hue_factor: float) -> Tensor:
     raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
 
 
-def adjust_gamma(img, gamma, gain=1):
+def adjust_gamma(img, gamma: float, gain: float = 1):
     r"""Perform gamma correction on an image.
 
     Also known as Power Law Transform. Intensities in RGB mode are adjusted
@@ -774,26 +774,16 @@ def adjust_gamma(img, gamma, gain=1):
     .. _Gamma Correction: https://en.wikipedia.org/wiki/Gamma_correction
 
     Args:
-        img (PIL Image): PIL Image to be adjusted.
+        img (PIL Image or Tensor): PIL Image to be adjusted.
         gamma (float): Non negative real number, same as :math:`\gamma` in the equation.
             gamma larger than 1 make the shadows darker,
             while gamma smaller than 1 make dark regions lighter.
         gain (float): The constant multiplier.
     """
-    if not F_pil._is_pil_image(img):
-        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
-
-    if gamma < 0:
-        raise ValueError('Gamma should be a non-negative real number')
+    if F_pil._is_pil_image(img):
+        return F_pil.adjust_gamma(img, gamma, gain)
 
-    input_mode = img.mode
-    img = img.convert('RGB')
-
-    gamma_map = [255 * gain * pow(ele / 255., gamma) for ele in range(256)] * 3
-    img = img.point(gamma_map)  # use PIL's point-function to accelerate this part
-
-    img = img.convert(input_mode)
-    return img
+    return F_t.adjust_gamma(img, gamma, gain)
 
 
 def rotate(img, angle, resample=False, expand=False, center=None, fill=None):

torchvision/transforms/functional_pil.py

@@ -164,6 +164,43 @@ def adjust_hue(img, hue_factor):
     return img
 
 
+@torch.jit.unused
+def adjust_gamma(img, gamma, gain=1):
+    r"""Perform gamma correction on an image.
+
+    Also known as Power Law Transform. Intensities in RGB mode are adjusted
+    based on the following equation:
+
+    .. math::
+        I_{\text{out}} = 255 \times \text{gain} \times \left(\frac{I_{\text{in}}}{255}\right)^{\gamma}
+
+    See `Gamma Correction`_ for more details.
+
+    .. _Gamma Correction: https://en.wikipedia.org/wiki/Gamma_correction
+
+    Args:
+        img (PIL Image): PIL Image to be adjusted.
+        gamma (float): Non negative real number, same as :math:`\gamma` in the equation.
+            gamma larger than 1 make the shadows darker,
+            while gamma smaller than 1 make dark regions lighter.
+        gain (float): The constant multiplier.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    if gamma < 0:
+        raise ValueError('Gamma should be a non-negative real number')
+
+    input_mode = img.mode
+    img = img.convert('RGB')
+
+    gamma_map = [255 * gain * pow(ele / 255., gamma) for ele in range(256)] * 3
+    img = img.point(gamma_map)  # use PIL's point-function to accelerate this part
+
+    img = img.convert(input_mode)
+    return img
+
+
 @torch.jit.unused
 def pad(img, padding, fill=0, padding_mode="constant"):
     r"""Pad the given PIL.Image on all sides with the given "pad" value.

torchvision/transforms/functional_tensor.py

@@ -194,6 +194,44 @@ def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor:
     return _blend(img, rgb_to_grayscale(img), saturation_factor)
 
 
+def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor:
+    r"""Adjust gamma of an RGB image.
+
+    Also known as Power Law Transform. Intensities in RGB mode are adjusted
+    based on the following equation:
+
+    .. math::
+        `I_{\text{out}} = 255 \times \text{gain} \times \left(\frac{I_{\text{in}}}{255}\right)^{\gamma}`
+
+    See `Gamma Correction`_ for more details.
+
+    .. _Gamma Correction: https://en.wikipedia.org/wiki/Gamma_correction
+
+    Args:
+        img (Tensor): PIL Image to be adjusted.
+        gamma (float): Non negative real number, same as :math:`\gamma` in the equation.
+            gamma larger than 1 make the shadows darker,
+            while gamma smaller than 1 make dark regions lighter.
+        gain (float): The constant multiplier.
+    """
+
+    if not isinstance(img, torch.Tensor):
+        raise TypeError('img should be a Tensor. Got {}'.format(type(img)))
+
+    if gamma < 0:
+        raise ValueError('Gamma should be a non-negative real number')
+
+    result = img
+    dtype = img.dtype
+    if torch.is_floating_point(img):
+        return gain * result ** gamma
+
+    result = 255.0 * gain * (result / 255.0) ** gamma
+    # PIL clamps, to(torch.uint8) would wrap
+    result = result.clamp(0, 255).to(dtype)
+    return result
+
+
 def center_crop(img: Tensor, output_size: BroadcastingList2[int]) -> Tensor:
     """Crop the Image Tensor and resize it to desired size.