Repository: Project-MONAI/MONAI

Issue:
update CVDecathlonDataset to use partition_dataset
**Is your feature request related to a problem? Please describe.**
As we already added `parittion_dataset` API, need to update the `CVDecathlonDataset` to use it.



Discussion:

Resolving patch:
diff --git a/monai/apps/datasets.py b/monai/apps/datasets.py
--- a/monai/apps/datasets.py
+++ b/monai/apps/datasets.py
@@ -15,8 +15,14 @@
 
 import numpy as np
 
-from monai.apps.utils import download_and_extract, split_dataset
-from monai.data import CacheDataset, load_decathlon_datalist, load_decathlon_properties
+from monai.apps.utils import download_and_extract
+from monai.data import (
+    CacheDataset,
+    load_decathlon_datalist,
+    load_decathlon_properties,
+    partition_dataset,
+    select_cross_validation_folds,
+)
 from monai.transforms import LoadNiftid, LoadPNGd, Randomizable
 from monai.utils import ensure_tuple
 
@@ -306,41 +312,34 @@ def _split_datalist(self, datalist: List[Dict]) -> List[Dict]:
             return [datalist[i] for i in self.indices]
 
 
-class CVDecathlonDataset:
+class CrossValidation:
     """
-    Cross vaidation dataset based on the general `DecathlonDataset`.
+    Cross validation dataset based on the general dataset which must have `_split_datalist` API.
 
     Args:
-        root_dir: user's local directory for caching and loading the MSD datasets.
-        task: which task to download and execute: one of list ("Task01_BrainTumour", "Task02_Heart",
-            "Task03_Liver", "Task04_Hippocampus", "Task05_Prostate", "Task06_Lung", "Task07_Pancreas",
-            "Task08_HepaticVessel", "Task09_Spleen", "Task10_Colon").
-        transform: transforms to execute operations on input data. the default transform is `LoadNiftid`,
-            which can load Nifti format data into numpy array with [H, W, D] or [H, W, D, C] shape.
-            for further usage, use `AddChanneld` or `AsChannelFirstd` to convert the shape to [C, H, W, D].
-        download: whether to download and extract the Decathlon from resource link, default is False.
-            if expected file already exists, skip downloading even set it to True.
+        dataset_cls: dataset class to be used to create the cross validation partitions.
+            It must have `_split_datalist` API.
+        nfolds: number of folds to split the data for cross validation.
         seed: random seed to randomly shuffle the datalist before splitting into N folds, default is 0.
-            note to set same seed for all the related datasets(`training` and `validation` sections in all N folds).
-        nsplits: number of folds, split the training data into N folds. each fold is then used once as a
-            validation while the N - 1 remaining folds form the training set.
-        cache_num: number of items to be cached. Default is `sys.maxsize`.
-            will take the minimum of (cache_num, data_length x cache_rate, data_length).
-        cache_rate: percentage of cached data in total, default is 1.0 (cache all).
-            will take the minimum of (cache_num, data_length x cache_rate, data_length).
-        num_workers: the number of worker threads to use.
-            if 0 a single thread will be used. Default is 0.
+        dataset_params: other additional parameters for the dataset_cls base class.
 
     Example of 5 folds cross validation training::
 
-        # use same random seed for all the datasets
-        cvdataset = CVDecathlonDataset(root_dir="./", task="Task09_Spleen", seed=12345, nsplts=5, download=True)
-        dataset_fold0_train = cvdataset.get_dataset(fold=0, section="training")
-        dataset_fold0_val = cvdataset.get_dataset(fold=0, section="validation")
+        cvdataset = CrossValidation(
+            dataset_cls=DecathlonDataset,
+            nfolds=5,
+            seed=12345,
+            root_dir="./",
+            task="Task09_Spleen",
+            section="training",
+            download=True,
+        )
+        dataset_fold0_train = cvdataset.get_dataset(folds=[1, 2, 3, 4])
+        dataset_fold0_val = cvdataset.get_dataset(folds=0)
         # execute training for fold 0 ...
 
-        dataset_fold1_train = cvdataset.get_dataset(fold=1, section="training")
-        dataset_fold1_val = cvdataset.get_dataset(fold=1, section="validation")
+        dataset_fold1_train = cvdataset.get_dataset(folds=[1])
+        dataset_fold1_val = cvdataset.get_dataset(folds=[0, 2, 3, 4])
         # execute training for fold 1 ...
 
         ...
@@ -352,63 +351,32 @@ class CVDecathlonDataset:
 
     def __init__(
         self,
-        root_dir: str,
-        task: str,
-        transform: Union[Sequence[Callable], Callable] = LoadNiftid(["image", "label"]),
-        download: bool = False,
+        dataset_cls,
+        nfolds: int = 5,
         seed: int = 0,
-        nsplits: int = 5,
-        cache_num: int = sys.maxsize,
-        cache_rate: float = 1.0,
-        num_workers: int = 0,
+        **dataset_params,
     ) -> None:
-        self.root_dir = root_dir
-        self.task = task
-        self.transform = transform
-        self.download = download
+        if not hasattr(dataset_cls, "_split_datalist"):
+            raise ValueError("dataset class must have _split_datalist API.")
+        self.dataset_cls = dataset_cls
+        self.nfolds = nfolds
         self.seed = seed
-        if nsplits < 2:
-            raise ValueError("nplits must be greater than 1 for cross validation.")
-        self.nsplits = nsplits
-        self.cache_num = cache_num
-        self.cache_rate = cache_rate
-        self.num_workers = num_workers
-
-    def get_dataset(self, fold: int, section: str):
+        self.dataset_params = dataset_params
+
+    def get_dataset(self, folds: Union[Sequence[int], int]):
         """
