Commit 22385bc6 authored by vfdev's avatar vfdev Committed by Francisco Massa
Browse files

Fix bug with to_tensor on the input of ndarray of float32 (#415)

parent ba7eb9f4
...@@ -269,11 +269,16 @@ class Tester(unittest.TestCase): ...@@ -269,11 +269,16 @@ class Tester(unittest.TestCase):
output = trans(img) output = trans(img)
assert np.allclose(input_data.numpy(), output.numpy()) assert np.allclose(input_data.numpy(), output.numpy())
ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)) ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
output = trans(ndarray) output = trans(ndarray)
expected_output = ndarray.transpose((2, 0, 1)) / 255.0 expected_output = ndarray.transpose((2, 0, 1)) / 255.0
assert np.allclose(output.numpy(), expected_output) assert np.allclose(output.numpy(), expected_output)
ndarray = np.random.rand(height, width, channels).astype(np.float32)
output = trans(ndarray)
expected_output = ndarray.transpose((2, 0, 1))
assert np.allclose(output.numpy(), expected_output)
@unittest.skipIf(accimage is None, 'accimage not available') @unittest.skipIf(accimage is None, 'accimage not available')
def test_accimage_to_tensor(self): def test_accimage_to_tensor(self):
trans = transforms.ToTensor() trans = transforms.ToTensor()
......
...@@ -47,7 +47,10 @@ def to_tensor(pic): ...@@ -47,7 +47,10 @@ def to_tensor(pic):
# handle numpy array # handle numpy array
img = torch.from_numpy(pic.transpose((2, 0, 1))) img = torch.from_numpy(pic.transpose((2, 0, 1)))
# backward compatibility # backward compatibility
return img.float().div(255) if isinstance(img, torch.ByteTensor):
return img.float().div(255)
else:
return img
if accimage is not None and isinstance(pic, accimage.Image): if accimage is not None and isinstance(pic, accimage.Image):
nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32) nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)
......
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