fluencyCAD/drawing_modules/vtk_widget.py
2024-07-04 22:58:52 +02:00

448 lines
16 KiB
Python

import sys
import numpy as np
import vtk
from PySide6 import QtCore, QtWidgets
from PySide6.QtCore import Signal
from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
class VTKWidget(QtWidgets.QWidget):
face_data = Signal(dict)
def __init__(self, parent=None):
super().__init__(parent)
self.vtk_widget = QVTKRenderWindowInteractor(self)
# Create layout and add VTK widget
layout = QtWidgets.QVBoxLayout()
layout.addWidget(self.vtk_widget)
self.setLayout(layout)
# Create VTK pipeline
self.renderer = vtk.vtkRenderer()
self.vtk_widget.GetRenderWindow().AddRenderer(self.renderer)
self.interactor = self.vtk_widget.GetRenderWindow().GetInteractor()
# Set up the camera
self.camera = self.renderer.GetActiveCamera()
self.camera.SetPosition(5, 5, 5)
self.camera.SetFocalPoint(0, 0, 0)
# Set up picking
self.picker = vtk.vtkCellPicker()
self.picker.SetTolerance(0.0005)
# Create a mapper and actor for picked cells
self.picked_mapper = vtk.vtkDataSetMapper()
self.picked_actor = vtk.vtkActor()
self.picked_actor.SetMapper(self.picked_mapper)
self.picked_actor.GetProperty().SetColor(1.0, 0.0, 0.0) # Red color for picked faces
self.renderer.AddActor(self.picked_actor)
# Set up interactor style
self.style = vtk.vtkInteractorStyleTrackballCamera()
self.interactor.SetInteractorStyle(self.style)
# Add observer for mouse clicks
self.interactor.AddObserver("RightButtonPressEvent", self.on_click)
def create_cube_mesh(self):
cube_source = vtk.vtkCubeSource()
print(cube_source)
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputConnection(cube_source.GetOutputPort())
actor = vtk.vtkActor()
actor.SetMapper(mapper)
self.renderer.AddActor(actor)
def simplify_mesh(self, input_mesh, target_reduction):
# Create the quadric decimation filter
decimate = vtk.vtkDecimatePro()
decimate.SetInputData(input_mesh)
# Set the reduction factor (0 to 1, where 1 means maximum reduction)
decimate.SetTargetReduction(target_reduction)
# Optional: Preserve topology (if needed)
decimate.PreserveTopologyOn()
# Perform the decimation
decimate.Update()
return decimate.GetOutput()
def combine_coplanar_faces(self, input_polydata, tolerance=0.001):
# Clean the polydata to merge duplicate points
clean = vtk.vtkCleanPolyData()
clean.SetInputData(input_polydata)
clean.SetTolerance(tolerance)
clean.Update()
# Generate normals and merge coplanar polygons
normals = vtk.vtkPolyDataNormals()
normals.SetInputConnection(clean.GetOutputPort())
normals.SplittingOff() # Disable splitting of sharp edges
normals.ConsistencyOn() # Ensure consistent polygon ordering
normals.AutoOrientNormalsOn() # Automatically orient normals
normals.ComputePointNormalsOff() # We only need face normals
normals.ComputeCellNormalsOn() # Compute cell normals
normals.Update()
return normals.GetOutput()
def poisson_reconstruction(self, points):
# Create a polydata object from points
point_polydata = vtk.vtkPolyData()
point_polydata.SetPoints(points)
# Create a surface reconstruction filter
surf = vtk.vtkSurfaceReconstructionFilter()
surf.SetInputData(point_polydata)
surf.Update()
# Create a contour filter to extract the surface
cf = vtk.vtkContourFilter()
cf.SetInputConnection(surf.GetOutputPort())
cf.SetValue(0, 0.0)
cf.Update()
# Reverse normals
reverse = vtk.vtkReverseSense()
reverse.SetInputConnection(cf.GetOutputPort())
reverse.ReverseCellsOn()
reverse.ReverseNormalsOn()
reverse.Update()
return reverse.GetOutput()
def load_interactor_mesh(self, simp_mesh):
vertices, faces = simp_mesh
self.load_custom_mesh(vertices, faces)
def load_custom_mesh(self, vertices, faces):
### Load meshes by own module
# Create a vtkPoints object and store the points in it
points = vtk.vtkPoints()
for vertex in vertices:
points.InsertNextPoint(vertex)
# Create a vtkCellArray to store the faces
cells = vtk.vtkCellArray()
for face in faces:
triangle = vtk.vtkTriangle()
triangle.GetPointIds().SetId(0, face[0])
triangle.GetPointIds().SetId(1, face[1])
triangle.GetPointIds().SetId(2, face[2])
cells.InsertNextCell(triangle)
# Create a polydata object
polydata = vtk.vtkPolyData()
polydata.SetPoints(points)
polydata.SetPolys(cells)
# Create mapper and actor
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputData(polydata) # Make sure this line is present
actor = vtk.vtkActor()
actor.SetMapper(mapper)
# Add to renderer
self.renderer.AddActor(actor)
# Force an update of the pipeline
mapper.Update()
self.vtk_widget.GetRenderWindow().Render()
def create_simplified_outline(self, polydata):
# 1. Extract the outer surface
surface_filter = vtk.vtkDataSetSurfaceFilter()
surface_filter.SetInputData(polydata)
surface_filter.Update()
# 2. Extract feature edges (only boundary edges)
feature_edges = vtk.vtkFeatureEdges()
feature_edges.SetInputConnection(surface_filter.GetOutputPort())
feature_edges.BoundaryEdgesOn()
feature_edges.FeatureEdgesOff()
feature_edges.NonManifoldEdgesOff()
feature_edges.ManifoldEdgesOff()
feature_edges.Update()
# 3. Clean the edges to merge duplicate points
cleaner = vtk.vtkCleanPolyData()
cleaner.SetInputConnection(feature_edges.GetOutputPort())
cleaner.Update()
# 4. Optional: Smooth the outline
smooth = vtk.vtkSmoothPolyDataFilter()
smooth.SetInputConnection(cleaner.GetOutputPort())
smooth.SetNumberOfIterations(15)
smooth.SetRelaxationFactor(0.1)
smooth.FeatureEdgeSmoothingOff()
smooth.BoundarySmoothingOn()
smooth.Update()
return smooth
def render_from_points_direct_with_faces(self, points):
# Create a vtkPoints object and store the points in it
vtk_points = vtk.vtkPoints()
for point in points:
vtk_points.InsertNextPoint(point)
# Create a vtkCellArray to store the triangles
triangles = vtk.vtkCellArray()
# Assuming points are organized as triplets forming triangles
for i in range(0, len(points), 3):
triangle = vtk.vtkTriangle()
triangle.GetPointIds().SetId(0, i)
triangle.GetPointIds().SetId(1, i + 1)
triangle.GetPointIds().SetId(2, i + 2)
triangles.InsertNextCell(triangle)
# Create a polydata object
polydata = vtk.vtkPolyData()
polydata.SetPoints(vtk_points)
polydata.SetPolys(triangles)
# Optional: Merge duplicate points
cleaner = vtk.vtkCleanPolyData()
cleaner.SetInputData(polydata)
cleaner.Update()
# Optional: Combine coplanar faces
normals = vtk.vtkPolyDataNormals()
normals.SetInputConnection(cleaner.GetOutputPort())
normals.SplittingOff()
normals.ConsistencyOn()
normals.AutoOrientNormalsOn()
normals.ComputePointNormalsOff()
normals.ComputeCellNormalsOn()
normals.Update()
# Create a mapper and actor
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputConnection(normals.GetOutputPort())
actor = vtk.vtkActor()
actor.SetMapper(mapper)
actor.GetProperty().SetColor(1, 1, 1) # Set color (white in this case)
actor.GetProperty().EdgeVisibilityOn() # Show edges
actor.GetProperty().SetLineWidth(2) # Set line width
feature_edges = self.create_simplified_outline(polydata)
# Create a mapper for the feature edges
edge_mapper = vtk.vtkPolyDataMapper()
# Already wiht output
edge_mapper.SetInputConnection(feature_edges.GetOutputPort())
# Create an actor for the feature edges
edge_actor = vtk.vtkActor()
edge_actor.SetMapper(edge_mapper)
# Set the properties of the edge actor
edge_actor.GetProperty().SetColor(1, 0, 0) # Set color (red in this case)
edge_actor.GetProperty().SetLineWidth(2) # Set line width
# Optionally, if you want to keep the original mesh visible:
# (assuming you have the original mesh mapper and actor set up)
self.renderer.AddActor(actor) # Add the original mesh actor
# Add the edge actor to the renderer
self.renderer.AddActor(edge_actor)
# Force an update of the pipeline
mapper.Update()
self.vtk_widget.GetRenderWindow().Render()
# Print statistics
print(f"Original points: {len(points)}")
print(f"Number of triangles: {triangles.GetNumberOfCells()}")
print(f"Final number of points: {normals.GetOutput().GetNumberOfPoints()}")
print(f"Final number of cells: {normals.GetOutput().GetNumberOfCells()}")
def render_from_points_direct(self, points):
### Rendermethod for SDF mesh (output)
# Create a vtkPoints object and store the points in it
vtk_points = vtk.vtkPoints()
for point in points:
vtk_points.InsertNextPoint(point)
# Create a polydata object
point_polydata = vtk.vtkPolyData()
point_polydata.SetPoints(vtk_points)
# Surface reconstruction
surf = vtk.vtkSurfaceReconstructionFilter()
surf.SetInputData(point_polydata)
surf.Update()
# Create a contour filter to extract the surface
cf = vtk.vtkContourFilter()
cf.SetInputConnection(surf.GetOutputPort())
cf.SetValue(0, 0.0)
cf.Update()
# Reverse the normals
reverse = vtk.vtkReverseSense()
reverse.SetInputConnection(cf.GetOutputPort())
reverse.ReverseCellsOn()
reverse.ReverseNormalsOn()
reverse.Update()
# Get the reconstructed mesh
reconstructed_mesh = reverse.GetOutput()
"""# Simplify the mesh
target_reduction = 1 # Adjust this value as needed
simplified_mesh = self.simplify_mesh(reconstructed_mesh, target_reduction)
combinded_faces = self.combine_coplanar_faces(simplified_mesh, 0.001)"""
# Create a mapper and actor for the simplified mesh
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputData(reconstructed_mesh)
actor = vtk.vtkActor()
actor.SetMapper(mapper)
actor.GetProperty().SetColor(1, 1, 1) # Set color (white in this case)
actor.GetProperty().EdgeVisibilityOn() # Show edges
actor.GetProperty().SetLineWidth(2) # Set line width
# Add the actor to the renderer
self.renderer.AddActor(actor)
# Force an update of the pipeline
#mapper.Update()
self.vtk_widget.GetRenderWindow().Render()
# Print statistics
print(f"Original points: {len(points)}")
print(
f"Reconstructed mesh: {reconstructed_mesh.GetNumberOfPoints()} points, {reconstructed_mesh.GetNumberOfCells()} cells")
"""print(
f"Simplified mesh: {simplified_mesh.GetNumberOfPoints()} points, {simplified_mesh.GetNumberOfCells()} cells")"""
def on_click(self, obj, event):
click_pos = self.interactor.GetEventPosition()
# Perform pick
self.picker.Pick(click_pos[0], click_pos[1], 0, self.renderer)
# Get picked cell
cell_id = self.picker.GetCellId()
if cell_id != -1:
print(f"Picked face ID: {cell_id}")
# Get the polydata and the picked cell
polydata = self.picker.GetActor().GetMapper().GetInput()
cell = polydata.GetCell(cell_id)
# Project2D
renderer = self.vtk_widget.GetRenderWindow().GetRenderers().GetFirstRenderer()
camera = renderer.GetActiveCamera()
# Get cell type
cell_type = cell.GetCellType()
print(f"Cell type: {cell_type}")
# Get points of the cell
points = cell.GetPoints()
num_points = points.GetNumberOfPoints()
print(f"Number of points in the cell: {num_points}")
vec_points = []
# Get coordinates of each point
for i in range(num_points):
point = points.GetPoint(i)
print(f"Point {i}: {point}")
vec_points.append(point)
# Get normal of the cell (if it's a polygon)
if cell_type == vtk.VTK_TRIANGLE:
normal = [0, 0, 0]
vtk.vtkPolygon.ComputeNormal(points, normal)
print(f"Face normal: {normal}")
# Get cell data
cell_data = polydata.GetCellData()
if cell_data:
num_arrays = cell_data.GetNumberOfArrays()
print(f"Number of cell data arrays: {num_arrays}")
for i in range(num_arrays):
array = cell_data.GetArray(i)
array_name = array.GetName()
num_components = array.GetNumberOfComponents()
value = [0] * num_components
array.GetTuple(cell_id, value)
print(f"Cell data '{array_name}': {value}")
# Get point data (average of all points in the cell)
point_data = polydata.GetPointData()
if point_data:
num_arrays = point_data.GetNumberOfArrays()
print(f"Number of point data arrays: {num_arrays}")
for i in range(num_arrays):
array = point_data.GetArray(i)
array_name = array.GetName()
num_components = array.GetNumberOfComponents()
avg_value = np.zeros(num_components)
for j in range(num_points):
point_id = cell.GetPointId(j)
value = [0] * num_components
array.GetTuple(point_id, value)
avg_value += np.array(value)
avg_value /= num_points
print(f"Average point data '{array_name}': {avg_value}")
if num_points and cell_data:
face_orient = {'cell_data': cell_data, 'points': vec_points }
print(face_orient)
self.face_data.emit(face_orient)
# Highlight picked face (your existing code)
ids = vtk.vtkIdTypeArray()
ids.SetNumberOfComponents(1)
ids.InsertNextValue(cell_id)
selection_node = vtk.vtkSelectionNode()
selection_node.SetFieldType(vtk.vtkSelectionNode.CELL)
selection_node.SetContentType(vtk.vtkSelectionNode.INDICES)
selection_node.SetSelectionList(ids)
selection = vtk.vtkSelection()
selection.AddNode(selection_node)
extract_selection = vtk.vtkExtractSelection()
extract_selection.SetInputData(0, polydata)
extract_selection.SetInputData(1, selection)
extract_selection.Update()
self.picked_mapper.SetInputData(extract_selection.GetOutput())
self.vtk_widget.GetRenderWindow().Render()
def start(self):
self.interactor.Initialize()
self.interactor.Start()
class MainWindow(QtWidgets.QMainWindow):
def __init__(self, parent=None):
super().__init__(parent)
self.vtk_widget = VTKWidget()
self.setCentralWidget(self.vtk_widget)
self.setWindowTitle("VTK Mesh Viewer")
self.vtk_widget.create_cube_mesh()
self.show()
self.vtk_widget.start()
if __name__ == "__main__":
app = QtWidgets.QApplication(sys.argv)
window = MainWindow()
sys.exit(app.exec())