From 5711b128dbebe87b465ad67bedbea3206d2353c0 Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Wed, 17 Jul 2024 00:20:28 -0400
Subject: [PATCH] feat: add vtk mesh display to zmesh meshes

---
 setup.py      |  3 +++
 zmesh/mesh.py | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/setup.py b/setup.py
index 42d3b24..4bcae17 100755
--- a/setup.py
+++ b/setup.py
@@ -34,6 +34,9 @@ def __repr__(self):
   python_requires=">=3.8",
   pbr=True,
   define_macros=[ ("NDEBUG", 1) ],
+  extras_require={
+    "viewer": [ "vtk" ],
+  },
   ext_modules=[
     setuptools.Extension(
       'zmesh._zmesh',
diff --git a/zmesh/mesh.py b/zmesh/mesh.py
index da5946c..7ce7a8b 100644
--- a/zmesh/mesh.py
+++ b/zmesh/mesh.py
@@ -210,3 +210,51 @@ def to_ply(self):
     )
 
     return plydata
+
+  def viewer(self):
+    # thanks to ChatGPT for making it easy to figure out
+    # how to display VTK meshes.
+    import vtk
+    polydata = _create_vtk_mesh(self.vertices, self.faces)
+
+    mapper = vtk.vtkPolyDataMapper()
+    mapper.SetInputData(polydata)
+
+    actor = vtk.vtkActor()
+    actor.SetMapper(mapper)
+
+    renderer = vtk.vtkRenderer()
+    render_window = vtk.vtkRenderWindow()
+    render_window.AddRenderer(renderer)
+    render_window_interactor = vtk.vtkRenderWindowInteractor()
+    render_window_interactor.SetRenderWindow(render_window)
+
+    render_window.SetSize(1024, 1024)
+
+    renderer.AddActor(actor)
+    renderer.SetBackground(0.1, 0.2, 0.3)  # Background color
+
+    render_window.Render()
+    render_window_interactor.Start()
+
+def _create_vtk_mesh(vertices, faces):
+  import vtk
+  from vtk.util.numpy_support import numpy_to_vtk, numpy_to_vtkIdTypeArray
+
+  vtk_points = vtk.vtkPoints()
+  vtk_points.SetData(numpy_to_vtk(vertices))
+  
+  polydata = vtk.vtkPolyData()
+  polydata.SetPoints(vtk_points)
+  
+  vtk_faces = vtk.vtkCellArray()
+  vtk_faces.SetCells(
+    faces.shape[0], 
+    numpy_to_vtkIdTypeArray(np.hstack([np.full((faces.shape[0], 1), 3), faces]).flatten())
+  )
+
+  polydata.SetPolys(vtk_faces)
+  
+  return polydata
+
+