diff --git a/monai/data/__init__.py b/monai/data/__init__.py index 8d8297deaf..0e9759aaf1 100644 --- a/monai/data/__init__.py +++ b/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 diff --git a/monai/data/folder_layout.py b/monai/data/folder_layout.py index 190a07334d..4403855030 100644 --- a/monai/data/folder_layout.py +++ b/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``. diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py index a21a070b15..325283c945 100644 --- a/monai/transforms/io/array.py +++ b/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( diff --git a/monai/transforms/io/dictionary.py b/monai/transforms/io/dictionary.py index 47dfbf7e28..b343faadb0 100644 --- a/monai/transforms/io/dictionary.py +++ b/monai/transforms/io/dictionary.py @@ -18,14 +18,16 @@ from __future__ import annotations from pathlib import Path +from typing import Callable import numpy as np +import monai from monai.config import DtypeLike, KeysCollection from monai.data import image_writer from monai.data.image_reader import ImageReader from monai.transforms.io.array import LoadImage, SaveImage -from monai.transforms.transform import MapTransform +from monai.transforms.transform import MapTransform, Transform from monai.utils import GridSamplePadMode, ensure_tuple, ensure_tuple_rep from monai.utils.deprecate_utils import deprecated_arg_default from monai.utils.enums import PostFix @@ -189,16 +191,19 @@ class SaveImaged(MapTransform): keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` meta_keys: explicitly indicate the key of the corresponding metadata dictionary. - For example, for data with key `image`, the metadata by default is in `image_meta_dict`. - The metadata is a dictionary contains values such as filename, original_shape. - This argument can be a sequence of string, map to the `keys`. - If `None`, will try to construct meta_keys by `key_{meta_key_postfix}`. - meta_key_postfix: if `meta_keys` is `None`, use `key_{meta_key_postfix}` to retrieve the metadict. + For example, for data with key ``image``, the metadata by default is in ``image_meta_dict``. + The metadata is a dictionary contains values such as ``filename``, ``original_shape``. + This argument can be a sequence of strings, mapped to the ``keys``. + If ``None``, will try to construct ``meta_keys`` by ``key_{meta_key_postfix}``. + meta_key_postfix: if ``meta_keys`` is ``None``, use ``key_{meta_key_postfix}`` to retrieve the metadict. output_dir: output image directory. - output_postfix: a string appended to all output file names, default to `trans`. - output_ext: output file extension name, available extensions: `.nii.gz`, `.nii`, `.png`. + 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, available extensions: ``.nii.gz``, ``.nii``, ``.png``. + Handled by ``folder_layout`` instead, if ``folder_layout`` not ``None``. 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: @@ -211,9 +216,9 @@ class SaveImaged(MapTransform): 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`. + if None, use the data type of input data. To set the output data type, use ``output_dtype``. output_dtype: data type for saving data. Defaults to ``np.float32``. allow_missing_keys: don't raise exception if key is missing. squeeze_end_dims: if True, any trailing singleton dimensions will be removed (after the channel @@ -221,31 +226,37 @@ class SaveImaged(MapTransform): 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 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; the supported built-in writer classes are ``"NibabelWriter"``, ``"ITKWriter"``, ``"PILWriter"``. output_name_formatter: a callable function (returning a kwargs dict) to format the output file name. see also: :py:func:`monai.data.folder_layout.default_name_formatter`. - + If using a custom ``folder_layout``, consider providing your own 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") @@ -270,7 +281,9 @@ def __init__( print_log: bool = True, output_format: str = "", writer: type[image_writer.ImageWriter] | str | None = None, - output_name_formatter=None, + output_name_formatter: Callable[[dict, Transform], dict] | None = None, + folder_layout: monai.data.FolderLayoutBase | None = None, + savepath_in_metadict: bool = False, ) -> None: super().__init__(keys, allow_missing_keys) self.meta_keys = ensure_tuple_rep(meta_keys, len(self.keys)) @@ -292,6 +305,8 @@ def __init__( output_format=output_format, writer=writer, output_name_formatter=output_name_formatter, + folder_layout=folder_layout, + savepath_in_metadict=savepath_in_metadict, ) def set_options(self, init_kwargs=None, data_kwargs=None, meta_kwargs=None, write_kwargs=None): diff --git a/tests/test_save_imaged.py b/tests/test_save_imaged.py index 676eb74678..ab0b9c0d9f 100644 --- a/tests/test_save_imaged.py +++ b/tests/test_save_imaged.py @@ -14,10 +14,13 @@ import os import tempfile import unittest +from pathlib import Path import torch from parameterized import parameterized +from monai.config import PathLike +from monai.data.folder_layout import FolderLayoutBase from monai.data.meta_tensor import MetaTensor from monai.transforms import SaveImaged from monai.utils import optional_import @@ -68,6 +71,57 @@ def test_saved_content(self, test_data, output_ext, resample): filepath = os.path.join("testfile0", "testfile0" + "_trans" + patch_index + output_ext) self.assertTrue(os.path.exists(os.path.join(tempdir, filepath))) + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) + def test_custom_folderlayout(self, test_data, output_ext, resample): + class TestFolderLayout(FolderLayoutBase): + def __init__(self, basepath: Path, extension: str, makedirs: bool): + self.basepath = basepath + self.ext = extension + self.makedirs = makedirs + + def filename(self, **kwargs) -> PathLike: + p = self.basepath / str(kwargs["subdirectory"]) + if not p.exists() and self.makedirs: + p.mkdir() + + return p / (str(kwargs["filename"]) + self.ext) + + def name_formatter(metadict: dict, _) -> dict: + # "[filename].[ext]" + # quick and dirty split on . + base_filename = metadict["filename_or_obj"].split(".")[0] + + return {"subdirectory": base_filename, "filename": "image"} + + with tempfile.TemporaryDirectory() as tempdir: + trans = SaveImaged( + keys=["img", "pred"], + resample=resample, + allow_missing_keys=True, + output_name_formatter=name_formatter, + folder_layout=TestFolderLayout(basepath=Path(tempdir), extension=output_ext, makedirs=True), + ) + trans(test_data) + + filepath = os.path.join("testfile0", "image" + output_ext) + self.assertTrue(os.path.exists(os.path.join(tempdir, filepath))) + + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) + def test_includes_metadata(self, test_data, output_ext, resample): + with tempfile.TemporaryDirectory() as tempdir: + trans = SaveImaged( + keys=["img", "pred"], + output_dir=tempdir, + output_ext=output_ext, + resample=resample, + allow_missing_keys=True, + savepath_in_metadict=True, + ) + trans(test_data) + + self.assertTrue("saved_to" in test_data["img"].meta.keys()) + self.assertTrue(os.path.exists(test_data["img"].meta["saved_to"])) + if __name__ == "__main__": unittest.main()