Repository: Project-MONAI/MONAI

Issue:
add ZipDataset to combine datasets
**Is your feature request related to a problem? Please describe.**
In order to combine several datasets and load in DataLoader together, we need to add ZipDataset.


Discussion:
Nic-Ma: /black
Nic-Ma: /black
Nic-Ma: Hi @wyli ,

Thanks for your review.
I updated the PR according to your comments, could you please help review it again?
Thanks.

Patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
index fda70efa..b78ab357 100644
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from .csv_saver import CSVSaver
-from .dataset import Dataset, PersistentDataset, CacheDataset
+from .dataset import Dataset, PersistentDataset, CacheDataset, ZipDataset, ArrayDataset
 from .grid_dataset import GridPatchDataset
 from .nifti_reader import load_nifti, NiftiDataset
 from .nifti_saver import NiftiSaver
diff --git a/monai/data/dataset.py b/monai/data/dataset.py
index 23be6d53..08b36551 100644
--- a/monai/data/dataset.py
+++ b/monai/data/dataset.py
@@ -23,12 +23,12 @@ import threading
 from monai.transforms import Compose, Randomizable
 from monai.transforms.utils import apply_transform
 from monai.utils import process_bar
-from monai.utils.misc import ensure_tuple
 
 
 class Dataset(torch.utils.data.Dataset):
     """
-    Generic dataset to handle dictionary format data, it can operate transforms for specific fields.
+    A generic dataset with a length property and an optional callable data transform
+    when fetching a data sample.
     For example, typical input data can be a list of dictionaries::
 
         [{                            {                            {
@@ -42,13 +42,10 @@ class Dataset(torch.utils.data.Dataset):
         """
         Args:
             data (Iterable): input data to load and transform to generate dataset for model.
-            transform (Callable, optional): transforms to execute operations on input data.
+            transform (Callable, optional): a callable data transform on input data.
         """
         self.data = data
-        if isinstance(transform, Compose):
-            self.transform = transform
-        else:
-            self.transform = Compose(ensure_tuple(transform))
+        self.transform = transform
 
     def __len__(self):
         return len(self.data)
@@ -107,6 +104,8 @@ class PersistentDataset(Dataset):
                 may share a common cache dir provided that the trasnsforms pre-processing is
                 consistent.
         """
+        if not isinstance(transform, Compose):
+            transform = Compose(transform)
         super().__init__(data=data, transform=transform)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else None
 
@@ -289,3 +288,105 @@ class CacheDataset(Dataset):
             # no cache for this data, execute all the transforms directly
             data = super(CacheDataset, self).__getitem__(index)
         return data
+
+
+class ZipDataset(torch.utils.data.Dataset):
+    """
+    Zip several PyTorch datasets and output data(with the same index) together in a tuple.
+    If the output of single dataset is alreay a tuple, flatten it and extend to the result.
+    For example: if datasetA returns (img, imgmeta), datsetB returns (seg, segmeta),
+    finally return (img, imgmeta, seg, segmeta).
+
+    Note:
+        Expect all the datasets have same length.
+
+    """
+
+    def __init__(self, datasets, transform=None):
+        """
+        Args:
+            datasets (list or tuple): list of datasets to zip together.
+        """
+        self.datasets = list(datasets)
+        self.len = min([len(dataset) for dataset in self.datasets])
+        self.transform = transform
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, index):
+        def to_list(x):
+            return list(x) if isinstance(x, (tuple, list)) else [x]
+
+        data = list()
+        for dataset in self.datasets:
+            data.extend(to_list(dataset[index]))
+        if self.transform is not None:
+            data = self.transform(data)
+        return data
+
+
+class ArrayDataset(ZipDataset, Randomizable):
+    """
+    Dataset for segmentation and classification tasks based on array format input data and transforms.
+    It can apply same random operations for both image transforms and segmentation label transforms.
+    The `transform` can be :py:class:`monai.transforms.Compose` or any other callable object.
+    For example:
+    If train based on Nifti format images without metadata, all transforms can be composed::
+        img_transform = Compose(
+            [
+                LoadNifti(image_only=True),
+                AddChannel(),
+                RandAdjustContrast()
+            ]
+        )
+    If train based on Nifti format images and the metadata, the array transforms can not be composed
+    because several transforms receives multiple parameters or return multiple values. Then Users need
+    to define their own callable method to parse metadata from `LoadNifti` or set `affine` matrix
+    to `Spacing` transform::
+        class TestCompose(Compose):
+            def __call__(self, input_):
+                img, metadata = self.transforms[0](input_)
+                img = self.transforms[1](img)
+                img, _, _ = self.transforms[2](img, metadata["affine"])
+                return self.transforms[3](img), metadata
+        img_transform = TestCompose(
+            [
+                LoadNifti(image_only=False),
+                AddChannel(),
+                Spacing(pixdim=(1.5, 1.5, 3.0)),
+                RandAdjustContrast()
+            ]
+        )
+    Recommend to use dictionary Datasets for complicated data pre-processing.
+    """
+
+    def __init__(
+        self, img_files, img_transform=None, seg_files=None, seg_transform=None, labels=None, label_transform=None
+    ):
+        """
+        Initializes the dataset with the filename lists. The transform `img_transform` is applied
+        to the images and `seg_transform` to the segmentations.
+        Args:
+            img_files (iterable, list of str): list of image filenames
+            img_transform (Callable, optional): transform to apply to image arrays
+            seg_files (iterable, list of str): if in segmentation task, list of segmentation filenames
+            seg_transform (Callable, optional): transform to apply to segmentation arrays
+            labels (iterable, list or array): if in classification task, list of classification labels
+            label_transform (Callable, optional): transform to apply to label arrays
+
+        """
+
+        super().__init__(
+            [Dataset(img_files, img_transform), Dataset(seg_files, seg_transform), Dataset(labels, label_transform)]
+        )
+
+    def randomize(self):
+        self.seed = self.R.randint(2147483647)
+
+    def __getitem__(self, index):
+        self.randomize()
+        for dataset in self.datasets:
+            if isinstance(dataset.transform, Randomizable):
+                dataset.transform.set_random_state(seed=self.seed)
+        return super().__getitem__(index)


