update saveimage and writer selector

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

docs/source/data.rst

@@ -153,6 +153,14 @@ WSIReader
 Image writer
 ------------
 
+resolve_writer
+~~~~~~~~~~~~~~
+.. autofunction:: resolve_writer
+
+register_writer
+~~~~~~~~~~~~~~~
+.. autofunction:: register_writer
+
 ImageWriter
 ~~~~~~~~~~~
 .. autoclass:: ImageWriter

monai/data/__init__.py

@@ -35,7 +35,16 @@
 from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter
 from .image_dataset import ImageDataset
 from .image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader, WSIReader
-from .image_writer import ImageWriter, ITKWriter, NibabelWriter, PILWriter, logger
+from .image_writer import (
+    SUPPORTED_WRITERS,
+    ImageWriter,
+    ITKWriter,
+    NibabelWriter,
+    PILWriter,
+    logger,
+    register_writer,
+    resolve_writer,
+)
 from .iterable_dataset import CSVIterableDataset, IterableDataset, ShuffleBuffer
 from .nifti_saver import NiftiSaver
 from .nifti_writer import write_nifti

monai/data/image_reader.py

@@ -18,12 +18,10 @@
 from torch.utils.data._utils.collate import np_str_obj_array_pattern
 
 from monai.config import DtypeLike, KeysCollection, PathLike
-from monai.data.utils import correct_nifti_header_if_necessary
+from monai.data.utils import correct_nifti_header_if_necessary, is_supported_format
 from monai.transforms.utility.array import EnsureChannelFirst
 from monai.utils import ensure_tuple, ensure_tuple_rep, optional_import, require_pkg
 
-from .utils import is_supported_format
-
 if TYPE_CHECKING:
     import itk
     import nibabel as nib

monai/data/image_writer.py

@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Mapping, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Dict, Mapping, Optional, Sequence, Union
 
 import numpy as np
 
@@ -22,6 +22,7 @@
     GridSampleMode,
     GridSamplePadMode,
     InterpolateMode,
+    OptionalImportError,
     convert_data_type,
     look_up_option,
     optional_import,
@@ -41,7 +42,69 @@
     PILImage, _ = optional_import("PIL.Image")
 
 
-__all__ = ["ImageWriter", "ITKWriter", "NibabelWriter", "PILWriter", "logger"]
+__all__ = [
+    "ImageWriter",
+    "ITKWriter",
+    "NibabelWriter",
+    "PILWriter",
+    "SUPPORTED_WRITERS",
+    "register_writer",
+    "resolve_writer",
+    "logger",
+]
+
+SUPPORTED_WRITERS: Dict = {}
+
+
+def register_writer(ext_name, *im_writer):
+    """
+    Register ``ImageWriter``, so that writing a file with filename extension ``ext_name``
+    could be resolved to a tuple of potentially appropriate ``ImageWriter``.
+    The customised writers could be registered by:
+
+    .. code-block:: python
+
+        from monai.data import image_writer
+        # `MyWriter` must implement `ImageWriter` interface
+        image_writer.register_writer(".nii", MyWriter)
+
+    Args:
+        ext_name: the filename extension of the image.
+            As an indexing key, it will be converted to a lower case string.
+        im_writer: one or multiple ImageWriter classes with high priority ones first.
+    """
+    fmt = f"{ext_name}".lower()
+    existing = look_up_option(fmt, SUPPORTED_WRITERS, default=())
+    all_writers = im_writer + existing
+    SUPPORTED_WRITERS[fmt] = all_writers
+
+
+def resolve_writer(ext_name, error_if_not_found=True) -> Sequence:
+    """
+    Resolves to a tuple of available ``ImageWriter`` in ``SUPPORTED_WRITERS``
+    according to the filename extension key ``ext_name``.
+
+    Args:
+        ext_name: the filename extension of the image.
+            As an indexing key it will be converted to a lower case string.
+        error_if_not_found: whether to raise an error if no suitable image writer is found.
+            if True , raise an ``OptionalImportError``, otherwise return an empty tuple. Default is ``True``.
+    """
+    if not SUPPORTED_WRITERS:
+        init()
+    fmt = f"{ext_name}".lower()
+    avail_writers = []
+    for _writer in look_up_option(fmt, SUPPORTED_WRITERS, default=SUPPORTED_WRITERS["*"]):
+        try:
+            _writer()  # this triggers `monai.utils.module.require_pkg` to check the system availability
+            avail_writers.append(_writer)
+        except OptionalImportError:
+            pass
+    if not avail_writers and error_if_not_found:
+        raise OptionalImportError(f"No ImageWriter backend found for {fmt}.")
+    writer_tuple = ensure_tuple(avail_writers)
+    SUPPORTED_WRITERS[fmt] = writer_tuple
+    return writer_tuple
 
 
 class ImageWriter:
