#6188 Allow user to define own FolderLayout in SaveImage and SaveImaged transforms

monai/data/__init__.py

@@ -46,7 +46,7 @@
     load_decathlon_datalist,
     load_decathlon_properties,
 )
-from .folder_layout import FolderLayout
+from .folder_layout import FolderLayout, FolderLayoutBase
 from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter, PatchIterd
 from .image_dataset import ImageDataset
 from .image_reader import ImageReader, ITKReader, NibabelReader, NrrdReader, NumpyReader, PILReader, PydicomReader

monai/data/folder_layout.py

@@ -11,14 +11,16 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+
 import monai
 from monai.config import PathLike
 from monai.data.utils import create_file_basename
 
-__all__ = ["FolderLayout", "default_name_formatter"]
+__all__ = ["FolderLayoutBase", "FolderLayout", "default_name_formatter"]
 
 
-def default_name_formatter(metadict, saver):
+def default_name_formatter(metadict: dict, saver: monai.transforms.Transform) -> dict:
     """Returns a kwargs dict for :py:meth:`FolderLayout.filename`,
     according to the input metadata and SaveImage transform."""
     subject = (
@@ -30,7 +32,58 @@ def default_name_formatter(metadict, saver):
     return {"subject": f"{subject}", "idx": patch_index}
 
 
-class FolderLayout:
+class FolderLayoutBase(ABC):
+    """
+    Abstract base class to define a common interface for FolderLayout and derived classes
+    Mainly, defines the ``filename(**kwargs) -> PathLike`` function, which must be defined
+    by the deriving class.
+
+    Example:
+
+    .. code-block:: python
+
+        from monai.data import FolderLayoutBase
+
+        class MyFolderLayout(FolderLayoutBase):
+            def __init__(
+                self,
+                basepath: Path,
+                extension: str = "",
+                makedirs: bool = False
+            ):
+                self.basepath = basepath
+                if not extension:
+                    self.extension = ""
+                elif extension.startswith("."):
+                    self.extension = extension:
+                else:
+                    self.extension = f".{extension}"
+                self.makedirs = makedirs
+
+            def filename(self, patient_no: int, image_name: str, **kwargs) -> Path:
+                sub_path = self.basepath / patient_no
+                if not sub_path.exists():
+                    sub_path.mkdir(parents=True)
+
+                file = image_name
+                for k, v in kwargs.items():
+                    file += f"_{k}-{v}"
+
+                file +=  self.extension
+                return sub_path / file
+
+    """
+
+    @abstractmethod
+    def filename(self, **kwargs) -> PathLike:
+        """
+        Create a filename with path based on the input kwargs.
+        Abstract method, implement your own.
+        """
+        raise NotImplementedError
+
+
+class FolderLayout(FolderLayoutBase):
     """
     A utility class to create organized filenames within ``output_dir``. The
     ``filename`` method could be used to create a filename following the folder structure.
@@ -81,7 +134,7 @@ def __init__(
         self.makedirs = makedirs
         self.data_root_dir = data_root_dir
 
-    def filename(self, subject: PathLike = "subject", idx=None, **kwargs):
+    def filename(self, subject: PathLike = "subject", idx=None, **kwargs) -> PathLike:
         """
         Create a filename based on the input ``subject`` and ``idx``.
 

monai/transforms/io/array.py

@@ -23,13 +23,14 @@
 from collections.abc import Sequence
 from pathlib import Path
 from pydoc import locate
+from typing import Callable
 
 import numpy as np
 import torch
 
 from monai.config import DtypeLike, NdarrayOrTensor, PathLike
 from monai.data import image_writer
-from monai.data.folder_layout import FolderLayout, default_name_formatter
+from monai.data.folder_layout import FolderLayout, FolderLayoutBase, default_name_formatter
 from monai.data.image_reader import (
     ImageReader,
     ITKReader,
@@ -315,11 +316,14 @@ class SaveImage(Transform):
 
     Args:
         output_dir: output image directory.
+        Handled by ``folder_layout`` instead, if ``folder_layout`` is not ``None``.
         output_postfix: a string appended to all output file names, default to `trans`.
+        Handled by ``folder_layout`` instead, if ``folder_layout`` is not ``None``.
         output_ext: output file extension name.
+        Handled by ``folder_layout`` instead, if ``folder_layout`` is not ``None``.
         output_dtype: data type (if not None) for saving data. Defaults to ``np.float32``.
         resample: whether to resample image (if needed) before saving the data array,
-            based on the `spatial_shape` (and `original_affine`) from metadata.
+            based on the ``"spatial_shape"`` (and ``"original_affine"``) from metadata.
         mode: This option is used when ``resample=True``. Defaults to ``"nearest"``.
             Depending on the writers, the possible options are
 
@@ -332,40 +336,49 @@ class SaveImage(Transform):
             Possible options are {``"zeros"``, ``"border"``, ``"reflection"``}
             See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
         scale: {``255``, ``65535``} postprocess data by clipping to [0, 1] and scaling
-            [0, 255] (uint8) or [0, 65535] (uint16). Default is `None` (no scaling).
+            [0, 255] (``uint8``) or [0, 65535] (``uint16``). Default is ``None`` (no scaling).
         dtype: data type during resampling computation. Defaults to ``np.float64`` for best precision.
-            if None, use the data type of input data. To set the output data type, use `output_dtype`.
-        squeeze_end_dims: if True, any trailing singleton dimensions will be removed (after the channel
+            if ``None``, use the data type of input data. To set the output data type, use ``output_dtype``.
+        squeeze_end_dims: if ``True``, any trailing singleton dimensions will be removed (after the channel
             has been moved to the end). So if input is (C,H,W,D), this will be altered to (H,W,D,C), and
-            then if C==1, it will be saved as (H,W,D). If D is also 1, it will be saved as (H,W). If `false`,
+            then if C==1, it will be saved as (H,W,D). If D is also 1, it will be saved as (H,W). If ``False``,
             image will always be saved as (H,W,D,C).
         data_root_dir: if not empty, it specifies the beginning parts of the input file's
-            absolute path. It's used to compute `input_file_rel_path`, the relative path to the file from
-            `data_root_dir` to preserve folder structure when saving in case there are files in different
+            absolute path. It's used to compute ``input_file_rel_path``, the relative path to the file from
+            ``data_root_dir`` to preserve folder structure when saving in case there are files in different
             folders with the same file names. For example, with the following inputs:
 
-            - input_file_name: `/foo/bar/test1/image.nii`
-            - output_postfix: `seg`
-            - output_ext: `.nii.gz`
-            - output_dir: `/output`
-            - data_root_dir: `/foo/bar`
+            - input_file_name: ``/foo/bar/test1/image.nii``
+            - output_postfix: ``seg``
+            - output_ext: ``.nii.gz``
+            - output_dir: ``/output``
+            - data_root_dir: ``/foo/bar``
 
-            The output will be: /output/test1/image/image_seg.nii.gz
+            The output will be: ``/output/test1/image/image_seg.nii.gz``
 
+            Handled by ``folder_layout`` instead, if ``folder_layout`` is not ``None``.
         separate_folder: whether to save every file in a separate folder. For example: for the input filename
-            `image.nii`, postfix `seg` and folder_path `output`, if `separate_folder=True`, it will be saved as:
-            `output/image/image_seg.nii`, if `False`, saving as `output/image_seg.nii`. Default to `True`.
-        print_log: whether to print logs when saving. Default to `True`.
+            ``image.nii``, postfix ``seg`` and ``folder_path`` ``output``, if ``separate_folder=True``, it will be
+            saved as: ``output/image/image_seg.nii``, if ``False``, saving as ``output/image_seg.nii``.
+            Default to ``True``.
+            Handled by ``folder_layout`` instead, if ``folder_layout`` is not ``None``.
+        print_log: whether to print logs when saving. Default to ``True``.
         output_format: an optional string of filename extension to specify the output image writer.
-            see also: `monai.data.image_writer.SUPPORTED_WRITERS`.
-        writer: a customised `monai.data.ImageWriter` subclass to save data arrays.
-            if `None`, use the default writer from `monai.data.image_writer` according to `output_ext`.
+            see also: ``monai.data.image_writer.SUPPORTED_WRITERS``.
+        writer: a customised ``monai.data.ImageWriter`` subclass to save data arrays.
+            if ``None``, use the default writer from ``monai.data.image_writer`` according to ``output_ext``.
             if it's a string, it's treated as a class name or dotted path (such as ``"monai.data.ITKWriter"``);
             the supported built-in writer classes are ``"NibabelWriter"``, ``"ITKWriter"``, ``"PILWriter"``.
-        channel_dim: the index of the channel dimension. Default to `0`.
-            `None` to indicate no channel dimension.
+        channel_dim: the index of the channel dimension. Default to ``0``.
+            ``None`` to indicate no channel dimension.
         output_name_formatter: a callable function (returning a kwargs dict) to format the output file name.
+            If using a custom ``monai.data.FolderLayoutBase`` class in ``folder_layout``, consider providing
+            your own formatter.
             see also: :py:func:`monai.data.folder_layout.default_name_formatter`.
+        folder_layout: A customized ``monai.data.FolderLayoutBase`` subclass to define file naming schemes.
+            if ``None``, uses the default ``FolderLayout``.
+        savepath_in_metadict: if ``True``, adds a key ``"saved_to"`` to the metadata, which contains the path
+            to where the input image has been saved.
     """
 
     @deprecated_arg_default("resample", True, False, since="1.1", replaced="1.3")
@@ -387,16 +400,28 @@ def __init__(
         output_format: str = "",
         writer: type[image_writer.ImageWriter] | str | None = None,
         channel_dim: int | None = 0,
-        output_name_formatter=None,
+        output_name_formatter: Callable[[dict, Transform], dict] | None = None,
+        folder_layout: FolderLayoutBase | None = None,
+        savepath_in_metadict: bool = False,
     ) -> None:
-        self.folder_layout = FolderLayout(
-            output_dir=output_dir,
-            postfix=output_postfix,
-            extension=output_ext,
-            parent=separate_folder,
-            makedirs=True,
-            data_root_dir=data_root_dir,
-        )
+        self.folder_layout: FolderLayoutBase
+        if folder_layout is None:
+            self.folder_layout = FolderLayout(
+                output_dir=output_dir,
+                postfix=output_postfix,
+                extension=output_ext,
+                parent=separate_folder,
+                makedirs=True,
+                data_root_dir=data_root_dir,
+            )
+        else:
+            self.folder_layout = folder_layout
+
+        self.fname_formatter: Callable
+        if output_name_formatter is None:
+            self.fname_formatter = default_name_formatter
+        else:
+            self.fname_formatter = output_name_formatter
 
         self.output_ext = output_ext.lower() or output_format.lower()
         if isinstance(writer, str):
@@ -418,8 +443,8 @@ def __init__(
         self.data_kwargs = {"squeeze_end_dims": squeeze_end_dims, "channel_dim": channel_dim}
         self.meta_kwargs = {"resample": resample, "mode": mode, "padding_mode": padding_mode, "dtype": dtype}
         self.write_kwargs = {"verbose": print_log}
-        self.fname_formatter = default_name_formatter if output_name_formatter is None else output_name_formatter
         self._data_index = 0
+        self.savepath_in_metadict = savepath_in_metadict
 
     def set_options(self, init_kwargs=None, data_kwargs=None, meta_kwargs=None, write_kwargs=None):
         """
@@ -478,6 +503,8 @@ def __call__(self, img: torch.Tensor | np.ndarray, meta_data: dict | None = None
                 )
             else:
                 self._data_index += 1
+                if self.savepath_in_metadict and meta_data is not None:
+                    meta_data["saved_to"] = filename
                 return img
         msg = "\n".join([f"{e}" for e in err])
         raise RuntimeError(