From 0cf0f15b63cb23dc26eb3b9c1aed93ec49b87cfa Mon Sep 17 00:00:00 2001
From: Andreas Schuh <aschuh@heartflow.com>
Date: Fri, 18 Oct 2024 10:19:36 +0000
Subject: [PATCH 1/3] enh: Modify numpy_to_vtk_matrix4x4 to accept different
 input shapes

---
 src/deepali/utils/vtk/numpy.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/deepali/utils/vtk/numpy.py b/src/deepali/utils/vtk/numpy.py
index 0fc5497..48a686e 100644
--- a/src/deepali/utils/vtk/numpy.py
+++ b/src/deepali/utils/vtk/numpy.py
@@ -55,10 +55,12 @@
 
 def numpy_to_vtk_matrix4x4(arr: np.ndarray) -> vtkMatrix4x4:
     """Create vtkMatrix4x4 from NumPy array."""
-    assert arr.shape == (4, 4)
+    if arr.shape not in [(3, 3), (3, 4), (4, 4)]:
+        raise ValueError("numpy_to_vtk_matrix4x4() 'arr' must have shape (3, 3), (3, 4), or (4, 4)")
     matrix = vtkMatrix4x4()
-    for i in range(4):
-        for j in range(4):
+    matrix.Identity()
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
             matrix.SetElement(i, j, arr[i, j])
     return matrix
 

From 1eab935527a73b3a436b59a8dc29c766e0ce9de7 Mon Sep 17 00:00:00 2001
From: Andreas Schuh <aschuh@heartflow.com>
Date: Fri, 18 Oct 2024 10:23:54 +0000
Subject: [PATCH 2/3] fix: GridAttrs initialization and grid.center property
 getter

---
 src/deepali/utils/simpleitk/grid.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/deepali/utils/simpleitk/grid.py b/src/deepali/utils/simpleitk/grid.py
index b7d7c18..ac7b34d 100644
--- a/src/deepali/utils/simpleitk/grid.py
+++ b/src/deepali/utils/simpleitk/grid.py
@@ -50,8 +50,8 @@ def __init__(
             if center is None:
                 origin = (0.0,) * ndim
             else:
-                offset = np.array(0.5 * n if n > 0 else 0 for n in self._int_size(size))
-                rotation = np.array(self.direction).reshape(ndim, ndim)
+                offset = np.array([0.5 * n if n > 0 else 0 for n in self._int_size(size)])
+                rotation = np.array(direction).reshape(ndim, ndim)
                 scaling = np.diag(spacing)
                 coords: NDArray = np.asanyarray(center) - np.matmul(rotation @ scaling, offset)
                 origin = tuple(float(x) for x in coords)
@@ -68,7 +68,7 @@ def __init__(
     @property
     def center(self) -> Tuple[float, ...]:
         r"""Get grid center point coordinates in world space."""
-        offset = np.array(0.5 * n if n > 0 else 0 for n in self.size)
+        offset = np.array([0.5 * n if n > 0 else 0 for n in self.size])
         coords: NDArray = self.origin + np.matmul(self.transform[:-1, :-1], offset)
         return tuple(float(x) for x in coords)
 

From e554ec36de0faf02056da8e3997f41a3af61c1f8 Mon Sep 17 00:00:00 2001
From: Andreas Schuh <aschuh@heartflow.com>
Date: Fri, 18 Oct 2024 10:20:22 +0000
Subject: [PATCH 3/3] fix: Rotation applied in surface_image_stencil() based on
 image orientation

---
 src/deepali/utils/vtk/image.py | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/deepali/utils/vtk/image.py b/src/deepali/utils/vtk/image.py
index e536c50..e90bac3 100644
--- a/src/deepali/utils/vtk/image.py
+++ b/src/deepali/utils/vtk/image.py
@@ -8,9 +8,9 @@
     vtkImageData,
     vtkImageStencilData,
     vtkImageStencilToImage,
-    vtkMatrixToLinearTransform,
     vtkPolyData,
     vtkPolyDataToImageStencil,
+    vtkTransform,
     vtkTransformPolyDataFilter,
 )
 
@@ -42,20 +42,21 @@ def surface_mesh_grid(*mesh: vtkPolyData, resolution: Optional[float] = None) ->
 
 
 def surface_image_stencil(mesh: vtkPolyData, grid: Grid) -> vtkImageStencilData:
-    r"""Convert vtkPolyData surface mesh to image stencil."""
-    max_index = [n - 1 for n in grid.size().tolist()]
-
-    rot = np.eye(4, dtype=np.float)
-    rot[:3, :3] = np.array(grid.direction).reshape(3, 3)
-    rot = numpy_to_vtk_matrix4x4(rot)
-
-    transform = vtkMatrixToLinearTransform()
-    transform.SetInput(rot)
-
+    r"""Convert vtkPolyData surface mesh to image stencil."""    
+    # Create the transform
+    transform = vtkTransform()
+    transform.Translate(grid.center().tolist())
+    transform.Concatenate(numpy_to_vtk_matrix4x4(grid.direction().numpy().T))  # type: ignore
+    transform.Translate(grid.center().neg().tolist())
+
+    # Apply the transform to the polydata
     transformer = vtkTransformPolyDataFilter()
     transformer.SetInputData(mesh)
     transformer.SetTransform(transform)
 
+    # Convert the transformed polydata to an image stencil
+    grid = Grid(size=grid.size(), spacing=grid.spacing(), center=grid.center())
+    max_index = [n - 1 for n in grid.size()]
     converter = vtkPolyDataToImageStencil()
     converter.SetInputConnection(transformer.GetOutputPort())
     converter.SetOutputOrigin(grid.origin().tolist())
@@ -63,6 +64,7 @@ def surface_image_stencil(mesh: vtkPolyData, grid: Grid) -> vtkImageStencilData:
     converter.SetOutputWholeExtent([0, max_index[0], 0, max_index[1], 0, max_index[2]])
     converter.Update()
 
+    # Get the output stencil
     stencil = vtkImageStencilData()
     stencil.DeepCopy(converter.GetOutput())
     return stencil