From 05888724c85fc3c61f2b7e473f7e5d5992073f77 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 | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/deepali/utils/vtk/image.py b/src/deepali/utils/vtk/image.py
index e536c50..04c0ae6 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,27 +42,29 @@ 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
+    stencil_grid = Grid(size=grid.size(), spacing=grid.spacing(), center=grid.center())
+    stencil_extent = [0, grid.size(0) - 1, 0, grid.size(1) - 1, 0, grid.size(2) - 1]
     converter = vtkPolyDataToImageStencil()
     converter.SetInputConnection(transformer.GetOutputPort())
-    converter.SetOutputOrigin(grid.origin().tolist())
-    converter.SetOutputSpacing(grid.spacing().tolist())
-    converter.SetOutputWholeExtent([0, max_index[0], 0, max_index[1], 0, max_index[2]])
+    converter.SetOutputOrigin(stencil_grid.origin().tolist())
+    converter.SetOutputSpacing(stencil_grid.spacing().tolist())
+    converter.SetOutputWholeExtent(stencil_extent)
     converter.Update()
 
+    # Get the output stencil
     stencil = vtkImageStencilData()
     stencil.DeepCopy(converter.GetOutput())
     return stencil