Geometric transform implementations

monai/data/meta_tensor.py

@@ -24,7 +24,7 @@
 from monai.data.meta_obj import MetaObj, get_track_meta
 from monai.data.utils import affine_to_spacing, decollate_batch, list_data_collate, remove_extra_metadata
 from monai.utils import look_up_option
-from monai.utils.enums import LazyAttr, MetaKeys, PostFix, SpaceKeys
+from monai.utils.enums import KindKeys, LazyAttr, MetaKeys, PostFix, SpaceKeys
 from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_numpy, convert_to_tensor
 
 __all__ = ["MetaTensor"]
@@ -345,6 +345,10 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
     def get_default_affine(dtype=torch.float64) -> torch.Tensor:
         return torch.eye(4, device=torch.device("cpu"), dtype=dtype)
 
+    @staticmethod
+    def get_default_kind() -> str:
+        return KindKeys.PIXEL
+
     def as_tensor(self) -> torch.Tensor:
         """
         Return the `MetaTensor` as a `torch.Tensor`.
@@ -469,13 +473,30 @@ def affine(self, d: NdarrayTensor) -> None:
         """Set the affine."""
         self.meta[MetaKeys.AFFINE] = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.float64)
 
+    @property
+    def kind(self) -> str:
+        """Get the data kind. Defaults to ``KindKeys.PIXEL``"""
+        return self.meta.get(MetaKeys.KIND, self.get_default_kind())  # type: ignore
+
+    @kind.setter
+    def kind(self, d: str) -> None:
+        """Set the data kind."""
+        self.meta[MetaKeys.KIND] = d
+
     @property
     def pixdim(self):
         """Get the spacing"""
         if self.is_batch:
             return [affine_to_spacing(a) for a in self.affine]
         return affine_to_spacing(self.affine)
 
+    # @property
+    # def ndims(self):
+    #     # TODO: this will be wrong when there are batches; review
+    #     if self.kind == KindKeys.POINT:
+    #         return self.shape[2] - 1
+    #     return len(self.shape)
+
     def peek_pending_shape(self):
         """
         Get the currently expected spatial shape as if all the pending operations are executed.

monai/transforms/io/array.py

@@ -46,6 +46,7 @@
 from monai.utils import GridSamplePadMode
 from monai.utils import ImageMetaKey as Key
 from monai.utils import OptionalImportError, convert_to_dst_type, ensure_tuple, look_up_option, optional_import
+from monai.utils.enums import KindKeys, MetaKeys
 
 nib, _ = optional_import("nibabel")
 Image, _ = optional_import("PIL.Image")
