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

Code before:

    def __init__(self, keys: KeysCollection, name: str, allow_missing_keys: bool = False, *args, **kwargs) -> None:
        MapTransform.__init__(self, keys, allow_missing_keys)
        self.name = name
        self.trans = TorchVision(name, *args, **kwargs)

    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            d[key] = self.trans(d[key])
        return d


class TorchIOd(MapTransform, RandomizableTrait):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.TorchIO` for transforms.
    All transforms in TorchIO can be applied randomly with probability p by specifying the `p=` argument.
    """

    backend = TorchIO.backend

    def __init__(
        self,
        keys: KeysCollection,
        name: str,
        apply_same_transform: bool = False,
        allow_missing_keys: bool = False,
        *args,
        **kwargs,
    ) -> None:
        """
        Args:
            keys: keys of the corresponding items to be transformed.
                See also: :py:class:`monai.transforms.compose.MapTransform`
            name: The transform name in TorchIO package.
            apply_same_transform: whether to apply the same transform for all the items specified by `keys`.
            allow_missing_keys: don't raise exception if key is missing.
            args: parameters for the TorchIO transform.
            kwargs: parameters for the TorchIO transform.

        """
        super().__init__(keys, allow_missing_keys)
        self.name = name
        self.apply_same_transform = apply_same_transform

        if self.apply_same_transform:
            kwargs["include"] = self.keys

        self.trans = TorchIO(name, *args, **kwargs)

    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:

Review comment (refactor):
The same comments about the RandomizableTrait inherit with the array version.
And what's this `apply_same_transform` arg used for? In MONAI, we always apply the same transform on the  sequence keys, such as image and label. Is there a typical case here users may want apply different transforms on the different keys? If not, perhaps we can remove this one.

Code after:
        for key in self.key_iterator(d):
            d[key] = self.trans(d[key])
        return d


class RandTorchVisiond(MapTransform, RandomizableTrait):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.TorchVision` for randomized transforms.
    For deterministic non-randomized transforms of TorchVision use :py:class:`monai.transforms.TorchVisiond`.

    Args:
        keys: keys of the corresponding items to be transformed.
            See also: :py:class:`monai.transforms.compose.MapTransform`
        name: The transform name in TorchVision package.
        allow_missing_keys: don't raise exception if key is missing.
        args: parameters for the TorchVision transform.
        kwargs: parameters for the TorchVision transform.

    Note:

        - As most of the TorchVision transforms only work for PIL image and PyTorch Tensor, this transform expects input
          data to be dict of PyTorch Tensors. Users should call `ToTensord` transform first to convert Numpy to Tensor.
        - This class inherits the ``Randomizable`` purely to prevent any dataset caching to skip the transform
          computation. If the random factor of the underlying torchvision transform is not derived from `self.R`,
          the results may not be deterministic. See Also: :py:class:`monai.transforms.Randomizable`.

    """

    backend = TorchVision.backend

    def __init__(self, keys: KeysCollection, name: str, allow_missing_keys: bool = False, *args, **kwargs) -> None:
        MapTransform.__init__(self, keys, allow_missing_keys)
        self.name = name
        self.trans = TorchVision(name, *args, **kwargs)

    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            d[key] = self.trans(d[key])
        return d


class TorchIOd(MapTransform):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.TorchIO` for non-randomized transforms.
    For randomized transforms of TorchIO use :py:class:`monai.transforms.RandTorchIOd`.
    """

    backend = TorchIO.backend

    def __init__(self, keys: KeysCollection, name: str, allow_missing_keys: bool = False, *args, **kwargs) -> None: