fix: Rotation applied in surface_image_stencil() based on image orientation

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)
 

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
+    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())
     converter.SetOutputSpacing(grid.spacing().tolist())
     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

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