Test patch:
diff --git a/tests/test_arraydataset.py b/tests/test_arraydataset.py
new file mode 100644
index 00000000..f6dca3e9
--- /dev/null
+++ b/tests/test_arraydataset.py
@@ -0,0 +1,74 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import numpy as np
+import tempfile
+import nibabel as nib
+from parameterized import parameterized
+from monai.data import ArrayDataset
+from monai.transforms import Compose, LoadNifti, AddChannel, RandAdjustContrast, Spacing
+
+TEST_CASE_1 = [
+    Compose([LoadNifti(image_only=True), AddChannel(), RandAdjustContrast()]),
+    Compose([LoadNifti(image_only=True), AddChannel(), RandAdjustContrast()]),
+    (0, 1),
+    (1, 128, 128, 128),
+]
+
+
+class TestCompose(Compose):
+    def __call__(self, input_):
+        img, metadata = self.transforms[0](input_)
+        img = self.transforms[1](img)
+        img, _, _ = self.transforms[2](img, metadata["affine"])
+        return self.transforms[3](img), metadata
+
+
+TEST_CASE_2 = [
+    TestCompose([LoadNifti(image_only=False), AddChannel(), Spacing(pixdim=(2, 2, 4)), RandAdjustContrast()]),
+    TestCompose([LoadNifti(image_only=False), AddChannel(), Spacing(pixdim=(2, 2, 4)), RandAdjustContrast()]),
+    (0, 2),
+    (1, 64, 64, 33),
+]
+
+
+class TestArrayDataset(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_shape(self, img_transform, label_transform, indexes, expected_shape):
+        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
+        tempdir = tempfile.mkdtemp()
+        test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
+        test_seg1 = os.path.join(tempdir, "test_seg1.nii.gz")
+        test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
+        test_seg2 = os.path.join(tempdir, "test_seg2.nii.gz")
+        nib.save(test_image, test_image1)
+        nib.save(test_image, test_seg1)
+        nib.save(test_image, test_image2)
+        nib.save(test_image, test_seg2)
+        test_images = [test_image1, test_image2]
+        test_segs = [test_seg1, test_seg2]
+        test_labels = [1, 1]
+        dataset = ArrayDataset(test_images, img_transform, test_segs, label_transform, test_labels, None)
+        data1 = dataset[0]
+        data2 = dataset[1]
+
+        self.assertTupleEqual(data1[indexes[0]].shape, expected_shape)
+        self.assertTupleEqual(data1[indexes[1]].shape, expected_shape)
+        np.testing.assert_allclose(data1[indexes[0]], data1[indexes[1]])
+        self.assertTupleEqual(data2[indexes[0]].shape, expected_shape)
+        self.assertTupleEqual(data2[indexes[1]].shape, expected_shape)
+        np.testing.assert_allclose(data2[indexes[0]], data2[indexes[0]])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_zipdataset.py b/tests/test_zipdataset.py
new file mode 100644
index 00000000..09b28ab4
--- /dev/null
+++ b/tests/test_zipdataset.py
@@ -0,0 +1,56 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import torch
+from parameterized import parameterized
+from monai.data import ZipDataset
+
+
+class Dataset_(torch.utils.data.Dataset):
+    def __init__(self, length, index_only=True):
+        self.len = length
+        self.index_only = index_only
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, index):
+        if self.index_only:
+            return index
+        else:
+            return 1, 2, index
+
+
+TEST_CASE_1 = [[Dataset_(5), Dataset_(5), Dataset_(5)], None, [0, 0, 0], 5]
+
+TEST_CASE_2 = [[Dataset_(3), Dataset_(4), Dataset_(5)], None, [0, 0, 0], 3]
+
+TEST_CASE_3 = [[Dataset_(3), Dataset_(4, index_only=False), Dataset_(5)], None, [0, 1, 2, 0, 0], 3]
+
+TEST_CASE_4 = [
+    [Dataset_(3), Dataset_(4, index_only=False), Dataset_(5)],
+    lambda x: [i + 1 for i in x],
+    [1, 2, 3, 1, 1],
+    3,
+]
+
+
+class TestZipDataset(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    def test_value(self, datasets, transform, expected_output, expected_length):
+        test_dataset = ZipDataset(datasets=datasets, transform=transform)
+        self.assertEqual(test_dataset[0], expected_output)
+        self.assertEqual(len(test_dataset), expected_length)
+
+
+if __name__ == "__main__":
+    unittest.main()