From c20f6b76a0e55721c27da787f6d25007e8091338 Mon Sep 17 00:00:00 2001 From: Artur Lacerda Date: Sat, 14 Apr 2018 16:03:27 -0300 Subject: [PATCH 1/3] Add case in test_to_tensor for PIL Images mode '1' --- test/test_transforms.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/test/test_transforms.py b/test/test_transforms.py index e2232e2491b..19ff522873c 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -404,6 +404,12 @@ def test_to_tensor(self): expected_output = ndarray.transpose((2, 0, 1)) assert np.allclose(output.numpy(), expected_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') + output = trans(img) + assert np.allclose(input_data.numpy(), output.numpy()) + @unittest.skipIf(accimage is None, 'accimage not available') def test_accimage_to_tensor(self): trans = transforms.ToTensor() From 3a5bc4ccb844a324589808b7bb2f31d7f7be771a Mon Sep 17 00:00:00 2001 From: Artur Lacerda Date: Sat, 14 Apr 2018 16:04:56 -0300 Subject: [PATCH 2/3] Add support in ToTensor for PIL Images mode '1' --- torchvision/transforms/functional.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py index 5d5325078be..7b090b6262c 100644 --- a/torchvision/transforms/functional.py +++ b/torchvision/transforms/functional.py @@ -64,9 +64,11 @@ def to_tensor(pic): img = torch.from_numpy(np.array(pic, np.int16, copy=False)) elif pic.mode == 'F': img = torch.from_numpy(np.array(pic, np.float32, copy=False)) + elif pic.mode == '1': + img = 255*torch.from_numpy(np.array(pic, np.uint8, copy=False)) else: img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) - # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK + # PIL image mode: L, P, I, F, RGB, YCbCr, RGBA, CMYK if pic.mode == 'YCbCr': nchannel = 3 elif pic.mode == 'I;16': From f0708406a5b6d4a0664e32b484fe8f7dc9ec7ca6 Mon Sep 17 00:00:00 2001 From: Artur Lacerda Date: Mon, 16 Apr 2018 10:42:47 -0300 Subject: [PATCH 3/3] Fix pep8 issues --- test/test_transforms.py | 2 +- torchvision/transforms/functional.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_transforms.py b/test/test_transforms.py index 19ff522873c..21305f5fa8f 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -404,7 +404,7 @@ def test_to_tensor(self): expected_output = ndarray.transpose((2, 0, 1)) assert np.allclose(output.numpy(), expected_output) - #separate test for mode '1' PIL images + # separate test for mode '1' PIL images input_data = torch.ByteTensor(1, height, width).bernoulli_() img = transforms.ToPILImage()(input_data.mul(255)).convert('1') output = trans(img) diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py index 7b090b6262c..7d1a8921a6b 100644 --- a/torchvision/transforms/functional.py +++ b/torchvision/transforms/functional.py @@ -65,7 +65,7 @@ def to_tensor(pic): elif pic.mode == 'F': img = torch.from_numpy(np.array(pic, np.float32, copy=False)) elif pic.mode == '1': - img = 255*torch.from_numpy(np.array(pic, np.uint8, copy=False)) + img = 255 * torch.from_numpy(np.array(pic, np.uint8, copy=False)) else: img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) # PIL image mode: L, P, I, F, RGB, YCbCr, RGBA, CMYK