Multiple transform return values are not properly passed as parameters for next transform

[Spenhouet]

Describe the bug

transform.apply_transform does currently not handle the case where the return value of a transform call includes multiple return values.

if isinstance(data, (list, tuple)) and map_items:
    return [transform(item) for item in data]

It just passes the return values of the last transform as tuple.

To Reproduce
Steps to reproduce the behavior:

1. Compose the LoadImage and Spacing transforms
2. LoadImage will return the array and meta_data
3. Spacing expects both parameters but will fail

Expected behavior
The parameters are unpacked:

if isinstance(data, (list, tuple)) and map_items:
    return [transform(*items) for items in data]

It might be necessary to either always return a list from all transforms or to make a check here to not unpack something that should not be unpacked.
The compose could accept a unpack boolean flag which is provided to the apply_transform function.

Environment

================================
Printing MONAI config...
================================
MONAI version: 0.5.2
Numpy version: 1.20.1
Pytorch version: 1.6.0+cu101
MONAI flags: HAS_EXT = False, USE_COMPILED = False
MONAI rev id: feb3a334b7bbf302b13a6da80e0b022a4cf75a4e

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: 3.2.1
scikit-image version: 0.17.2
Pillow version: 8.1.0
Tensorboard version: 2.2.2
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: NOT INSTALLED or UNKNOWN VERSION.
ITK version: NOT INSTALLED or UNKNOWN VERSION.
tqdm version: 4.61.1
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: NOT INSTALLED or UNKNOWN VERSION.

For details about installing the optional dependencies, please visit:
    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies


================================
Printing system config...
================================
`psutil` required for `print_system_info`

================================
Printing GPU config...
================================
Num GPUs: 1
Has CUDA: True
CUDA version: 10.1
cuDNN enabled: True
cuDNN version: 7604
Current device: 0
Library compiled for CUDA architectures: ['sm_37', 'sm_50', 'sm_60', 'sm_61', 'sm_70', 'sm_75', 'compute_37']
GPU 0 Name: Quadro K2200
GPU 0 Is integrated: False
GPU 0 Is multi GPU board: False
GPU 0 Multi processor count: 5
GPU 0 Total memory (GB): 4.0
GPU 0 CUDA capability (maj.min): 5.0

[Nic-Ma]

Hi @Spenhouet ,

Thanks for your interest here.

1. If you want to use the image meta data in a composed transform chain, I would recommend the dict version transforms instead of these array version transforms.
2. If you really want to use array version transforms in a chain, you can set "image_only=True" arg to many of them, then it will only return the image array.

Thanks.

[Spenhouet]

Hi @Nic-Ma,

The image_only option does not really make sense. See also the issue I raised here: #2424
The meta_data is important for medical images. To just throw it away does seem to be a bad idea in my opinion. Not sure if you should actually support this behavior at all.

[Nic-Ma]

Hi @Spenhouet ,

If you need to use the meta data in your case, I think you should consider to use the dict version transforms.

Thanks.

[wyli]

Agree that the metadata is important. I think the issue here is, with the metadata the array-based transforms do not work well with the default Compose operation, so the intended usage is to build your own Compose for example in the ImageDataset https://github.com/Project-MONAI/tutorials/blob/master/modules/image_dataset.ipynb. In this case map_item=False is used:

MONAI/monai/data/image_dataset.py

Lines 114 to 116 in a0afa60

	img = apply_transform(
	self.transform, (img, meta_data) if self.transform_with_metadata else img, map_items=False
	)

Currently by default, we don't use map_item=False or transform(*items) because there are transforms that generate a list of image patches (such as the patch samplings)

[Spenhouet]

With map_item=False it is no longer possible to pass in a list of items to the compose call, right? Or am I missing something?

Adding an additional boolean flag for unpack as proposed above:

def apply_transform(transform: Callable, data, map_items: bool = True, unpack=False):
    if isinstance(data, (list, tuple)) and map_items:
        return [
            transform(*item) if unpack and isinstance(item, tuple) else transform(item)
            for item in data
        ]
    return transform(data)

class ComposeUnpack(Compose):

    def __init__(self,
                 transforms: Optional[Union[Sequence[Callable], Callable]] = None,
                 unpack=False) -> None:
        super().__init__(transforms)
        self.unpack = unpack

    def __call__(self, input_):
        for _transform in self.transforms:
            input_ = apply_transform(_transform, input_, unpack=self.unpack)
        return input_

This works as expected.

compose_labels = ComposeUnpack(
    [
        LoadImage(),
        EnsureChannelFirstWithMeta(),
        SpacingWithMeta(pixdim=min_pixdim, mode=GridSampleMode.NEAREST),
    ],
    unpack=True,
)

labels = compose_labels(files_to_load)

@wyli Would such an additional parameter for compose be welcome? I can open up a pull request for that.

Here also the fix for the transforms with meta data:

class EnsureChannelFirstWithMeta(EnsureChannelFirst):

    def __call__(self, data_array: np.ndarray,
                 meta_data: Dict[str, Any]) -> Tuple[np.ndarray, Dict[str, Any]]:
        output_data = super().__call__(data_array, meta_data)

        dim = meta_data['dim'].copy()
        dim[0] = 4
        dim[4] = 1

        return output_data, {
            **meta_data,
            "dim": dim,
            "spatial_shape": np.array(output_data.shape, dtype=np.int16),
        }


class SpacingWithMeta(Spacing):

    def __call__(self, data_array: np.ndarray,
                 meta_data: Dict[str, Any]) -> Tuple[np.ndarray, Dict[str, Any]]:
        output_data, _, new_affine = super().__call__(data_array, meta_data['affine'])

        dim = meta_data['dim'].copy()
        dim[1:4] = output_data.shape[1:]

        pixdim = meta_data['pixdim'].copy()
        pixdim[1:4] = self.pixdim

        return output_data, {
            **meta_data,
            "dim": dim,
            "pixdim": pixdim,
            "affine": new_affine,
            "spatial_shape": np.array(output_data.shape, dtype=np.int16),
        }

[wyli]

sure, please feel free to have a non-breaking change to introduce unpack.
currently map_items=True means applying the transform to each image patch respectively [xform(patch_1), xform(patch_2)], and map_items=False means xform([patch_1, patch_2])

[Spenhouet]

@wyli I would like to come back to the example you posted before:

MONAI/monai/data/image_dataset.py

Lines 114 to 116 in a0afa60

	img = apply_transform(
	self.transform, (img, meta_data) if self.transform_with_metadata else img, map_items=False
	)

This will only work with the _TestCompose implementation in test_image_dataset.py. This does not work via the map_item=False flag. As I see it the implementation of ImageDataset isn't really useable at the moment. Maybe I'm missing something.

I will adjust this in my pull request.