Non-breaking refactoring of the array transforms to dispatch by input data type (and by additional info.)

[wyli]

Multiple feature requests have described common use cases of metadata handling and computational backends for pre-/post-processing transformations.

The proposed solution

We agreed that these features are essential MONAI core capabilities, and as a first step, they should be implemented at the primary level of the array-based ‘vanilla’ transforms.

The current array-level transforms were initially (in MONAI v0.1) designed to facilitate NumPy ndarray processing using the NumPy APIs. These were further extended to support both torch tensor and numpy ndarray inputs and outputs (in MONAI v0.7). We have agreed to further this extension with clearer APIs at a finer granularity.

Goals and benefits

The goal is to decouple data handling logic depending on input data type, so that dictionary-level transform wrappers can leverage the APIs to simultaneously process images, coordinates, metadata, transform randomisable info, transform traceable and invertible info. And this potentially enables optimised transform chains based on the backend dispatched.

From the developer's perspective, extension/subclasses of the array-level transforms can choose to implement one or multiple backends, or simply overwrite the __call__ to bypassing the dispatch logic.

Details

A potential implementation in pseudo-code would be:

def dispatch_by_type(xform, data, **kwargs):
    xform_func = xform.dispatch_type_dict[str(type(data))]
    return xform_func(data, **kwargs)


class Transform:
    dispatch_type_dict={
        "np.ndarray": self.call_impl, 
        "torch.Tensor": self.call_torch,
        "dict": self.call_meta_dict,
    }
    
    def __init__(self):
        ...

    def call_impl(self, img_np, **kwargs):
        ...

    def call_torch(self, img_torch, **kwargs):
        ...

    def call_meta_dict(self, meta, **kwargs):
        ...

    def __call__(self, data, **kwargs):
        return dispatch_by_type(self, data, **kwargs)