Unverified Commit 9b526dbc authored by Francisco Massa's avatar Francisco Massa Committed by GitHub
Browse files

Fix ToTensor when PIL Image has mode F

Fixes https://github.com/pytorch/vision/issues/399.

The only case of floating point supported by PIL seems to be `F`, so this should fix it.
parent 23ea65cd
...@@ -59,6 +59,8 @@ def to_tensor(pic): ...@@ -59,6 +59,8 @@ def to_tensor(pic):
img = torch.from_numpy(np.array(pic, np.int32, copy=False)) img = torch.from_numpy(np.array(pic, np.int32, copy=False))
elif pic.mode == 'I;16': elif pic.mode == 'I;16':
img = torch.from_numpy(np.array(pic, np.int16, copy=False)) 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))
else: else:
img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
# PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
......
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