I'm having some issues understanding SaveImage and SaveImaged output path and filename configuration

[MathijsdeBoer]

I'm trying to set up a simple dataset/model management suite to simplify running experiments with supervised and unsupervised training techniques.
Part of this is the dataset preprocessing step, where I want to run all collected samples through some Transforms, and save them to another location at the end.
I'm having some issues with this:

A) I want to know where the SaveImaged Transform has actually saved my sample, this data does not seem to be included in the output metadict.
Only the original input image seems to be present in the dict.

B) I want to specify with more granularity where this sample ends up.

Solving either problem will make my life easier because I can use the solution of one to work around the other issue.
Solving both would be awesome.

---

Consider my dataset with which I'm testing my implementation is the BraTS2021 data, as it's provided on the challenge page:

RSNA_ASNR_MICCAI_BraTS2021_TrainingData_16July2021
    |-  BraTS2021_00000
    |-  BraTS2021_00002
            |-  BraTS2021_00002_flair.nii.gz
            |-  BraTS2021_00002_seg.nii.gz
            |-  BraTS2021_00002_t1.nii.gz
            |-  BraTS2021_00002_t1ce.nii.gz
            |-  BraTS2021_00002_t2.nii.gz
    |-  ...
    |-  BraTS2021_01666
RSNA_ASNR_MICCAI_BraTS2021_ValidationData
    |-  BraTS2021_00001
    |-  BraTS2021_00013
            |-  BraTS2021_00013_flair.nii.gz
            |-  BraTS2021_00013_t1.nii.gz
            |-  BraTS2021_00013_t1ce.nii.gz
            |-  BraTS2021_00013_t2.nii.gz
    |-  ...
    |-  BraTS2021_01798

I've written some code to manage this dataset, where _SerializableClass is a simple abstract class which defines to_dict() and from_dict() to simplify saving these objects to JSON.

@dataclass(init=True, repr=True, eq=False, kw_only=True, slots=True)
class Sample(_SerializableClass):
    """
    Represents a single sample in our dataset.

    Represents a single sample in our dataset.
    Can be supervised or unsupervised, depending on the presence of a label path.

    :param image: The path to the image file
    :param label: The path to the label file, optional
    :param name: The name of this sample, optional.
                 If not given, will default to the name of the image.
    :param is_preprocessed: Whether this sample has been pre-processed or not
    """

    image: Path
    label: Path = field(default=None)
    name: str = field(default=None)
    is_preprocessed: bool = field(default=False)

    ...

@dataclass(init=True, repr=True, eq=False, kw_only=True, slots=True)
class Dataset(_SerializableClass):
    """
    A representation of an entire dataset.

    A representation of an entire dataset, with all accompanying parameters.
    This object manages samples, pre-processing and general dataset information.

    :param name: The name of this dataset
    :param authors: The authors of this dataset
    :param license: The license of this dataset.
    :param creation_date: The moment when this dataset was first instanciated
    :param modified_date: The last time this dataset was altered
    :param channels: The channels of the input images, and what they represent
    :param classes: The classes of the labels, and what they represent
    :param samples: A list of the samples of this dataset.
    :param spacing_override: If set, will override median spacing of this dataset for pre-processing.
    """

    name: str = field(repr=True)
    path: Path = field(init=False, repr=False)
    authors: list[str] = field(default_factory=list, repr=True)
    license: str = field(default="All rights reserved", repr=True)
    creation_date: Arrow = field(default_factory=now, repr=True)
    modified_date: Arrow = field(default_factory=now, repr=True)
    channels: dict[int, str] = field(default_factory=dict, init=True, repr=False)
    classes: dict[int, str] = field(default_factory=dict, init=True, repr=False)
    samples: list[Sample] = field(default_factory=list, init=True, repr=False)
    spacing_override: Optional[tuple[float, float, float]] = field(
        init=True, default=None
    )

    ...

The JSON representations of such objects looks like:

