Image space and World Space convert transforms

monai/transforms/utility/array.py

@@ -31,7 +31,7 @@
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.meta_obj import get_track_meta
 from monai.data.meta_tensor import MetaTensor
-from monai.data.utils import is_no_channel, no_collation
+from monai.data.utils import is_no_channel, no_collation, orientation_ras_lps
 from monai.networks.layers.simplelayers import (
     ApplyFilter,
     EllipticalFilter,
@@ -42,16 +42,17 @@
     SharpenFilter,
     median_filter,
 )
-from monai.transforms.inverse import InvertibleTransform
+from monai.transforms.inverse import InvertibleTransform, TraceableTransform
 from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import Randomizable, RandomizableTrait, RandomizableTransform, Transform
 from monai.transforms.utils import (
+    apply_affine_to_points,
     extreme_points_to_image,
     get_extreme_points,
     map_binary_to_indices,
     map_classes_to_indices,
 )
-from monai.transforms.utils_pytorch_numpy_unification import concatenate, in1d, moveaxis, unravel_indices
+from monai.transforms.utils_pytorch_numpy_unification import concatenate, in1d, linalg_inv, moveaxis, unravel_indices
 from monai.utils import (
     MetaKeys,
     TraceKeys,
@@ -66,7 +67,7 @@
 )
 from monai.utils.enums import TransformBackends
 from monai.utils.misc import is_module_ver_at_least
-from monai.utils.type_conversion import convert_to_dst_type, get_equivalent_dtype
+from monai.utils.type_conversion import convert_to_dst_type, get_dtype_string, get_equivalent_dtype
 
 PILImageImage, has_pil = optional_import("PIL.Image", name="Image")
 pil_image_fromarray, _ = optional_import("PIL.Image", name="fromarray")
@@ -1715,3 +1716,111 @@ def __call__(self, img: NdarrayOrTensor, meta_dict: Mapping | None = None) -> Nd
         if self._do_transform:
             img = self.filter(img)
         return img
