When loading data for train allow to read data that has channel on first position

test/data/test_image_collection_dataset.py

@@ -3,6 +3,8 @@
 from glob import glob
 from shutil import rmtree
 
+import tifffile
+import numpy as np
 from torch_em.util.test import create_image_collection_test_data
 
 
@@ -37,5 +39,49 @@ def test_dataset(self):
             self.assertEqual(y.shape, expected_shape)
 
 
+def generate_sample_data(folder, n_images, shape):
+    im_folder = os.path.join(folder, "images")
+    label_folder = os.path.join(folder, "labels")
+    os.makedirs(im_folder)
+    os.makedirs(label_folder)
+    for i in range(n_images):
+        raw = np.empty(shape, dtype=np.uint8)
+        label = np.ones(shape, dtype=np.float32)
+        tifffile.imwrite(os.path.join(im_folder, f"test_{i}.tif"), raw)
+        tifffile.imwrite(os.path.join(label_folder, f"test_{i}.tif"), label)
+
+
+class TestChannelsDataset(unittest.TestCase):
+    def test_channel_end(self):
+        from torch_em.data import ImageCollectionDataset
+
+        patch_shape = (256, 256)
+
+        with tempfile.TemporaryDirectory() as td:
+            im_folder = glob(os.path.join(td, "images", "*.tif"))
+            label_folder = glob(os.path.join(td, "labels", "*.tif"))
+
+            generate_sample_data(td, 10, (64, 64, 2))
+            ds = ImageCollectionDataset(raw_paths, label_paths,
+                                    patch_shape=patch_shape)
+            self.assertEqual(len(ds), 10)
+            self.assertEqual(ds._get_sample(0)[0].shape[0], 2)
+
+    def test_channel_begin(self):
+        from torch_em.data import ImageCollectionDataset
+
+        patch_shape = (256, 256)
+
+        with tempfile.TemporaryDirectory() as td:
+            im_folder = glob(os.path.join(td, "images", "*.tif"))
+            label_folder = glob(os.path.join(td, "labels", "*.tif"))
+
+            generate_sample_data(td, 10, (2, 64, 64))
+            ds = ImageCollectionDataset(raw_paths, label_paths,
+                                    patch_shape=patch_shape)
+            self.assertEqual(len(ds), 10)
+            self.assertEqual(ds._get_sample(0)[0].shape[0], 2)
+
+
 if __name__ == '__main__':
     unittest.main()

torch_em/data/image_collection_dataset.py

@@ -102,17 +102,23 @@ def _get_sample(self, index):
 
         shape = raw.shape
         # we assume images are loaded with channel last!
+        prefix_box = tuple()
         if have_raw_channels:
-            shape = shape[:-1]
+            if shape[-1] < 16:
+                shape = shape[:-1]
+            else:
+                shape = shape[1:]
+                prefix_box = (slice(None), )
+
 
         # sample random bounding box for this image
         bb = self._sample_bounding_box(shape)
 
-        raw = np.array(raw[bb])
+        raw = np.array(raw[prefix_box + bb])
         label = np.array(label[bb])
 
         # to channel first
-        if have_raw_channels:
+        if have_raw_channels and len(prefix_box) == 0:
             raw = raw.transpose((2, 0, 1))
 
         return raw, label