{
    "name": "BraTS2021",
    "authors": [],
    "license": "All rights reserved",
    "creation_date": "2023-03-20T17:21:39.470708+01:00",
    "modified_date": "2023-03-20T17:21:58.437291+01:00",
    "channels": {},
    "classes": {},
    "spacing_override": null,
    "samples": [
        {
            "name": "BraTS2021_00000",
            "is_preprocessed": false,
            "image": "E:/BraTS2021/RSNA_ASNR_MICCAI_BraTS2021_TrainingData_16July2021/BraTS2021_00000/BraTS2021_00000_t1ce.nii.gz",
            "label": "E:/BraTS2021/RSNA_ASNR_MICCAI_BraTS2021_TrainingData_16July2021/BraTS2021_00000/BraTS2021_00000_seg.nii.gz"
        },

    ...

        {
            "name": "BraTS2021_01767",
            "is_preprocessed": false,
            "image": "E:/BraTS2021/RSNA_ASNR_MICCAI_BraTS2021_ValidationData/BraTS2021_01767/BraTS2021_01767_t1ce.nii.gz"
        },
    ...

Now that I have the relevant samples, I want to be able to take these and feed them through a preprocessing pipeline

from pprint import pprint

import monai.transforms as mt

# ---default_transforms.py---

def init_base_transforms(
    image_key: str = "image",
    label_key: Optional[str] = None,
    spacing: tuple[float, float, float] = (0.5, 0.5, 0.5),
    allow_missing: bool = False,
) -> mt.Compose:
    keys = [image_key]
    interp_mode = ["bilinear"]
    if label_key is not None:
        keys.append(label_key)
        interp_mode.append("nearest")

    transforms = [
        mt.LoadImaged(keys=keys, allow_missing_keys=allow_missing),
        mt.EnsureChannelFirstd(keys=keys, allow_missing_keys=allow_missing),
        mt.EnsureTyped(keys=keys, allow_missing_keys=allow_missing),
        mt.Orientationd(keys=keys, axcodes="LPS", allow_missing_keys=allow_missing),
        mt.Spacingd(
            keys=keys,
            pixdim=spacing,
            mode=interp_mode,
            allow_missing_keys=allow_missing,
        ),
        mt.NormalizeIntensityd(
            keys=[image_key],
            nonzero=True,
            channel_wise=True,
            allow_missing_keys=allow_missing,
        ),
    ]

    return mt.Compose(transforms)

# ---------------------------

preprocessing_transforms = init_base_transforms(
    image_key="image",
    label_key="label",
    spacing=dataset.dataset_spacing,
    allow_missing=True,
)

# Helper function to create a list[dict[str, Path]] for use in dictionary transforms
samples = dataset.as_monai_list(get_supervised=True, get_unsupervised=True)

save_transform = mt.SaveImaged(
    keys=["image", "label"],
    allow_missing_keys=True,
    # dataset.json location -> ./Preprocessed
    output_dir=dataset.processed_samples_path,
    output_postfix="",
    resample=False,
    separate_folder=False,
    print_log=True,
)

for sample in samples:
    output = preprocessing_transforms(sample)
    pprint(save_transform(output))

Running this code gives me:

2023-03-20 17:30:30,667 INFO image_writer.py:194 - writing: E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00000_t1ce.nii.gz
2023-03-20 17:30:30,970 INFO image_writer.py:194 - writing: E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00000_seg.nii.gz

{'image': *snip*,
 'image_meta_dict': {affine: tensor([[ -1.,  -0.,  -0.,   0.],
        [ -0.,  -1.,  -0., 239.],
        [  0.,   0.,   1.,   0.],
        [  0.,   0.,   0.,   1.]], dtype=torch.float64),
                     'as_closest_canonical': False,
                     'bitpix': array(16, dtype=int16),
                     'cal_max': array(0., dtype=float32),
                     'cal_min': array(0., dtype=float32),
                     'datatype': array(4, dtype=int16),
                     'dim': array([  3, 240, 240, 155,   1,   1,   1,   1], dtype=int16),
                     'dim_info': array(0, dtype=uint8),
                     'extents': array(0),
                     'filename_or_obj': 'E:\\BraTS2021\\RSNA_ASNR_MICCAI_BraTS2021_TrainingData_16July2021\\BraTS2021_00000\\BraTS2021_00000_t1ce.nii.gz',
                     'glmax': array(0),
                     'glmin': array(0),
                     'intent_code': array(0, dtype=int16),
                     'intent_p1': array(0., dtype=float32),
                     'intent_p2': array(0., dtype=float32),
                     'intent_p3': array(0., dtype=float32),
                     original_affine: array([[ -1.,  -0.,  -0.,   0.],
       [ -0.,  -1.,  -0., 239.],
       [  0.,   0.,   1.,   0.],
       [  0.,   0.,   0.,   1.]]),
                     original_channel_dim: 'no_channel',
                     'pixdim': array([1., 1., 1., 1., 0., 0., 0., 0.], dtype=float32),
                     'qform_code': array(1, dtype=int16),
                     'qoffset_x': array(0., dtype=float32),
                     'qoffset_y': array(239., dtype=float32),
                     'qoffset_z': array(0., dtype=float32),
                     'quatern_b': array(0., dtype=float32),
                     'quatern_c': array(0., dtype=float32),
                     'quatern_d': array(1., dtype=float32),
                     'scl_inter': array(nan, dtype=float32),
                     'scl_slope': array(nan, dtype=float32),
                     'session_error': array(0, dtype=int16),
                     'sform_code': array(1, dtype=int16),
                     'sizeof_hdr': array(348),
                     'slice_code': array(0, dtype=uint8),
                     'slice_duration': array(0., dtype=float32),
                     'slice_end': array(0, dtype=int16),
                     'slice_start': array(0, dtype=int16),
                     space: RAS,
                     spatial_shape: array([240, 240, 155], dtype=int16),
                     'srow_x': array([-1., -0., -0.,  0.], dtype=float32),
                     'srow_y': array([ -0.,  -1.,  -0., 239.], dtype=float32),
                     'srow_z': array([0., 0., 1., 0.], dtype=float32),
                     'toffset': array(0., dtype=float32),
                     'vox_offset': array(0., dtype=float32),
                     'xyzt_units': array(2, dtype=uint8)},
 'label': *snip*,
 'label_meta_dict': *snip, similar to 'image_meta_dict'*

If I set separate_folder=True, I get:

2023-03-20 17:36:03,351 INFO image_writer.py:194 - writing: E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00000_t1ce\BraTS2021_00000_t1ce.nii.gz
2023-03-20 17:36:03,673 INFO image_writer.py:194 - writing: E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00000_seg\BraTS2021_00000_seg.nii.gz

I've been digging into the output_name_formatter argument, including the FolderLayout object, but I can't seem to understand how these are supposed to function.
Digging into the code for these, they seem to be relatively rigid in their function.
Ideally, I'd be able to write my output as follows:

Image: E:\BraTS2021\RSNA_ASNR_MICCAI_BraTS2021_TrainingData_16July2021\BraTS2021_00000\BraTS2021_00000_t1ce.nii.gz
Label: E:\\BraTS2021\\RSNA_ASNR_MICCAI_BraTS2021_TrainingData_16July2021\\BraTS2021_00000\\BraTS2021_00000_seg.nii.gz

Preprocessing

Saves to

Image:E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00000\image.nii.gz
Label:E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00000\label.nii.gz

Where E:\...\Datasets\BraTS2021\Preprocessed\ is given in output_dir, and ...\BraTS2021_0000\ is inferred from the sample name.
My second step would then to either update my Sample objects to point to the new location, or take a leaf out of nnU-Net's book and create a second list with the pre-processed samples.
Either way, I'll need to find a way to find out where the SaveImaged has put my samples.

---

I hope I haven't overwhelmed you with logs and code, and that I have paired down my issue to the minimal example.
I could of course write my own pre-processing pipeline, but if MONAI already offers most steps it'd strike me as wasteful to implement it from the ground up.

[MathijsdeBoer]

I've found a solution, but it feels very convoluted, and prone to breaking on any future updates in MONAI. In the current version, V1.1.0, I can do the following.

First, I make a CustomFolderLayout, based on monai.data.folder_layout.FolderLayout:

class CustomFolderLayout:
    def __init__(
        self,
        output_dir: Path,
        extension: str = "",
        makedirs: bool = True,
        **kwargs: dict
    ):
        self.output_dir: Path = output_dir

        if not extension:
            self.ext = ""
        elif extension.startswith("."):
            self.ext = extension
        else:
            self.ext = f".{extension}"

        self.makedirs = makedirs

    def filename(self, subdirectory: str, filename: str = "subject", **kwargs):
        path = self.output_dir / subdirectory

        for k, v in kwargs.items():
            filename += f"_{k}-{v}"

        if self.ext is not None:
            filename += self.ext

        return path / filename

In here I need the **kwargs to catch off any additional arguments that SaveImage might include. This far I'm not too concerned, it's a little round-peg-in-an-oval-hole, because I have to replace the original FolderStructure object in SaveImageD.saver post-hoc, but it'll do. The part that makes me feel a little iffy is the following:

Now that I have my CustomFolderLayout, I'll need to figure out what image we're dealing with, the image, or the label. Unfortunately, at no point is this information included in the metadata. So I figure I'd add it myself:

def _name_formatter(metadict: dict, saver: mt.SaveImaged) -> dict:
    return {
        "subdirectory": metadict["sample_name"],
        "filename": metadict["dict_key"]
    }

...

save_transform = mt.SaveImaged(
    keys=["image", "label"],
    allow_missing_keys=True,
    output_dir=dataset.processed_samples_path,
    output_postfix="",
    resample=False,
    separate_folder=False,
    print_log=True,
    output_name_formatter=_name_formatter,
)

# Please forgive me for this absolute hack to get path naming to work as I want
save_transform.saver.folder_layout = CustomFolderLayout(
    output_dir=save_transform.saver.folder_layout.output_dir,
    extension=save_transform.saver.output_ext,
    makedirs=True
)

sample: Sample
for sample in tqdm(dataset.samples, desc="Processing samples", leave=False):
    output: dict = preprocessing_transforms(sample.to_monai())
 
   # And this one, I beg of you
    for k in save_transform.keys:
        output[f"{k}_meta_data"]["dict_key"] = k
        output[f"{k}_meta_data"]["sample_name"] = sample.name

    save_transform(output)

But this gave me problems, as the metadata dict in _name_formatter() didn't include my own inserted items. It took me a little more time to find monai.transforms.io.array.py:434, which states the existance of a MetaTensor overrides any {key}_meta_dict information. I hadn't noticed that my data was a MetaTensor during working this out, because when I was printing my values to console (yes, yes, print-debugging strikes again. I'm not experienced enough to hang a debugging session to a terminal in PyCharm yet) the __repr__ prints as if it was a normal torch.Tensor instead.

Idea is the same, we're modifying the MetaTensor.meta dict instead:

...
save_transform = mt.SaveImaged(
    keys=["image", "label"],
    allow_missing_keys=True,
    output_dir=dataset.processed_samples_path,
    output_postfix="",
    resample=False,
    separate_folder=False,
    print_log=True,
    output_name_formatter=_name_formatter,
)

# Please forgive me for this absolute hack to get path naming to work as I want
save_transform.saver.folder_layout = CustomFolderLayout(
    output_dir=save_transform.saver.folder_layout.output_dir,
    extension=save_transform.saver.output_ext,
    makedirs=True
)

sample: Sample
for sample in tqdm(dataset.samples, desc="Processing samples", leave=False):
    output: dict = preprocessing_transforms(sample.to_monai())

    # And this one, I beg of you
    for k in save_transform.keys:
        output[f"{k}"].meta["dict_key"] = k
        output[f"{k}"].meta["sample_name"] = sample.name

    save_transform(output)

This gives me the output I wanted:

2023-03-20 23:53:45,118 INFO image_writer.py:194 - writing: E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00016\image.nii.gz
2023-03-20 23:53:45,539 INFO image_writer.py:194 - writing: E:\...\Datasets\BraTS2021\Preprocessed\BraTS2021_00016\label.nii.gz

And with me standardizing the file names, and knowing the subdirectory, I can go back and update my dataset records to point to the pre-processed data.

My question remains, however, I've managed to solve it this way, but there's no guarantee that this will remain a workable solution, or won't cause conflicts in the future. Did I just completely misunderstand the FileLayout class and associated functions and reinvented the wheel for no good reason?

[KumoLiu]

Hi @MathijsdeBoer, I think you already find a good way to do this. Thanks!

[MathijsdeBoer]

Can I make the assumption then that hitchhiking my own data to the MetaTensor.meta dict will not cause widespread death and destruction? If I don't conflict with existing keys, of course.

[KumoLiu]

Yes, I think so.