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] 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