@@ -716,3 +779,15 @@ def create_backend_obj(
             data = np.moveaxis(data, 0, 1)
 
         return PILImage.fromarray(data, mode=kwargs.pop("image_mode", None))
+
+
+def init():
+    """
+    Initialize the image writer modules according to the filename extension.
+    """
+    for ext in (".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif"):
+        register_writer(ext, PILWriter)  # TODO: test 16-bit
+    for ext in (".nii.gz", ".nii"):
+        register_writer(ext, NibabelWriter, ITKWriter)
+    register_writer(".nrrd", ITKWriter, NibabelWriter)
+    register_writer("*", ITKWriter, NibabelWriter)

monai/data/nifti_saver.py

@@ -19,8 +19,10 @@
 from monai.data.utils import create_file_basename
 from monai.utils import GridSampleMode, GridSamplePadMode
 from monai.utils import ImageMetaKey as Key
+from monai.utils import deprecated
 
 
+@deprecated(since="0.8", msg_suffix="use monai.transforms.SaveImage instead.")
 class NiftiSaver:
     """
     Save the data as NIfTI file, it can support single data content or a batch of data.
@@ -32,6 +34,9 @@ class NiftiSaver:
 
     Note: image should include channel dimension: [B],C,H,W,[D].
 
+    .. deprecated:: 0.8
+        Use :py:class:`monai.transforms.SaveImage` instead.
+
     """
 
     def __init__(

monai/data/nifti_writer.py

@@ -19,12 +19,13 @@
 from monai.data.utils import compute_shape_offset, to_affine_nd
 from monai.networks.layers import AffineTransform
 from monai.transforms.utils_pytorch_numpy_unification import allclose
-from monai.utils import GridSampleMode, GridSamplePadMode, optional_import
+from monai.utils import GridSampleMode, GridSamplePadMode, deprecated, optional_import
 from monai.utils.type_conversion import convert_data_type
 
 nib, _ = optional_import("nibabel")
 
 
+@deprecated(since="0.8", msg_suffix="use monai.data.NibabelWriter instead.")
 def write_nifti(
     data: NdarrayOrTensor,
     file_name: str,
@@ -98,6 +99,10 @@ def write_nifti(
         dtype: data type for resampling computation. Defaults to ``np.float64`` for best precision.
             If None, use the data type of input data.
         output_dtype: data type for saving data. Defaults to ``np.float32``.
+
+    .. deprecated:: 0.8
+        Use :py:meth:`monai.data.NibabelWriter` instead.
+
     """
     data, *_ = convert_data_type(data, np.ndarray)
     affine, *_ = convert_data_type(affine, np.ndarray)

monai/data/png_saver.py

@@ -18,9 +18,10 @@
 from monai.data.png_writer import write_png
 from monai.data.utils import create_file_basename
 from monai.utils import ImageMetaKey as Key
-from monai.utils import InterpolateMode, look_up_option
+from monai.utils import InterpolateMode, deprecated, look_up_option
 
 
+@deprecated(since="0.8", msg_suffix="use monai.transforms.SaveImage instead.")
 class PNGSaver:
     """
     Save the data as png file, it can support single data content or a batch of data.
@@ -30,6 +31,9 @@ class PNGSaver:
     where the input image name is extracted from the provided meta data dictionary.
     If no meta data provided, use index from 0 as the filename prefix.
 
+    .. deprecated:: 0.8
+        Use :py:class:`monai.transforms.SaveImage` instead.
+
     """
 
     def __init__(

monai/data/png_writer.py

@@ -14,11 +14,12 @@
 import numpy as np
 
 from monai.transforms.spatial.array import Resize
-from monai.utils import InterpolateMode, ensure_tuple_rep, look_up_option, optional_import
+from monai.utils import InterpolateMode, deprecated, ensure_tuple_rep, look_up_option, optional_import
 
 Image, _ = optional_import("PIL", name="Image")
 
 
+@deprecated(since="0.8", msg_suffix="use monai.data.PILWriter instead.")
 def write_png(
     data: np.ndarray,
     file_name: str,
@@ -46,6 +47,9 @@ def write_png(
     Raises:
         ValueError: When ``scale`` is not one of [255, 65535].
 
+    .. deprecated:: 0.8
+        Use :py:meth:`monai.data.PILWriter` instead.
+
     """
     if not isinstance(data, np.ndarray):
         raise ValueError("input data must be numpy array.")