@@ -280,6 +281,7 @@ def __call__(self, filename: Sequence[PathLike] | PathLike, reader: ImageReader
 
         img_array: NdarrayOrTensor
         img_array, meta_data = reader.get_data(img)
+        meta_data[MetaKeys.KIND] = KindKeys.PIXEL
         img_array = convert_to_dst_type(img_array, dst=img_array, dtype=self.dtype)[0]
         if not isinstance(meta_data, dict):
             raise ValueError(f"`meta_data` must be a dict, got type {type(meta_data)}.")

monai/transforms/io/functional.py

@@ -0,0 +1,84 @@
+import json
+
+import numpy as np
+
+import torch
+
+from monai.data.meta_tensor import MetaTensor
+
+from monai.utils.enums import KindKeys
+
+
+def load_geometry(file, image, origin):
+    """
+    Load geometry from a file and optionally map it to another coordinate space.
+    """
+    geometry = json.load(file)
+    geometry_schema = geometry.get("schema", None)
+    if geometry_schema is None:
+        raise ValueError("Geometry import issue: missing 'schema' entry")
+    elif "geometry" not in geometry_schema:
+        raise ValueError(f"Geometry import issue: 'schema' entry must contain 'geometry' key, got: {geometry_schema}")
+
+    if "points" not in geometry:
+        raise ValueError("Geometry import issue: missing 'points' entry")
+
+    points = geometry["points"]
+    if not isinstance(points, list):
+        raise ValueError(f"Geometry import issue: 'points' entry must be a list, got: {type(points)}")
+
+    if len(points) > 0:
+        first_len = None
+        for p in points:
+            if first_len is None:
+                first_len = len(p)
+            if len(p) != first_len:
+                raise ValueError("Geometry import issue: 'points' entry contains inconsistent point lengths")
+
+    points = np.asarray(points)
+    points = np.concatenate((points, np.ones((points.shape[0], 1))), axis=1)
+    points = torch.as_tensor(points, dtype=torch.float32)
+    points = MetaTensor(points)
+    points.kind = KindKeys.POINT
+
+    return points
+
+"""
+{
+            "schema": {
+                "geometry": "point"
+            },
+            "points": [
+                [0, 0, 0],
+                [0, 1, 0],
+                [0, 0, 1],
+                [0, 1, 1],
+                [1, 0, 0],
+                [1, 1, 0],
+                [1, 0, 1],
+                [1, 1, 1],
+            ]
+        }
+"""
+
+def save_geometry(data, file, image, origin):
+    """
+    Load geometry from a file and optionally map it to another coordinate space.
+    """
+    if not isinstance(data, MetaTensor):
+        raise ValueError(f"Geometry export issue: data must be a MetaTensor, got: {type(data)}")
+    if data.kind != KindKeys.POINT:
+        raise ValueError(f"Geometry export issue: geometry must be a point {KindKeys.POINT}")
+    geometry = data.detach().cpu().numpy()
+    geometry = geometry[:, :-1].tolist()
+
+    schema = {
+        "schema": {
+            "geometry": "point"
+        },
+        "points":
+            geometry
+    }
+
+    geometry = json.dump(schema, file)
+    return None

monai/transforms/lazy/functional.py

@@ -17,6 +17,7 @@
 
 from monai.apps.utils import get_logger
 from monai.config import NdarrayOrTensor
+from monai.data.meta_obj import get_track_meta, MetaObj
 from monai.data.meta_tensor import MetaTensor
 from monai.data.utils import to_affine_nd
 from monai.transforms.lazy.utils import (
@@ -28,9 +29,9 @@
 )
 from monai.transforms.traits import LazyTrait
 from monai.transforms.transform import MapTransform
-from monai.utils import LazyAttr, look_up_option
+from monai.utils import LazyAttr, MetaKeys, convert_to_tensor, look_up_option
 
-__all__ = ["apply_pending_transforms", "apply_pending_transforms_in_order", "apply_pending"]
+__all__ = ["apply_pending_transforms", "apply_pending_transforms_in_order", "apply_pending", "apply_to_geometry"]
 
 __override_keywords = {"mode", "padding_mode", "dtype", "align_corners", "resample_mode", "device"}
 
@@ -293,3 +294,56 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None,
         for p in pending:
             data.push_applied_operation(p)
     return data, pending
+
+
+def apply_to_geometry(
+        data: torch.Tensor,
+        meta_info: dict | MetaObj | None = None,
+        transform: torch.Tensor | None = None,
+):
+    """
+    Apply an affine geometric transform or deformation field to geometry.
+    At present this is limited to the transformation of points.
+
+    The points must be provided as a tensor and must be  compatible with a homogeneous
+    transform. This means that:
+     - 2D points are of the form (x, y, 1)
+     - 3D points are of the form (x, y, z, 1)
+
+    The affine transform or deformation field is applied to the the points and a tensor of
+    the same shape as the input tensor is returned.
+
+    Args:
+        data: the tensor of points to be transformed.
+        meta_info: the metadata containing the affine transformation
+    """
+
+    if meta_info is None and transform is None:
+        raise ValueError("either meta_info or transform must be provided")
+    if meta_info is not None and transform is not None:
+        raise ValueError("only one of meta_info or transform can be provided")
+
+    if not isinstance(data, (torch.Tensor, MetaTensor)):
+        raise TypeError(f"data {type(data)} must be a torch.Tensor or MetaTensor")
+
+    data = convert_to_tensor(data, track_meta=get_track_meta())
+
+    if meta_info is not None:
+        transform_ = meta_info.meta[MetaKeys.AFFINE]
+    else:
+        transform_ = transform
+
+    if transform_.dtype != data.dtype:
+        transform_ = transform_.to(data.dtype)
+        if data.shape[-1] == 3 and transform_.shape[0] == 4:
+            transform_[2, 0:2] = transform_[3, 0:2]
+            transform_[2, 2] = transform_[3, 3]
+            transform_[0:2, 2] = transform_[0:2, 3]
+            transform_ = transform_[:-1, :-1]
+
+    if data.shape[-1] != transform_.shape[0]:
+        raise ValueError(f"final element of data.shape {data.shape} must match transform shape {transform_.shape}")
+
+    result = torch.matmul(data, transform_.T)
+
+    return result

monai/transforms/spatial/array.py

@@ -44,6 +44,7 @@
 )
 from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import LazyTransform, Randomizable, RandomizableTransform, Transform
+from monai.transforms.utility.functional import traced_no_op
 from monai.transforms.utils import (
     create_control_grid,
     create_grid,
@@ -54,6 +55,7 @@
     map_spatial_axes,
     resolves_modes,
     scale_affine,
+    spatial_dims_from_tensorlike,
 )
 from monai.transforms.utils_pytorch_numpy_unification import argsort, argwhere, linalg_inv, moveaxis
 from monai.utils import (
@@ -72,7 +74,7 @@
     issequenceiterable,
     optional_import,
 )
-from monai.utils.enums import GridPatchSort, PatchKeys, TraceKeys, TransformBackends
+from monai.utils.enums import GridPatchSort, KindKeys, PatchKeys, TraceKeys, TransformBackends
 from monai.utils.misc import ImageMetaKey as Key
 from monai.utils.module import look_up_option
 from monai.utils.type_conversion import convert_data_type, get_equivalent_dtype, get_torch_dtype_from_string
@@ -482,6 +484,12 @@ def __call__(
             data tensor or MetaTensor (resampled into `self.pixdim`).
 
         """
+        lazy_ = self.lazy if lazy is None else lazy
+        if isinstance(data_array, MetaTensor) and data_array.kind == KindKeys.POINT:
+            warnings.warn("Spacing transform is not applied to point data.")
+            data_array = traced_no_op(data_array, lazy_, self.get_transform_info())
+            return data_array
+
         original_spatial_shape = (
             data_array.peek_pending_shape() if isinstance(data_array, MetaTensor) else data_array.shape[1:]
         )
@@ -521,7 +529,6 @@ def __call__(
         new_affine[:sr, -1] = offset[:sr]
 
         actual_shape = list(output_shape) if output_spatial_shape is None else output_spatial_shape
-        lazy_ = self.lazy if lazy is None else lazy
         data_array = self.sp_resample(
             data_array,
             dst_affine=torch.as_tensor(new_affine),
@@ -1089,7 +1096,8 @@ def __call__(
                 during initialization for this call. Defaults to None.
         """
         img = convert_to_tensor(img, track_meta=get_track_meta())
-        _zoom = ensure_tuple_rep(self.zoom, img.ndim - 1)  # match the spatial image dim
+        spatial_dims = spatial_dims_from_tensorlike(img)
+        _zoom = ensure_tuple_rep(self.zoom, spatial_dims)  # match the spatial image dim
         _mode = self.mode if mode is None else mode
         _padding_mode = padding_mode or self.padding_mode
         _align_corners = self.align_corners if align_corners is None else align_corners