From d44b92ebcf3b8eb64c66a9f0bea9caf09ca686a7 Mon Sep 17 00:00:00 2001
From: Thomas Dowrick <tdowrick@gmail.com>
Date: Wed, 9 Sep 2020 14:23:47 +0100
Subject: [PATCH] Voxelise function can now take vtk data (e.g. vtkPolyData) as
 input, rather than just readgin from file

---
 sksurgeryvtk/models/voxelise.py          |  8 +++++--
 sksurgeryvtk/models/vtk_surface_model.py |  9 ++++++++
 tests/models/test_voxelise.py            | 29 ++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/sksurgeryvtk/models/voxelise.py b/sksurgeryvtk/models/voxelise.py
index 8c483a84..8d629345 100644
--- a/sksurgeryvtk/models/voxelise.py
+++ b/sksurgeryvtk/models/voxelise.py
@@ -211,8 +211,9 @@ def load_points_from_file(filename):
     mesh = reader.GetOutput()
     return mesh
 
-def voxelise(input_mesh: Union[np.ndarray, str],
-             output_grid: Union[vtk.vtkStructuredGrid, str] = None,
+def voxelise(input_mesh: Union[np.ndarray, vtk.vtkDataObject, str],
+             output_grid: Union[vtk.vtkStructuredGrid, str] \
+                 = None,
              array_name: str = "",
              size: float = 0.3,
              grid_elements: int = 64,
@@ -265,6 +266,9 @@ def voxelise(input_mesh: Union[np.ndarray, str],
     if isinstance(input_mesh, str):
         mesh = load_points_from_file(input_mesh)
 
+    elif isinstance(input_mesh, vtk.vtkDataObject):
+        mesh = input_mesh
+
     else:
         input_is_point_cloud = True
 
diff --git a/sksurgeryvtk/models/vtk_surface_model.py b/sksurgeryvtk/models/vtk_surface_model.py
index c570f2d4..b6800d4e 100644
--- a/sksurgeryvtk/models/vtk_surface_model.py
+++ b/sksurgeryvtk/models/vtk_surface_model.py
@@ -144,6 +144,15 @@ def get_points_as_numpy(self):
         as_numpy = numpy_support.vtk_to_numpy(vtk_points.GetData())
         return as_numpy
 
+    def get_vtk_data(self) -> vtk.vtkPolyData:
+        """Return vtk poly data for this object
+
+        :return: vtkPolyData
+        :rtype: vtk.vtkPolyData
+        """
+        self.transform_filter.Update()
+        return self.source
+
     def get_normals_as_numpy(self):
         """
          Returns the vtkPolyData point normals as a numpy array.
diff --git a/tests/models/test_voxelise.py b/tests/models/test_voxelise.py
index 23502bfb..96f62a91 100644
--- a/tests/models/test_voxelise.py
+++ b/tests/models/test_voxelise.py
@@ -2,6 +2,7 @@
 import pytest
 import numpy as np
 from sksurgeryvtk.models import voxelise
+from sksurgeryvtk.models import vtk_surface_model
 import vtk
 from vtk.util import numpy_support
 
@@ -101,6 +102,34 @@ def test_intraop_surface_voxelisation():
 
     assert np.count_nonzero(cells_on_surface) == 2059
 
+def test_preop_from_vtkdata():
+    """ Voxelise from already loaded vtkdata, rather than from file name."""
+    input_file = 'tests/data/voxelisation/liver_downsample.stl'
+    model = vtk_surface_model.VTKSurfaceModel(input_file, [0.0, 1.0, 0.0])
+
+    size = 0.3
+    grid_elements = 64
+
+    grid = voxelise.voxelise(input_mesh=model.get_vtk_data(),
+                            scale_input=0.001,
+                            center=True,
+                            size=size,
+                            grid_elements=grid_elements
+                            )
+
+    # Check dimensions correct
+    cell_dims = [0, 0, 0]
+    grid.GetCellDims(cell_dims)
+    assert cell_dims == [63, 63, 63]
+
+    # Check array name is correct
+    numpy_data = voxelise.extract_array_from_grid(grid, 'preoperativeSurface')
+    print("Numpy data" ,numpy_data)
+    
+    # Cells 'inside' the liver have negative values, so this should be
+    # consistent
+    cells_in_liver = numpy_data < 0
+    assert np.count_nonzero(cells_in_liver) == 14628
 
 def test_intraop_from_numpy():
     """ test_liver_stl_voxelisation needs to have run successfully for this