diff --git a/test/test_transforms.py b/test/test_transforms.py index b9fa620e187..4d857ae484a 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -288,6 +288,15 @@ def test_pil_to_tensor(self, channels): expected_output = (input_data * 255).byte() torch.testing.assert_close(output, expected_output) + for mode in ("I;16", "I;16B"): + input_data = torch.tensor([[1, 258], [3, 4]], dtype=torch.int16) + array = input_data.numpy() + if mode == "I;16B": + array = array.byteswap() + img = Image.fromarray(array, mode=mode) + output = trans(img) + torch.testing.assert_close(input_data.unsqueeze(0), output) + # separate test for mode '1' PIL images input_data = torch.ByteTensor(1, height, width).bernoulli_() img = transforms.ToPILImage()(input_data.mul(255)).convert("1") diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py index 7b950b0c45b..4550ace4d24 100644 --- a/torchvision/transforms/functional.py +++ b/torchvision/transforms/functional.py @@ -206,7 +206,13 @@ def pil_to_tensor(pic: Any) -> Tensor: return torch.as_tensor(nppic) # handle PIL Image - img = torch.as_tensor(np.array(pic, copy=True)) + mode_to_nptype = { + "I": np.int32, + "I;16": np.int16, + "I;16B": np.int16, + "F": np.float32, + } + img = torch.as_tensor(np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True)) img = img.view(pic.size[1], pic.size[0], F_pil.get_image_num_channels(pic)) # put it from HWC to CHW format img = img.permute((2, 0, 1))