Repository: Project-MONAI/MONAI
File: monai/transforms/utility/array.py (Python)

Code before:
        """
        Args:
            name: The transform name in TorchVision package.
            args: parameters for the TorchVision transform.
            kwargs: parameters for the TorchVision transform.

        """
        super().__init__()
        self.name = name
        transform, _ = optional_import("torchvision.transforms", "0.8.0", min_version, name=name)
        self.trans = transform(*args, **kwargs)

    def __call__(self, img: NdarrayOrTensor):
        """
        Args:
            img: PyTorch Tensor data for the TorchVision transform.

        """
        img_t, *_ = convert_data_type(img, torch.Tensor)

        out = self.trans(img_t)
        out, *_ = convert_to_dst_type(src=out, dst=img)
        return out


class TorchIO(Transform, RandomizableTrait):
    """
    This is a wrapper transform for TorchIO transforms based on the specified transform name and args.
    See https://torchio.readthedocs.io/transforms/transforms.html for more details.
    """

    backend = [TransformBackends.TORCH]

    def __init__(self, name: str, *args, **kwargs) -> None:
        """
        Args:
            name: The transform name in TorchIO package.
            args: parameters for the TorchIO transform.
            kwargs: parameters for the TorchIO transform.

        Note:
            The `p=` kwarg of TorchIO transforms control set the probability with which the transform is applied.
            You can specify the probability of applying the transform by passing either `prob` ot `p` in kwargs but'
            ' not both.
        """
        super().__init__()
        self.name = name
        transform, _ = optional_import("torchio.transforms", "0.18.0", min_version, name=name)

        if "prob" in kwargs:
            if "p" in kwargs:

Review comment (performance):
I would suggest remove this RandomizabelTrait inherit here otherwise all the torchio tranforms will be considered as random transform and will not be cached. I guess it's not what we expected? 
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/dataset.py#L910

To include random transforms, can we do something like torchvision transform wrapper?
https://github.com/Project-MONAI/MONAI/blob/4029c422f1da84cdf5e4ce546dbff33245cbe025/monai/transforms/utility/dictionary.py#L1385

Code after:


class TorchVision(Transform):
    """
    This is a wrapper transform for PyTorch TorchVision non-randomized transform based on the specified transform name and args.
    Data is converted to a torch.tensor before applying the transform and then converted back to the original data type.
    """

    backend = [TransformBackends.TORCH]

    def __init__(self, name: str, *args, **kwargs) -> None:
        """
        Args:
            name: The transform name in TorchVision package.
            args: parameters for the TorchVision transform.
            kwargs: parameters for the TorchVision transform.

        """
        super().__init__()
        self.name = name
        transform, _ = optional_import("torchvision.transforms", "0.8.0", min_version, name=name)
        self.trans = transform(*args, **kwargs)

    def __call__(self, img: NdarrayOrTensor):
        """
        Args:
            img: PyTorch Tensor data for the TorchVision transform.

        """
        img_t, *_ = convert_data_type(img, torch.Tensor)

        out = self.trans(img_t)
        out, *_ = convert_to_dst_type(src=out, dst=img)
        return out


class RandTorchVision(Transform, RandomizableTrait):
    """
    This is a wrapper transform for PyTorch TorchVision randomized transform based on the specified transform name and args.
    Data is converted to a torch.tensor before applying the transform and then converted back to the original data type.
    """

    backend = [TransformBackends.TORCH]

    def __init__(self, name: str, *args, **kwargs) -> None:
        """
        Args:
            name: The transform name in TorchVision package.
            args: parameters for the TorchVision transform.
            kwargs: parameters for the TorchVision transform.