+
+
+class ApplyTransformToPoints(InvertibleTransform, Transform):
+    """
+    Transform points between image coordinates and world coordinates.
+
+    Args:
+        dtype: The desired data type for the output.
+        affine: A 3x3 or 4x4 affine transformation matrix applied to points. Typically, this matrix originates from the image.
+        invert_affine: Whether to invert the affine transformation matrix applied to the points. Defaults to ``True``.
+            Typically, the affine matrix is derived from the image, while the points are in world coordinates.
+            If you want to align the points with the image, set this to ``True``. Otherwise, set it to ``False``.
+        affine_lps_to_ras: Defaults to ``False``. Set this to ``True`` if all of the following are true:
+            1) The image is read by `ITKReader`,
+            2) The `ITKReader` has `affine_lps_to_ras=True`,
+            3) The data is in world coordinates.
+            This ensures the correct application of the affine transformation between LPS (left-posterior-superior)
+            and RAS (right-anterior-superior) coordinate systems. This argument ensures the points and the affine
+            matrix are in the same coordinate system.
+    """
+
+    def __init__(
+        self,
+        dtype: DtypeLike | torch.dtype = torch.float64,
+        affine: torch.Tensor | None = None,
+        invert_affine: bool = True,
+        affine_lps_to_ras: bool = False,
+    ) -> None:
+        self.dtype = dtype
+        self.affine = affine
+        self.invert_affine = invert_affine
+        self.affine_lps_to_ras = affine_lps_to_ras
+
+    def transform_coordinates(self, data: torch.Tensor, affine: torch.Tensor | None = None):
+        """
+        Transform coordinates using an affine transformation matrix.
+
+        Args:
+            data: The input coordinates, assume to be in shape (C, N, 2 or 3).
+            affine: A 3x3 or 4x4 affine transformation matrix.
+
+        Returns:
+            Transformed coordinates.
+        """
+        data = convert_to_tensor(data, track_meta=get_track_meta())
+        applied_affine = data.affine if isinstance(data, MetaTensor) else None
+
+        if affine is None and self.invert_affine:
+            raise ValueError("affine must be provided when invert_affine is True.")
+
+        affine = applied_affine if affine is None else affine
+        affine = convert_data_type(affine, dtype=torch.float64)[0]  # always convert to float64 for affine
+        original_affine = affine
+        if self.affine_lps_to_ras:
+            affine = orientation_ras_lps(affine)
+
+        _affine = affine
+        if self.invert_affine:
+            _affine = linalg_inv(affine)
+            # consider the affine transformation already applied to the data in the world space
+            if applied_affine is not None:
+                _affine = _affine @ linalg_inv(applied_affine)
+        out = apply_affine_to_points(data, _affine, self.dtype)
+
+        extra_info = {
+            "invert_affine": self.invert_affine,
+            "dtype": get_dtype_string(self.dtype),
+            "image_affine": original_affine,
+            "affine_lps_to_ras": self.affine_lps_to_ras,
+        }
+        xform = original_affine if self.invert_affine else linalg_inv(original_affine)
+        meta_info = TraceableTransform.track_transform_meta(
+            data, affine=xform, extra_info=extra_info, transform_info=self.get_transform_info()
+        )
+
+        return out, meta_info
+
+    def __call__(self, data: torch.Tensor, affine: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Args:
+            data: The input coordinates, assume to be in shape (C, N, 2 or 3).
+            affine: A 3x3 or 4x4 affine transformation matrix, this argument will take precedence over ``self.affine``.
+        """
+        if data.ndim != 3 or data.shape[-1] not in (2, 3):
+            raise ValueError(f"data should be in shape (C, N, 2 or 3), got {data.shape}.")
+        affine = self.affine if affine is None else affine
+        if affine is not None and affine.shape not in ((3, 3), (4, 4)):
+            raise ValueError(f"affine should be in shape (3, 3) or (4, 4), got {affine.shape}.")
+
+        out, meta_info = self.transform_coordinates(data, affine)
+
+        return out.copy_meta_from(meta_info) if isinstance(out, MetaTensor) else out
+
+    def inverse(self, data: torch.Tensor) -> torch.Tensor:
+        transform = self.pop_transform(data)
+        # Create inverse transform
+        dtype = transform[TraceKeys.EXTRA_INFO]["dtype"]
+        invert_affine = not transform[TraceKeys.EXTRA_INFO]["invert_affine"]
+        affine = transform[TraceKeys.EXTRA_INFO]["image_affine"]
+        affine_lps_to_ras = transform[TraceKeys.EXTRA_INFO]["affine_lps_to_ras"]
+        inverse_transform = ApplyTransformToPoints(
+            dtype=dtype, invert_affine=invert_affine, affine_lps_to_ras=affine_lps_to_ras
+        )
+        # Apply inverse
+        with inverse_transform.trace_transform(False):
+            data = inverse_transform(data, affine)
+
+        return data

monai/transforms/utility/dictionary.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import re
+import warnings
 from collections.abc import Callable, Hashable, Mapping
 from copy import deepcopy
 from typing import Any, Sequence, cast
@@ -35,6 +36,7 @@
 from monai.transforms.utility.array import (
     AddCoordinateChannels,
     AddExtremePointsChannel,
+    ApplyTransformToPoints,
     AsChannelLast,
     CastToType,
     ClassesToIndices,
@@ -1740,6 +1742,68 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
         return d
 
 
+class ApplyTransformToPointsd(MapTransform, InvertibleTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.ApplyTransformToPoints`.
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: monai.transforms.MapTransform
+        refer_key: the key of the reference item used for transformation.
+        dtype: The desired data type for the output.
+        affine: A 3x3 or 4x4 affine transformation matrix applied to points. Typically, this matrix originates from the image.
+        invert_affine: Whether to invert the affine transformation matrix applied to the points. Defaults to ``True``.
+            Typically, the affine matrix is derived from the image, while the points are in world coordinates.
+            If you want to align the points with the image, set this to ``True``. Otherwise, set it to ``False``.
+        affine_lps_to_ras: Defaults to ``False``. Set this to ``True`` if all of the following are true:
+            1) The image is read by `ITKReader`,
+            2) The `ITKReader` has `affine_lps_to_ras=True`,
+            3) The data is in world coordinates.
+            This ensures the correct application of the affine transformation between LPS (left-posterior-superior)
+            and RAS (right-anterior-superior) coordinate systems. This argument ensures the points and the affine
+            matrix are in the same coordinate system.
+        allow_missing_keys: Don't raise exception if key is missing.
+    """
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        refer_key: str | None = None,
+        dtype: DtypeLike | torch.dtype = torch.float64,
+        affine: torch.Tensor | None = None,
+        invert_affine: bool = True,
+        affine_lps_to_ras: bool = False,
+        allow_missing_keys: bool = False,
+    ):
+        MapTransform.__init__(self, keys, allow_missing_keys)
+        self.refer_key = refer_key
+        self.affine = affine
+        if self.refer_key is None and self.affine is None:
+            warnings.warn("No reference data or affine matrix is provided, will use the affine derived from the data.")
+        self.converter = ApplyTransformToPoints(
+            dtype=dtype, invert_affine=invert_affine, affine_lps_to_ras=affine_lps_to_ras
+        )
+
+    def __call__(self, data: Mapping[Hashable, torch.Tensor]):
+        d = dict(data)
+        refer_data = d[self.refer_key] if self.refer_key is not None else None
+        if isinstance(refer_data, MetaTensor):
+            affine = refer_data.affine
+        else:
+            warnings.warn("No reference affine find in the refer key, will use the affine derived from the data.")
+            affine = self.affine
+        for key in self.key_iterator(d):
+            coords = d[key]
+            d[key] = self.converter(coords, affine)
+        return d
+
+    def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torch.Tensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.converter.inverse(d[key])
+        return d
+
+
 RandImageFilterD = RandImageFilterDict = RandImageFilterd
 ImageFilterD = ImageFilterDict = ImageFilterd
 IdentityD = IdentityDict = Identityd

