diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst index 529815ead9a..bd2b9cf540d 100644 --- a/docs/source/transforms.rst +++ b/docs/source/transforms.rst @@ -102,6 +102,13 @@ that can be represented in that dtype. Typically, images of dtype Use :class:`~torchvision.transforms.v2.ToDtype` to convert both the dtype and range of the inputs. +.. note:: + + ``torch.uint16``, ``torch.uint32``, and ``torch.uint64`` dtypes are not + officially supported by the torchvision transforms. While some operations + may work, most transforms expect ``torch.uint8`` or ``torch.float32`` + inputs. + .. _v1_or_v2: V1 or V2? Which one should I use? diff --git a/test/test_transforms.py b/test/test_transforms.py index d93800d59bc..5b2658c99a4 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -614,8 +614,8 @@ def _get_1_channel_tensor_various_types(): expected_output = img_data_byte.float().div(255.0).numpy() yield img_data_byte, expected_output, "L" - img_data_short = torch.ShortTensor(1, 4, 4).random_() - expected_output = img_data_short.numpy() + img_data_short = torch.ShortTensor(1, 4, 4).random_(0, 32767) + expected_output = img_data_short.numpy().astype(np.uint16) yield img_data_short, expected_output, "I;16" if sys.byteorder == "little" else "I;16B" img_data_int = torch.IntTensor(1, 4, 4).random_() @@ -631,8 +631,8 @@ def _get_2d_tensor_various_types(): expected_output = img_data_byte.float().div(255.0).numpy() yield img_data_byte, expected_output, "L" - img_data_short = torch.ShortTensor(4, 4).random_() - expected_output = img_data_short.numpy() + img_data_short = torch.ShortTensor(4, 4).random_(0, 32767) + expected_output = img_data_short.numpy().astype(np.uint16) yield img_data_short, expected_output, "I;16" if sys.byteorder == "little" else "I;16B" img_data_int = torch.IntTensor(4, 4).random_() diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index a87e601e1a6..a187811e340 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -6825,6 +6825,18 @@ def test_functional_error(self): F.pil_to_tensor(object()) +@pytest.mark.parametrize("f", [F.to_tensor, F.pil_to_tensor]) +def test_I16_to_tensor(f): + # See https://github.com/pytorch/vision/issues/8188 + I16_pil_img = PIL.Image.fromarray(np.random.randint(0, 2**16, (10, 10), dtype=np.uint16)) + assert I16_pil_img.mode == "I;16" + + cm = pytest.warns(UserWarning, match="deprecated") if f is F.to_tensor else contextlib.nullcontext() + with cm: + out = f(I16_pil_img) + assert out.dtype == torch.uint16 + + class TestLambda: @pytest.mark.parametrize("input", [object(), torch.empty(()), np.empty(()), "string", 1, 0.0]) @pytest.mark.parametrize("types", [(), (torch.Tensor, np.ndarray)]) diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py index 7b950b0c45b..dacdd283bfb 100644 --- a/torchvision/transforms/functional.py +++ b/torchvision/transforms/functional.py @@ -164,7 +164,7 @@ def to_tensor(pic: Union[PILImage, np.ndarray]) -> Tensor: return torch.from_numpy(nppic).to(dtype=default_float_dtype) # handle PIL Image - mode_to_nptype = {"I": np.int32, "I;16" if sys.byteorder == "little" else "I;16B": np.int16, "F": np.float32} + mode_to_nptype = {"I": np.int32, "I;16" if sys.byteorder == "little" else "I;16B": np.uint16, "F": np.float32} img = torch.from_numpy(np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True)) if pic.mode == "1":