Add To transform

Author: fepegar

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,46 @@
+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.
+
+    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],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.target = target
+        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