Repository: Project-MONAI/MONAI

Issue:
image data reader and LoadImage transform refactoring
This is a follow-up of #856 to track the progress of a preliminary IO refactoring.

**Goal**

The main goal of this ticket is to create a file format agnostic wrapper for the image inputs (`LoadImage` and `LoadImaged`).
A typical use case (enhancing the current [`NiftiDataset`](https://github.com/Project-MONAI/MONAI/blob/6a6f52f5f80e5c706f1f0976aa76b56af488d119/monai/data/nifti_reader.py#L21)) is:
```python
train_ds = ImageDataset(images, segs, transform=train_imtrans, seg_transform=train_segtrans)
```
where `images` and `segs` are filenames, `ImageDataset` internally invokes `LoadImage`.

Or using a generic monai Dataset API, specify the first component of `train_imtrans` as a `LoadImage` transform that extracts data array and meta data from the file.

More specifically:

**Reader** （file format-specific component)

1. `[format]reader` are themselves monai data components (users can use the readers without going through the `LoadImage` logic), they call external package APIs
2. `[format]reader` share the same superclass `ImageReader` interface (for handling physical/voxel resolution, orientation...)
```py
class ImageReader:
    def get_data(self, ...):
        # return the data in the designated resolution, orientation, 
        # image channel selector, the relevant coordinates, datatype
    
    def read(self, ...):
        # create an IO object instance or a set of instances
```
3. users can make customized `MyReader` by implementing the `ImageReader` interface
4. the customized `MyReader` could be registered with the `LoadImage` factory
```python
@register_reader(requires='my_external_library')
class MyReader(ImageReader):
    ...
```


**LoadImage**（file format agnostic component)

1. `LoadImage` is a monai transform
2. `LoadImaged` should return the same underlying data representation dictionary regardless of the image format.
3. `LoadImage` may choose a suitable `[format]reader` dynamically at runtime based on the following rules (ordered from high to low priority):
  - user's selection of a specific reader, such as `LoadImage(reader='ITkReader')`
  - try any user-registered readers
  - try resources map based on a global property file and filename extensions:
  ```
  MONAI.IO
  {
      FOO LoadFOO
      PNG  ITK/5.1/LoadITK
      TIFF,TIF  ITK/5.1/LoadITK
      * ITK/5.1/LoadITK
  }
  ```
  - fails early and preserves detailed error messages

**additional requirements**
- should be able to specify versions of external reader to ensure reproducibility
- need to have consistent coordinate systems, resolution conventions

**tasks**
- [x] implement interface `ImageReader`, factory (transforms) `LoadImage` `LoadImaged`
- [x] implement ITKReader, NibabelReader, NumpyReader, PILReader?
- [x] `LoadImage` selecting readers according to the preference list
- [x] new interfaces tutorial https://github.com/Project-MONAI/Tutorials/pull/19
- [x] deprecate `NiftiDataset`, add `ImageDataset`
- [x] deprecate `LoadNifti`, `LoadPNG`, `LoadNumpy`, `LoadDatad`
- [x] revise/check all examples demos


Discussion:
thanks for the discussion @Nic-Ma  let's use this ticket to raise any detailed technical issues
Mark: completed 2 tasks in the list, will work on `NumpyReader`, `PILReader` in the next PR.
Thanks.
before deprecating the previous loaders, would be very helpful to have a new tutorial to demonstrate the new interfaces, what do you think @Nic-Ma?
Hi @wyli ,

Sounds good a plan.
Please submit a ticket and assign it to me.

Thanks.
this ticket is currently blocked at https://github.com/InsightSoftwareConsortium/ITK/issues/2003 because the default is 
https://github.com/Project-MONAI/MONAI/blob/99f7a439020bac0c70de1a5096bb24aceeaf4ae5/monai/transforms/io/array.py#L62-L63
Here is the plan for this ticket in v0.4 development:
1. Change the default reader of `LoadImage` to Nibabel.
2. Add string support for the reader parameter.
3. Update all the examples and tutorials to use `LoadImage` transform.

Thanks.
Will deprecate `NiftiDataset`, add `ImageDataset` in another separate PR later.
Thanks.

Resolving patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -23,9 +23,9 @@
 )
 from .decathlon_datalist import load_decathlon_datalist, load_decathlon_properties
 from .grid_dataset import GridPatchDataset, PatchDataset
+from .image_dataset import ImageDataset
 from .image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader
 from .iterable_dataset import IterableDataset
-from .nifti_reader import NiftiDataset
 from .nifti_saver import NiftiSaver
 from .nifti_writer import write_nifti
 from .png_saver import PNGSaver
diff --git a/monai/data/nifti_reader.py b/monai/data/image_dataset.py
similarity index 71%
rename from monai/data/nifti_reader.py
rename to monai/data/image_dataset.py
--- a/monai/data/nifti_reader.py
+++ b/monai/data/image_dataset.py
@@ -9,19 +9,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Callable, Optional, Sequence, Union
 
 import numpy as np
 from torch.utils.data import Dataset
 
