Added transform to subselect channels from a multichannel image

bob/learn/pytorch/datasets/__init__.py

@@ -9,7 +9,8 @@ from .utils import RollChannels
 from .utils import ToTensor 
 from .utils import Normalize 
 from .utils import Resize 
-from .utils import ToGray 
+from .utils import ToGray
+from .utils import ChannelSelect 
 
 from .utils import map_labels
 from .utils import ConcatDataset

bob/learn/pytorch/datasets/utils.py

@@ -199,4 +199,44 @@ class ConcatDataset(Dataset):
 
       return sample
 
+class ChannelSelect(object):
+
+  """Subselects or re-orders channels in a multi-channel image. Expects a 
+  numpy.ndarray as input with size `HxWxnum_channels` and returns an image 
+  with size `HxWxlen(selected_channels)`, where the last dimension is subselected
+  using the indexes in the list `selected_channels`. 
+
+  Attributes
+  ----------
+
+  selected_channels: list
+    The indexes of the channels to be selected.
+
+  img: numpy.ndarray
+    A multi channel image, HxWxnum_channels
+  """
+
+  def __init__(self, selected_channels=[0,1,2,3]):
+
+    """
+    Parameters
+    ----------
+    selected_channels: list
+      The indexes of the channels to be selected.
+
+    """
+    self.selected_channels = selected_channels
+
+  def __call__(self, img):
+    """
+    Parameters
+    ----------
+    img: numpy.ndarray
+      A multi channel image, HxWxnum_channels
+
+    """
+    return img[:,:,self.selected_channels]
+
+  def __repr__(self):
+    return self.__class__.__name__ + '(selected_channels={}, output channels ={})'.format(self.selected_channels, len(self.selected_channels))
 

bob/learn/pytorch/test/test.py

@@ -109,6 +109,12 @@ def test_transforms():
 
   image = numpy.random.rand(3, 128, 128).astype("uint8")
 
+  img = numpy.random.rand(128, 128, 4).astype("uint8")
+
+  from ..datasets import ChannelSelect
+  cs = ChannelSelect(selected_channels=[0,1,2])
+  assert(cs(img).shape==(128,128,3))
+
   from ..datasets import RollChannels
   sample = {'image': image}
   rc = RollChannels()