Commit f6ab1079 authored by arturml's avatar arturml Committed by Francisco Massa
Browse files

Add support in transforms.ToTensor for PIL Images mod '1' (#471)

* Add case in test_to_tensor for PIL Images mode '1'

* Add support in ToTensor for PIL Images mode '1'

* Fix pep8 issues
parent 7bda0e85
......@@ -404,6 +404,12 @@ class Tester(unittest.TestCase):
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()
......
......@@ -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':
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment