make convert_image_dtype scriptable

Author: nairbv

test/test_transforms.py

@@ -526,6 +526,10 @@ def test_to_tensor(self):
         output = trans(img)
         self.assertTrue(np.allclose(input_data.numpy(), output.numpy()))
 
+    def test_max_value(self):
+        for dtype in int_dtypes():
+            self.assertEqual(F._max_value(dtype), torch.iinfo(dtype).max)
+
     def test_convert_image_dtype_float_to_float(self):
         for input_dtype, output_dtypes in cycle_over(float_dtypes()):
             input_image = torch.tensor((0.0, 1.0), dtype=input_dtype)

torchvision/transforms/functional.py

@@ -124,7 +124,24 @@ def pil_to_tensor(pic):
     return img
 
 
-def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) -> torch.Tensor:
+# torch.iinfo isn't scriptable so using this helper function
+# https://github.com/pytorch/pytorch/issues/41492
+def _max_value(dtype: int) -> int:
+    a = torch.tensor(2, dtype=dtype)
+    signed = 1 if torch.tensor(0, dtype=dtype).is_signed() else 0
+    bits = 1
+    max_value = torch.tensor(-signed, dtype=torch.long)
+    while(True):
+        next_value = a.pow(bits - signed).sub(1)
+        if next_value > max_value:
+            max_value = next_value
+            bits *= 2
+        else:
+            return max_value.item()
+    return max_value.item()
+
+
+def convert_image_dtype(image: torch.Tensor, dtype: int = torch.float) -> torch.Tensor:
     """Convert a tensor image to the given ``dtype`` and scale the values accordingly
 
     Args:
@@ -148,9 +165,9 @@ def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) -
     if image.dtype == dtype:
         return image
 
-    if image.dtype.is_floating_point:
+    if torch.empty(0, dtype=image.dtype).is_floating_point():
         # float to float
-        if dtype.is_floating_point:
+        if torch.tensor(0, dtype=dtype).is_floating_point():
             return image.to(dtype)
 
         # float to int
@@ -166,19 +183,19 @@ def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) -
         # `max + 1 - epsilon` provides more evenly distributed mapping of
         # ranges of floats to ints.
         eps = 1e-3
-        result = image.mul(torch.iinfo(dtype).max + 1 - eps)
+        max_val = _max_value(dtype)
+        result = image.mul(max_val + 1.0 - eps)
         return result.to(dtype)
     else:
+        input_max = _max_value(image.dtype)
+        output_max = _max_value(dtype)
+
         # int to float
-        if dtype.is_floating_point:
-            max = torch.iinfo(image.dtype).max
+        if torch.tensor(0, dtype=dtype).is_floating_point():
             image = image.to(dtype)
-            return image / max
+            return image / input_max
 
         # int to int
-        input_max = torch.iinfo(image.dtype).max
-        output_max = torch.iinfo(dtype).max
-
         if input_max > output_max:
             factor = (input_max + 1) // (output_max + 1)
             image = image // factor

torchvision/transforms/functional_tensor.py

@@ -6,6 +6,8 @@
 from torch.nn.functional import affine_grid, grid_sample
 from torch.jit.annotations import List, BroadcastingList2
 
+import torchvision.transforms.functional as F
+
 
 def _is_tensor_a_torch_image(x: Tensor) -> bool:
     return x.ndim >= 2
@@ -228,13 +230,11 @@ def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor:
     result = img
     dtype = img.dtype
     if not torch.is_floating_point(img):
-        result = result / 255.0
+        result = F.convert_image_dtype(result, torch.get_default_dtype())
 
     result = (gain * result ** gamma).clamp(0, 1)
 
-    if result.dtype != dtype:
-        eps = 1e-3
-        result = (255 + 1.0 - eps) * result
+    result = F.convert_image_dtype(result, dtype)
     result = result.to(dtype)
     return result