monai/transforms/utils.py

@@ -27,6 +27,7 @@
 import monai
 from monai.config import DtypeLike, IndexSelection
 from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
+from monai.data.utils import to_affine_nd
 from monai.networks.layers import GaussianFilter
 from monai.networks.utils import meshgrid_ij
 from monai.transforms.compose import Compose
@@ -35,6 +36,7 @@
 from monai.transforms.utils_pytorch_numpy_unification import (
     any_np_pt,
     ascontiguousarray,
+    concatenate,
     cumsum,
     isfinite,
     nonzero,
@@ -2509,5 +2511,25 @@ def distance_transform_edt(
     return convert_data_type(r_vals[0] if len(r_vals) == 1 else r_vals, output_type=type(img), device=device)[0]
 
 
+def apply_affine_to_points(data: torch.Tensor, affine: torch.Tensor, dtype: DtypeLike | torch.dtype):
+    """
+    apply affine transformation to a set of points.
+
+    Args:
+        data: input data to apply affine transformation, should be a tensor of shape (C, N, 2 or 3).
+        affine: affine matrix to be applied, should be a tensor of shape (3, 3) or (4, 4).
+        dtype: output data dtype.
+    """
+    data_: torch.Tensor = convert_to_tensor(data, track_meta=False, dtype=torch.float64)
+    affine = to_affine_nd(data_.shape[-1], affine)
+
+    homogeneous = concatenate((data_, torch.ones((data_.shape[0], data_.shape[1], 1))), axis=2)
+    transformed_homogeneous = torch.matmul(homogeneous, affine.T)
+    transformed_coordinates = transformed_homogeneous[:, :, :-1]
+    out, *_ = convert_to_dst_type(transformed_coordinates, data, dtype=dtype)
+
+    return out
+
+
 if __name__ == "__main__":
     print_transform_backends()

monai/utils/__init__.py

@@ -148,6 +148,7 @@
     dtype_numpy_to_torch,
     dtype_torch_to_numpy,
     get_dtype,
+    get_dtype_string,
     get_equivalent_dtype,
     get_numpy_dtype_from_string,
     get_torch_dtype_from_string,

monai/utils/type_conversion.py

@@ -33,6 +33,7 @@
     "get_equivalent_dtype",
     "convert_data_type",
     "get_dtype",
+    "get_dtype_string",
     "convert_to_cupy",
     "convert_to_numpy",
     "convert_to_tensor",
@@ -102,6 +103,13 @@ def get_dtype(data: Any) -> DtypeLike | torch.dtype:
     return type(data)
 
 
+def get_dtype_string(dtype: DtypeLike | torch.dtype) -> str:
+    """Get a string representation of the dtype."""
+    if isinstance(dtype, torch.dtype):
+        return str(dtype)[6:]
+    return str(dtype)[3:]
+
+
 def convert_to_tensor(
     data: Any,
     dtype: DtypeLike | torch.dtype = None,