+from monai.data.image_reader import ImageReader
 from monai.transforms import LoadImage, Randomizable, apply_transform
 from monai.utils import MAX_SEED, get_seed
 
 
-class NiftiDataset(Dataset, Randomizable):
+class ImageDataset(Dataset, Randomizable):
     """
-    Loads image/segmentation pairs of Nifti files from the given filename lists. Transformations can be specified
+    Loads image/segmentation pairs of files from the given filename lists. Transformations can be specified
     for the image and segmentation arrays separately.
+    The difference between this dataset and `ArrayDataset` is that this dataset can apply transform chain to images
+    and segs and return both the images and metadata, and no need to specify transform to load images from files.
+
     """
 
     def __init__(
@@ -29,11 +33,13 @@ def __init__(
         image_files: Sequence[str],
         seg_files: Optional[Sequence[str]] = None,
         labels: Optional[Sequence[float]] = None,
-        as_closest_canonical: bool = False,
         transform: Optional[Callable] = None,
         seg_transform: Optional[Callable] = None,
         image_only: bool = True,
         dtype: Optional[np.dtype] = np.float32,
+        reader: Optional[Union[ImageReader, str]] = None,
+        *args,
+        **kwargs,
     ) -> None:
         """
         Initializes the dataset with the image and segmentation filename lists. The transform `transform` is applied
@@ -43,14 +49,18 @@ def __init__(
             image_files: list of image filenames
             seg_files: if in segmentation task, list of segmentation filenames
             labels: if in classification task, list of classification labels
-            as_closest_canonical: if True, load the image as closest to canonical orientation
             transform: transform to apply to image arrays
             seg_transform: transform to apply to segmentation arrays
-            image_only: if True return only the image volume, other return image volume and header dict
+            image_only: if True return only the image volume, otherwise, return image volume and the metadata
             dtype: if not None convert the loaded image to this data type
+            reader: register reader to load image file and meta data, if None, will use the default readers.
+                If a string of reader name provided, will construct a reader object with the `*args` and `**kwargs`
+                parameters, supported reader name: "NibabelReader", "PILReader", "ITKReader", "NumpyReader"
+            args: additional parameters for reader if providing a reader name
+            kwargs: additional parameters for reader if providing a reader name
 
         Raises:
-            ValueError: When ``seg_files`` length differs from ``image_files``.
+            ValueError: When ``seg_files`` length differs from ``image_files``
 
         """
 
@@ -63,13 +73,11 @@ def __init__(
         self.image_files = image_files
         self.seg_files = seg_files
         self.labels = labels
-        self.as_closest_canonical = as_closest_canonical
         self.transform = transform
         self.seg_transform = seg_transform
         self.image_only = image_only
-        self.dtype = dtype
+        self.loader = LoadImage(reader, image_only, dtype, *args, **kwargs)
         self.set_random_state(seed=get_seed())
-
         self._seed = 0  # transform synchronization seed
 
     def __len__(self) -> int:
@@ -81,21 +89,18 @@ def randomize(self, data: Optional[Any] = None) -> None:
     def __getitem__(self, index: int):
         self.randomize()
         meta_data = None
-        img_loader = LoadImage(
-            reader="NibabelReader",
-            image_only=self.image_only,
-            dtype=self.dtype,
-            as_closest_canonical=self.as_closest_canonical,
-        )
-        if self.image_only:
-            img = img_loader(self.image_files[index])
-        else:
-            img, meta_data = img_loader(self.image_files[index])
         seg = None
-        if self.seg_files is not None:
-            seg_loader = LoadImage(image_only=True)
-            seg = seg_loader(self.seg_files[index])
         label = None
+
+        if self.image_only:
+            img = self.loader(self.image_files[index])
+            if self.seg_files is not None:
+                seg = self.loader(self.seg_files[index])
+        else:
+            img, meta_data = self.loader(self.image_files[index])
+            if self.seg_files is not None:
+                seg, _ = self.loader(self.seg_files[index])
+
         if self.labels is not None:
             label = self.labels[index]
 


Test patch:
diff --git a/tests/min_tests.py b/tests/min_tests.py
--- a/tests/min_tests.py
+++ b/tests/min_tests.py
@@ -69,7 +69,7 @@ def run_testsuit():
         "test_load_imaged",
         "test_load_spacing_orientation",
         "test_mednistdataset",
-        "test_nifti_dataset",
+        "test_image_dataset",
         "test_nifti_header_revise",
         "test_nifti_rw",
         "test_nifti_saver",
diff --git a/tests/test_nifti_dataset.py b/tests/test_image_dataset.py
similarity index 88%
rename from tests/test_nifti_dataset.py
rename to tests/test_image_dataset.py
--- a/tests/test_nifti_dataset.py
+++ b/tests/test_image_dataset.py
@@ -16,7 +16,7 @@
 import nibabel as nib
 import numpy as np
 
-from monai.data import NiftiDataset
+from monai.data import ImageDataset
 from monai.transforms import Randomizable
 
 FILENAMES = ["test1.nii.gz", "test2.nii", "test3.nii.gz"]
@@ -35,7 +35,7 @@ def __call__(self, data):
         return data + self._a
 
 
-class TestNiftiDataset(unittest.TestCase):
+class TestImageDataset(unittest.TestCase):
     def test_dataset(self):
         with tempfile.TemporaryDirectory() as tempdir:
             full_names, ref_data = [], []
@@ -47,46 +47,46 @@ def test_dataset(self):
                 nib.save(nib.Nifti1Image(test_image, np.eye(4)), save_path)
 
             # default loading no meta
-            dataset = NiftiDataset(full_names)
+            dataset = ImageDataset(full_names)
             for d, ref in zip(dataset, ref_data):
                 np.testing.assert_allclose(d, ref, atol=1e-3)
 
             # loading no meta, int
-            dataset = NiftiDataset(full_names, dtype=np.float16)
+            dataset = ImageDataset(full_names, dtype=np.float16)
             for d, _ in zip(dataset, ref_data):
                 self.assertEqual(d.dtype, np.float16)
 
             # loading with meta, no transform
-            dataset = NiftiDataset(full_names, image_only=False)
+            dataset = ImageDataset(full_names, image_only=False)
             for d_tuple, ref in zip(dataset, ref_data):
                 d, meta = d_tuple
                 np.testing.assert_allclose(d, ref, atol=1e-3)
                 np.testing.assert_allclose(meta["original_affine"], np.eye(4))
 
             # loading image/label, no meta
-            dataset = NiftiDataset(full_names, seg_files=full_names, image_only=True)
+            dataset = ImageDataset(full_names, seg_files=full_names, image_only=True)
             for d_tuple, ref in zip(dataset, ref_data):
                 img, seg = d_tuple
                 np.testing.assert_allclose(img, ref, atol=1e-3)
                 np.testing.assert_allclose(seg, ref, atol=1e-3)
 
             # loading image/label, no meta
-            dataset = NiftiDataset(full_names, transform=lambda x: x + 1, image_only=True)
+            dataset = ImageDataset(full_names, transform=lambda x: x + 1, image_only=True)
             for d, ref in zip(dataset, ref_data):
                 np.testing.assert_allclose(d, ref + 1, atol=1e-3)
 
             # set seg transform, but no seg_files
             with self.assertRaises(RuntimeError):
-                dataset = NiftiDataset(full_names, seg_transform=lambda x: x + 1, image_only=True)
+                dataset = ImageDataset(full_names, seg_transform=lambda x: x + 1, image_only=True)
                 _ = dataset[0]
 
             # set seg transform, but no seg_files
             with self.assertRaises(RuntimeError):
-                dataset = NiftiDataset(full_names, seg_transform=lambda x: x + 1, image_only=True)
+                dataset = ImageDataset(full_names, seg_transform=lambda x: x + 1, image_only=True)
                 _ = dataset[0]
 
             # loading image/label, with meta
-            dataset = NiftiDataset(
+            dataset = ImageDataset(
                 full_names,
                 transform=lambda x: x + 1,
                 seg_files=full_names,
@@ -100,7 +100,7 @@ def test_dataset(self):
                 np.testing.assert_allclose(meta["original_affine"], np.eye(4), atol=1e-3)
 
             # loading image/label, with meta
-            dataset = NiftiDataset(
+            dataset = ImageDataset(
                 full_names, transform=lambda x: x + 1, seg_files=full_names, labels=[1, 2, 3], image_only=False
             )
             for idx, (d_tuple, ref) in enumerate(zip(dataset, ref_data)):
@@ -111,7 +111,7 @@ def test_dataset(self):
                 np.testing.assert_allclose(meta["original_affine"], np.eye(4), atol=1e-3)
 
             # loading image/label, with sync. transform
-            dataset = NiftiDataset(
+            dataset = ImageDataset(
                 full_names, transform=RandTest(), seg_files=full_names, seg_transform=RandTest(), image_only=False
             )
             for d_tuple, ref in zip(dataset, ref_data):
diff --git a/tests/test_integration_sliding_window.py b/tests/test_integration_sliding_window.py
--- a/tests/test_integration_sliding_window.py
+++ b/tests/test_integration_sliding_window.py
@@ -19,7 +19,7 @@
 from ignite.engine import Engine
 from torch.utils.data import DataLoader
 
-from monai.data import NiftiDataset, create_test_image_3d
+from monai.data import ImageDataset, create_test_image_3d
 from monai.handlers import SegmentationSaver
 from monai.inferers import sliding_window_inference
 from monai.networks import eval_mode, predict_segmentation
@@ -30,7 +30,7 @@
 
 
 def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
-    ds = NiftiDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=False)
+    ds = ImageDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=False)
     loader = DataLoader(ds, batch_size=1, pin_memory=torch.cuda.is_available())
 
     net = UNet(