[test] added test for ToTensor transform on grayscale images

dee2c5de · Guillaume HEUSCH · 45246db3 · dee2c5de
Commit dee2c5de authored Feb 04, 2019 by Guillaume HEUSCH
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

bob/learn/pytorch/test/test.py bob/learn/pytorch/test/test.py +5 -0

No files found.
--- a/bob/learn/pytorch/test/test.py
+++ b/bob/learn/pytorch/test/test.py
@@ -111,6 +111,11 @@ def test_transforms():
  tt = ToTensor()
  tt(sample)
  assert isinstance(sample['image'], torch.Tensor)
+  # grayscale 
+  image_gray = numpy.random.rand(128, 128).astype("uint8")
+  sample_gray = {'image': image_gray}
+  tt(sample_gray)
+  assert isinstance(sample['image'], torch.Tensor)

  from ..datasets import Normalize
  image_copy = torch.Tensor(sample['image'])