-        Generate dataset for specified fold index and section in the cross validation group.
+        Generate dataset based on the specified fold indice in the cross validation group.
 
         Args:
-            fold: index of fold for training and validation.
-            section: expected data section, can be: `training` or `validation`.
+            folds: index of folds for training or validation, if a list of values, concatenate the data.
 
         """
-        if fold > self.nsplits - 1:
-            raise ValueError("fold index should be in [0, nsplit - 1].")
-        nsplits = self.nsplits
+        nfolds = self.nfolds
+        seed = self.seed
 
-        class _NsplitsDataset(DecathlonDataset):
+        class _NsplitsDataset(self.dataset_cls):  # type: ignore
             def _split_datalist(self, datalist: List[Dict]) -> List[Dict]:
-                if section not in ["training", "validation"]:
-                    raise ValueError("section must be training or validation.")
-                length = len(datalist)
-                indices = np.arange(length)
-                self.randomize(indices)
-                fold_indices = split_dataset(nsplits, length)
-                if section == "training":
-                    self.indices = np.delete(indices, np.s_[fold_indices[fold][0] : fold_indices[fold][1]], axis=None)
-                else:
-                    self.indices = indices[fold_indices[fold][0] : fold_indices[fold][1]]
-                return [datalist[i] for i in self.indices]
-
-        return _NsplitsDataset(
-            root_dir=self.root_dir,
-            task=self.task,
-            section=section,
-            transform=self.transform,
-            download=self.download,
-            seed=self.seed,
-            cache_num=self.cache_num,
-            cache_rate=self.cache_rate,
-            num_workers=self.num_workers,
-        )
+                data = partition_dataset(data=datalist, num_partitions=nfolds, shuffle=True, seed=seed)
+                return select_cross_validation_folds(partitions=data, folds=folds)
+
+        return _NsplitsDataset(**self.dataset_params)
diff --git a/monai/apps/utils.py b/monai/apps/utils.py
--- a/monai/apps/utils.py
+++ b/monai/apps/utils.py
@@ -20,8 +20,6 @@
 from urllib.error import ContentTooShortError, HTTPError, URLError
 from urllib.request import Request, urlopen, urlretrieve
 
-import numpy as np
-
 from monai.utils import min_version, optional_import
 
 gdown, has_gdown = optional_import("gdown", "3.6")
@@ -251,22 +249,3 @@ def download_and_extract(
     """
     download_url(url=url, filepath=filepath, hash_val=hash_val, hash_type=hash_type)
     extractall(filepath=filepath, output_dir=output_dir, hash_val=hash_val, hash_type=hash_type)
-
-
-def split_dataset(nsplits: int, length: int):
-    """
-    Split dataset into N folds, return a list of indices for every fold.
-
-    Args:
-        nsplits: expected fold number.
-        length: dataset size in total.
-
-    """
-    sizes = np.full(nsplits, length // nsplits, dtype=np.int)
-    sizes[: length % nsplits] += 1
-    indices = list()
-    pos = 0
-    for size in sizes:
-        indices.append([pos, pos + size])
-        pos += size
-    return indices
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -704,24 +704,25 @@ def partition_dataset_classes(
     return datasets
 
 
-def generate_cross_validation_fold(partitions: Sequence[List], fold_idx: int):
+def select_cross_validation_folds(partitions: Sequence[List], folds: Union[Sequence[int], int]) -> List:
     """
-    Generate cross validation fold based on data partitions and specified fold index.
-    Will select the partition with fold index as validaiton dataset and concat others as train dataset.
-    For example, `partitions`: [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]] and `fold_idx`: 2
-    The the output will be [1, 2, 3, 4, 7, 8, 9, 10], [5, 6]
+    Select cross validation data based on data partitions and specified fold indice.
+    if a list of folds provided, concatenate the partitions of these folds.
+    For example, `partitions`: [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]] and `folds`: 2
+    The the output will be [5, 6], if `folds`: [1, 2], output will be [3, 4, 5, 6].
 
     """
-    if fold_idx >= len(partitions):
-        raise ValueError(f"fold_idx: {fold_idx} is bigger than number of partitions.")
+    folds = ensure_tuple(folds)
+    for fold in folds:
+        if fold >= len(partitions):
+            raise ValueError(f"fold index: {fold} is bigger than number of partitions.")
 
-    val_list = partitions[fold_idx]
-    train_list = list()
+    data_list = list()
     for i, data in enumerate(partitions):
-        if i != fold_idx:
-            train_list += data
+        if i in folds:
+            data_list += data
 
-    return train_list, val_list
+    return data_list
 
 
 class DistributedSampler(_TorchDistributedSampler):


Test patch:
diff --git a/tests/test_cvdecathlondataset.py b/tests/test_cross_validation.py
similarity index 80%
rename from tests/test_cvdecathlondataset.py
rename to tests/test_cross_validation.py
--- a/tests/test_cvdecathlondataset.py
+++ b/tests/test_cross_validation.py
@@ -13,12 +13,12 @@
 import unittest
 from urllib.error import ContentTooShortError, HTTPError
 
-from monai.apps import CVDecathlonDataset
+from monai.apps import CrossValidation, DecathlonDataset
 from monai.transforms import AddChanneld, Compose, LoadNiftid, ScaleIntensityd, ToTensord
 from tests.utils import skip_if_quick
 
 
-class TestCVDecathlonDataset(unittest.TestCase):
+class TestCrossValidation(unittest.TestCase):
     @skip_if_quick
     def test_values(self):
         testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data")
@@ -38,17 +38,19 @@ def _test_dataset(dataset):
             self.assertTrue("image_meta_dict" in dataset[0])
             self.assertTupleEqual(dataset[0]["image"].shape, (1, 34, 49, 41))
 
-        cvdataset = CVDecathlonDataset(
+        cvdataset = CrossValidation(
+            dataset_cls=DecathlonDataset,
+            nfolds=5,
+            seed=12345,
             root_dir=testing_dir,
             task="Task04_Hippocampus",
+            section="validation",
             transform=transform,
             download=True,
-            seed=12345,
-            nsplits=5,
         )
 
         try:  # will start downloading if testing_dir doesn't have the Decathlon files
-            data = cvdataset.get_dataset(fold=0, section="validation")
+            data = cvdataset.get_dataset(folds=0)
         except (ContentTooShortError, HTTPError, RuntimeError) as e:
             print(str(e))
             if isinstance(e, RuntimeError):
@@ -59,14 +61,14 @@ def _test_dataset(dataset):
         _test_dataset(data)
 
         # test training data for fold 0 of 5 splits
-        data = cvdataset.get_dataset(fold=0, section="training")
-        self.assertTupleEqual(data[0]["image"].shape, (1, 34, 48, 40))
+        data = cvdataset.get_dataset(folds=[1, 2, 3, 4])
+        self.assertTupleEqual(data[0]["image"].shape, (1, 35, 52, 33))
         self.assertEqual(len(data), 208)
         # test train / validation for fold 4 of 5 splits
-        data = cvdataset.get_dataset(fold=4, section="validation")
-        self.assertTupleEqual(data[0]["image"].shape, (1, 33, 55, 29))
+        data = cvdataset.get_dataset(folds=[4])
+        self.assertTupleEqual(data[0]["image"].shape, (1, 38, 53, 30))
         self.assertEqual(len(data), 52)
-        data = cvdataset.get_dataset(fold=4, section="training")
+        data = cvdataset.get_dataset(folds=[0, 1, 2, 3])
         self.assertTupleEqual(data[0]["image"].shape, (1, 34, 49, 41))
         self.assertEqual(len(data), 208)
 
diff --git a/tests/test_generate_cross_validation_fold.py b/tests/test_select_cross_validation_folds.py
similarity index 81%
rename from tests/test_generate_cross_validation_fold.py
rename to tests/test_select_cross_validation_folds.py
--- a/tests/test_generate_cross_validation_fold.py
+++ b/tests/test_select_cross_validation_folds.py
@@ -13,7 +13,7 @@
 
 from parameterized import parameterized
 
-from monai.data import generate_cross_validation_fold, partition_dataset
+from monai.data import partition_dataset, select_cross_validation_folds
 
 TEST_CASE_1 = [
     {
@@ -40,12 +40,13 @@
 ]
 
 
-class TestGenerateCrossValidationFold(unittest.TestCase):
+class TestSelectCrossValidationFolds(unittest.TestCase):
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
     def test_value(self, input_param, result):
         partitions = partition_dataset(**input_param)
-        train, val = generate_cross_validation_fold(partitions, fold_idx=1)
+        train = select_cross_validation_folds(partitions=partitions, folds=[0] + list(range(2, len(partitions))))
         self.assertListEqual(train, result[0])
+        val = select_cross_validation_folds(partitions=partitions, folds=1)
         self.assertListEqual(val, result[1])
 
 
diff --git a/tests/test_split_dataset.py b/tests/test_split_dataset.py
deleted file mode 100644
--- a/tests/test_split_dataset.py
+++ /dev/null
@@ -1,36 +0,0 @@
-# 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
-
-from parameterized import parameterized
-
-from monai.apps import split_dataset
-
-TEST_CASE_1 = [{"nsplits": 5, "length": 10}, [[0, 2], [2, 4], [4, 6], [6, 8], [8, 10]]]
-
-TEST_CASE_2 = [{"nsplits": 5, "length": 13}, [[0, 3], [3, 6], [6, 9], [9, 11], [11, 13]]]
-
-TEST_CASE_3 = [{"nsplits": 5, "length": 11}, [[0, 3], [3, 5], [5, 7], [7, 9], [9, 11]]]
-
-TEST_CASE_4 = [{"nsplits": 1, "length": 10}, [[0, 10]]]
-
-
-class TestSplitDataset(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
-    def test_shape(self, input_param, expected_indices):
-        indices = split_dataset(**input_param)
-        for data, expected in zip(indices, expected_indices):
-            self.assertListEqual(data, expected)
-
-
-if __name__ == "__main__":
-    unittest.main()