Add To transform (#1345)

Author: fepegar

docs/source/transforms/preprocessing.rst

@@ -46,6 +46,13 @@ Intensity
     :show-inheritance:
 
 
+:class:`To`
+~~~~~~~~~~~
+
+.. autoclass:: To
+    :show-inheritance:
+
+
 .. currentmodule:: torchio.transforms.preprocessing.intensity
 
 
@@ -56,23 +63,20 @@ Intensity
     :show-inheritance:
 
 
-
-
-
 Spatial
 -------
 
 .. currentmodule:: torchio.transforms
 
 :class:`CropOrPad`
-~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: CropOrPad
     :show-inheritance:
     :members: _get_six_bounds_parameters
 
 :class:`ToOrientation`
-~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: ToOrientation
     :show-inheritance:
@@ -106,7 +110,7 @@ Spatial
 
 
 :class:`CopyAffine`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: CopyAffine
     :show-inheritance:
@@ -131,14 +135,14 @@ Label
 
 
 :class:`RemapLabels`
-~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: RemapLabels
     :show-inheritance:
 
 
 :class:`RemoveLabels`
-~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: RemoveLabels
     :show-inheritance:

src/torchio/transforms/__init__.py

@@ -53,6 +53,7 @@
 from .preprocessing import RescaleIntensity
 from .preprocessing import Resize
 from .preprocessing import SequentialLabels
+from .preprocessing import To
 from .preprocessing import ToCanonical
 from .preprocessing import ToOrientation
 from .preprocessing import ZNormalization
@@ -101,6 +102,7 @@
     'Crop',
     'Resize',
     'Resample',
+    'To',
     'ToCanonical',
     'ToOrientation',
     'ZNormalization',

src/torchio/transforms/preprocessing/__init__.py

@@ -2,6 +2,7 @@
 from .intensity.histogram_standardization import HistogramStandardization
 from .intensity.mask import Mask
 from .intensity.rescale import RescaleIntensity
+from .intensity.to import To
 from .intensity.z_normalization import ZNormalization
 from .label.contour import Contour
 from .label.keep_largest_component import KeepLargestComponent
@@ -31,6 +32,7 @@
     'EnsureShapeMultiple',
     'Mask',
     'RescaleIntensity',
+    'To',
     'Clamp',
     'ZNormalization',
     'HistogramStandardization',

src/torchio/transforms/preprocessing/intensity/to.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from ....data.image import ScalarImage
+from ....data.subject import Subject
+from ...intensity_transform import IntensityTransform
+
+
+class To(IntensityTransform):
+    """Convert the image tensor data type and/or device.
+
+    This transform is a thin wrapper around :func:`torch.Tensor.to`.
+
+    Args:
+        target: First argument to :func:`torch.Tensor.to`.
+        to_kwargs: Additional keyword arguments to pass to :func:`torch.Tensor.to`.
+
+    Example:
+        >>> import torchio as tio
+        >>> ct = tio.datasets.Slicer('CTChest').CT_chest
+        >>> clamp = tio.Clamp(out_min=-1000, out_max=1000)
+        >>> ct_clamped = clamp(ct)
+        >>> rescale = tio.RescaleIntensity(in_min_max=(-1000, 1000), out_min_max=(0, 255))
+        >>> ct_rescaled = rescale(ct_clamped)
+        >>> to_uint8 = tio.To(torch.uint8)
+        >>> ct_uint8 = to_uint8(ct_rescaled)
+    """
+
+    def __init__(
+        self,
+        target: str | torch.dtype | torch.device,
+        to_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.target = target
+        if to_kwargs is None:
+            to_kwargs = {}
+        self.to_kwargs = to_kwargs
+        self.args_names = ['target', 'to_kwargs']
+
+    def apply_transform(self, subject: Subject) -> Subject:
+        for image in self.get_images(subject):
+            assert isinstance(image, ScalarImage)
+            image.set_data(image.data.to(self.target, **self.to_kwargs))
+        return subject

tests/transforms/preprocessing/test_to.py

@@ -0,0 +1,17 @@
+import torch
+
+import torchio as tio
+
+from ...utils import TorchioTestCase
+
+
+class TestTo(TorchioTestCase):
+    """Tests for :class:`tio.To` class."""
+
+    def test_to(self):
+        transform = tio.To(torch.int)
+        tensor = 10 * torch.rand(2, 3, 4, 5)
+        image = tio.ScalarImage(tensor=tensor)
+        transformed = transform(image)
+        assert image.data.dtype == torch.float32
+        assert transformed.data.dtype == torch.int