Merge branch 'to-tensor-gray' into 'master'

modified the ToTensor transform to properly work with grasyscale images

See merge request !15

bob/learn/pytorch/datasets/utils.py

@@ -98,6 +98,8 @@ class ToTensor(object):
     self.op = transforms.ToTensor()
 
   def __call__(self, sample):
+    if len(sample['image'].shape) == 2:
+      sample['image'] = sample['image'][..., numpy.newaxis]
     sample['image'] = self.op(sample['image'])
     return sample
 

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'])