Fix closed surface creation for older vtk

Closes OpenwaterHealth#237

The workaround here is just to manually apply the transform
of the vtkImageData to the vtkPolyData that is created,
only if the vtk version is below 9.3.0.

Author: ebrahimebrahim

src/openlifu/seg/skinseg.py

@@ -19,13 +19,28 @@
 import skimage.filters
 import skimage.measure
 import vtk
+from packaging.version import parse
 from scipy.interpolate import LinearNDInterpolator
 from scipy.ndimage import distance_transform_edt
 from vtk.util.numpy_support import numpy_to_vtk
 
 from openlifu.geo import cartesian_to_spherical
 
 
+def apply_affine_to_polydata(affine:np.ndarray, polydata:vtk.vtkPolyData) -> vtk.vtkPolyData:
+    """Apply an affine transform to a vtkPolyData."""
+    affine_vtkmat = vtk.vtkMatrix4x4()
+    for i in range(4):
+        for j in range(4):
+            affine_vtkmat.SetElement(i, j, affine[i, j])
+    affine_vtktransform = vtk.vtkTransform()
+    affine_vtktransform.SetMatrix(affine_vtkmat)
+    transform_filter = vtk.vtkTransformPolyDataFilter()
+    transform_filter.SetTransform(affine_vtktransform)
+    transform_filter.SetInputData(polydata)
+    transform_filter.Update()
+    return transform_filter.GetOutput()
+
 def take_largest_connected_component(mask: np.ndarray) -> np.ndarray:
     """Given a boolean image array (or any integer numpy array), return a mask of the largest connected component."""
     mask_labeled = skimage.measure.label(mask)
@@ -145,6 +160,23 @@ def vtk_img_from_array_and_affine(vol_array:np.ndarray, affine:np.ndarray) -> vt
 
     return vtk_img
 
+def affine_from_vtk_image_data(vtk_img:vtk.vtkImageData) -> np.ndarray:
+    """Get a 4x4 affine matrix out of a vtkImageData, a partial reverse to `vtk_img_from_array_and_affine`"""
+    origin = np.array(vtk_img.GetOrigin())
+    spacing = np.array(vtk_img.GetSpacing())
+
+    direction_vtk = vtk_img.GetDirectionMatrix()
+    direction = np.eye(3)
+    for i in range(3):
+        for j in range(3):
+            direction[i, j] = direction_vtk.GetElement(i, j)
+
+    affine = np.eye(4, dtype=float)
+    affine[:3, :3] = direction @ np.diag(spacing)
+    affine[:3, 3] = origin
+
+    return affine
+
 def create_closed_surface_from_labelmap(
     binary_labelmap:vtk.vtkImageData,
     decimation_factor:float=0.,
@@ -164,6 +196,17 @@ def create_closed_surface_from_labelmap(
     https://github.com/Slicer/Slicer/blob/677932127c73a6c78654d4afd9458a655a4eef63/Libs/vtkSegmentationCore/vtkBinaryLabelmapToClosedSurfaceConversionRule.cxx#L246-L476
     """
 
+    affine = None # Only needed if vtk version is less than 9.3.0
+    if parse(vtk.__version__) < parse("9.3.0"):
+        # In these older versions of vtk, the labelmap would not work.
+        affine = affine_from_vtk_image_data(binary_labelmap)
+        binary_labelmap.SetOrigin([0,0,0])
+        binary_labelmap.SetSpacing([1,1,1])
+        direction_matrix_vtk = vtk.vtkMatrix3x3()
+        direction_matrix_vtk.Identity()
+        binary_labelmap.SetDirectionMatrix(direction_matrix_vtk)
+
+
     # step 1: pad by 1 pixel all around with 0s, to ensure that the surface is still closed
     # even if the labelmap runs up against the image boundary.
     padder = vtk.vtkImageConstantPad()
@@ -223,6 +266,14 @@ def create_closed_surface_from_labelmap(
     normals.Update()
     surface_mesh = normals.GetOutput()
 
+    if parse(vtk.__version__) < parse("9.3.0"):
+        # In these older versions of vtk, the labelmap internal affine transform is not used correctly,
+        # so we manually apply the transform after the fact
+        surface_mesh = apply_affine_to_polydata(
+            affine,
+            surface_mesh,
+        )
+
     return surface_mesh
 
 def spherical_interpolator_from_mesh(

tests/test_skinseg.py

@@ -9,6 +9,8 @@
 from scipy.stats import skew
 
 from openlifu.seg.skinseg import (
+    affine_from_vtk_image_data,
+    apply_affine_to_polydata,
     cartesian_to_spherical,
     compute_foreground_mask,
     create_closed_surface_from_labelmap,
@@ -37,20 +39,6 @@ def add_ball(
 
     volume[dist_sq <= radius**2] = value
 
-def apply_affine_to_polydata(affine:np.ndarray, polydata:vtk.vtkPolyData) -> vtk.vtkPolyData:
-    """Apply an affine transform to a vtkPolyData"""
-    affine_vtkmat = vtk.vtkMatrix4x4()
-    for i in range(4):
-        for j in range(4):
-            affine_vtkmat.SetElement(i, j, affine[i, j])
-    affine_vtktransform = vtk.vtkTransform()
-    affine_vtktransform.SetMatrix(affine_vtkmat)
-    transform_filter = vtk.vtkTransformPolyDataFilter()
-    transform_filter.SetTransform(affine_vtktransform)
-    transform_filter.SetInputData(polydata)
-    transform_filter.Update()
-    return transform_filter.GetOutput()
-
 def test_take_largest_connected_component():
     vol_array = np.zeros((20,20,20))
     add_ball(vol_array, (5,5,5), 4) # ball of radius 4 at (5,5,5)
@@ -84,6 +72,16 @@ def test_vtk_img_from_array_and_affine():
 
     assert vtk_img.GetPointData().GetScalars().GetTuple1(point_id) == pytest.approx(vol_array[i,j,k])
 
+def test_affine_from_vtk_image_data():
+    rng = np.random.default_rng(716)
+    vol_array = rng.random((5,4,3))
+    affine = np.eye(4)
+    affine[:3,:3] = expm((lambda A: (A - A.T)/2)(rng.normal(size=(3,3)))) # generate a random orthogonal matrix
+    affine[:3,3] = rng.random(3) # generate a random origin
+    vtk_img = vtk_img_from_array_and_affine(vol_array, affine)
+    affine_reconstructed = affine_from_vtk_image_data(vtk_img)
+    assert np.allclose(affine, affine_reconstructed)
+
 def test_create_closed_surface_from_labelmap():
     # create a ball of radius 7 for a labelmap
     labelmap = np.zeros((20,20,20))