418 lines
15 KiB
Python
418 lines
15 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
|
|
from vtkmodules.util.numpy_support import vtk_to_numpy, numpy_to_vtk
|
|
|
|
|
|
class VTKWidget(QtWidgets.QWidget):
|
|
face_data = Signal(dict)
|
|
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent)
|
|
self.access_selected_points = []
|
|
self.selected_normal = None
|
|
self.selected_edges = []
|
|
self.cell_normals = None
|
|
|
|
self.project_tosketch_edge = []
|
|
|
|
self.vtk_widget = QVTKRenderWindowInteractor(self)
|
|
|
|
self.picked_edge_actors = []
|
|
|
|
# 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, 100)
|
|
self.camera.SetFocalPoint(0, 0, 0)
|
|
|
|
# Set up picking
|
|
self.picker = vtk.vtkCellPicker()
|
|
self.picker.SetTolerance(0.005)
|
|
|
|
# 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.picked_actor.VisibilityOff() # Initially hide the actor
|
|
self.renderer.AddActor(self.picked_actor)
|
|
|
|
# Create an extract selection filter
|
|
self.extract_selection = vtk.vtkExtractSelection()
|
|
|
|
# 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)
|
|
|
|
# Add axis gizmo (smaller size)
|
|
self.axes = vtk.vtkAxesActor()
|
|
self.axes.SetTotalLength(0.5, 0.5, 0.5) # Reduced size
|
|
self.axes.SetShaftType(0)
|
|
self.axes.SetAxisLabels(1)
|
|
|
|
# Create an orientation marker
|
|
self.axes_widget = vtk.vtkOrientationMarkerWidget()
|
|
self.axes_widget.SetOrientationMarker(self.axes)
|
|
self.axes_widget.SetInteractor(self.interactor)
|
|
self.axes_widget.SetViewport(0.0, 0.0, 0.2, 0.2) # Set position and size
|
|
self.axes_widget.EnabledOn()
|
|
self.axes_widget.InteractiveOff()
|
|
|
|
# Start the interactor
|
|
self.interactor.Initialize()
|
|
self.interactor.Start()
|
|
|
|
@staticmethod
|
|
def compute_normal_from_lines(line1, line2):
|
|
vec1 = line1[1] - line1[0]
|
|
vec2 = line2[1] - line2[0]
|
|
normal = np.cross(vec1, vec2)
|
|
print(normal)
|
|
normal = normal / np.linalg.norm(normal)
|
|
return normal
|
|
|
|
def load_interactor_mesh(self, edges):
|
|
# Create vtkPoints to store all points
|
|
points = vtk.vtkPoints()
|
|
|
|
# Create vtkCellArray to store the lines
|
|
lines = vtk.vtkCellArray()
|
|
|
|
for edge in edges:
|
|
# Add points for this edge
|
|
point_id1 = points.InsertNextPoint(edge[0])
|
|
point_id2 = points.InsertNextPoint(edge[1])
|
|
|
|
# Create a line using the point IDs
|
|
line = vtk.vtkLine()
|
|
line.GetPointIds().SetId(0, point_id1)
|
|
line.GetPointIds().SetId(1, point_id2)
|
|
|
|
# Add the line to the cell array
|
|
lines.InsertNextCell(line)
|
|
|
|
# Create vtkPolyData to store the geometry
|
|
polydata = vtk.vtkPolyData()
|
|
polydata.SetPoints(points)
|
|
polydata.SetLines(lines)
|
|
|
|
# Create a mapper and actor
|
|
mapper = vtk.vtkPolyDataMapper()
|
|
mapper.SetInputData(polydata)
|
|
|
|
|
|
actor = vtk.vtkActor()
|
|
actor.SetMapper(mapper)
|
|
actor.GetProperty().SetColor(0.0, 0.0, 1.0) # Set color to red
|
|
actor.GetProperty().SetLineWidth(2) # Set line width
|
|
|
|
# Add the actor to the scene
|
|
self.renderer.AddActor(actor)
|
|
#self.renderer.SetBackground(0.1, 0.2, 0.4) # Set background color
|
|
|
|
mapper.Update()
|
|
self.vtk_widget.GetRenderWindow().Render()
|
|
|
|
def render_from_points_direct_with_faces(self, vertices, faces):
|
|
points = vtk.vtkPoints()
|
|
for i in range(vertices.shape[0]):
|
|
points.InsertNextPoint(vertices[i])
|
|
|
|
# Create a vtkCellArray to store the triangles
|
|
triangles = vtk.vtkCellArray()
|
|
for i in range(faces.shape[0]):
|
|
triangle = vtk.vtkTriangle()
|
|
triangle.GetPointIds().SetId(0, faces[i, 0])
|
|
triangle.GetPointIds().SetId(1, faces[i, 1])
|
|
triangle.GetPointIds().SetId(2, faces[i, 2])
|
|
triangles.InsertNextCell(triangle)
|
|
|
|
# Create a polydata object
|
|
polydata = vtk.vtkPolyData()
|
|
polydata.SetPoints(points)
|
|
polydata.SetPolys(triangles)
|
|
|
|
# Calculate normals
|
|
normalGenerator = vtk.vtkPolyDataNormals()
|
|
normalGenerator.SetInputData(polydata)
|
|
normalGenerator.ComputePointNormalsOn()
|
|
normalGenerator.ComputeCellNormalsOn()
|
|
normalGenerator.Update()
|
|
|
|
### There might be aproblem earlier but this fixes the drawing for now.
|
|
#TODO: Investigate upstream conversion errors.
|
|
# Create a transform for mirroring across the x-axis
|
|
mirror_transform = vtk.vtkTransform()
|
|
mirror_transform.Scale(-1, -1, 1) # This mirrors across the x-axis
|
|
|
|
# Apply the transform to the polydata
|
|
transformFilter = vtk.vtkTransformPolyDataFilter()
|
|
transformFilter.SetInputData(polydata)
|
|
transformFilter.SetTransform(mirror_transform)
|
|
transformFilter.Update()
|
|
|
|
self.cell_normals = vtk_to_numpy(normalGenerator.GetOutput().GetCellData().GetNormals())
|
|
|
|
# Create a mapper and actor
|
|
mapper = vtk.vtkPolyDataMapper()
|
|
mapper.SetInputData(transformFilter.GetOutput())
|
|
|
|
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
|
|
|
|
|
|
# (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
|
|
|
|
# Force an update of the pipeline
|
|
#mapper.Update()
|
|
self.vtk_widget.GetRenderWindow().Render()
|
|
|
|
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 get_points_and_edges_from_polydata(self, polydata) -> list:
|
|
# Extract points
|
|
points = {}
|
|
vtk_points = polydata.GetPoints()
|
|
for i in range(vtk_points.GetNumberOfPoints()):
|
|
point = vtk_points.GetPoint(i)
|
|
points[i] = np.array(point)
|
|
|
|
# Extract edges
|
|
edges = []
|
|
for i in range(polydata.GetNumberOfCells()):
|
|
cell = polydata.GetCell(i)
|
|
if cell.GetCellType() == vtk.VTK_LINE:
|
|
point_ids = cell.GetPointIds()
|
|
edge = (point_ids.GetId(0), point_ids.GetId(1))
|
|
edges.append(edge)
|
|
|
|
return points, edges
|
|
|
|
def project_mesh_to_plane(self, input_mesh, normal, origin):
|
|
# Create the projector
|
|
projector = vtk.vtkProjectPointsToPlane()
|
|
projector.SetInputData(input_mesh)
|
|
projector.SetProjectionTypeToSpecifiedPlane()
|
|
|
|
# Set the normal and origin of the plane
|
|
projector.SetNormal(normal)
|
|
projector.SetOrigin(origin)
|
|
|
|
# Execute the projection
|
|
projector.Update()
|
|
|
|
# Get the projected mesh
|
|
projected_mesh = projector.GetOutput()
|
|
return projected_mesh
|
|
|
|
def compute_2d_coordinates(self, projected_mesh, normal):
|
|
# Compute centroid of projected points
|
|
center_of_mass = vtk.vtkCenterOfMass()
|
|
center_of_mass.SetInputData(projected_mesh)
|
|
center_of_mass.SetUseScalarsAsWeights(False)
|
|
center_of_mass.Update()
|
|
centroid = center_of_mass.GetCenter()
|
|
|
|
# Create a coordinate system on the plane
|
|
z_axis = np.array(normal)
|
|
x_axis = np.cross(z_axis, [0, 0, 1])
|
|
if np.allclose(x_axis, 0):
|
|
x_axis = np.cross(z_axis, [0, 1, 0])
|
|
x_axis = x_axis / np.linalg.norm(x_axis)
|
|
y_axis = np.cross(z_axis, x_axis)
|
|
|
|
# Create transformation matrix
|
|
matrix = vtk.vtkMatrix4x4()
|
|
for i in range(3):
|
|
matrix.SetElement(i, 0, x_axis[i])
|
|
matrix.SetElement(i, 1, y_axis[i])
|
|
matrix.SetElement(i, 2, z_axis[i])
|
|
matrix.SetElement(i, 3, centroid[i])
|
|
matrix.Invert()
|
|
|
|
# Transform points to 2D coordinates
|
|
transform = vtk.vtkTransform()
|
|
transform.SetMatrix(matrix)
|
|
|
|
transformer = vtk.vtkTransformPolyDataFilter()
|
|
transformer.SetInputData(projected_mesh)
|
|
transformer.SetTransform(transform)
|
|
transformer.Update()
|
|
|
|
transformed_mesh = transformer.GetOutput()
|
|
points = transformed_mesh.GetPoints()
|
|
|
|
xy_coordinates = []
|
|
for i in range(points.GetNumberOfPoints()):
|
|
point = points.GetPoint(i)
|
|
xy_coordinates.append((point[0], point[1]))
|
|
|
|
return xy_coordinates
|
|
|
|
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 ID
|
|
cell_id = self.picker.GetCellId()
|
|
|
|
if cell_id != -1:
|
|
print(f"Picked cell ID: {cell_id}")
|
|
|
|
# Get the polydata and the picked cell
|
|
polydata = self.picker.GetActor().GetMapper().GetInput()
|
|
cell = polydata.GetCell(cell_id)
|
|
|
|
# Ensure it's a line
|
|
if cell.GetCellType() == vtk.VTK_LINE:
|
|
# Get the two points of the line
|
|
point_id1 = cell.GetPointId(0)
|
|
point_id2 = cell.GetPointId(1)
|
|
|
|
proj_point1 = polydata.GetPoint(point_id1)
|
|
proj_point2 = polydata.GetPoint(point_id2)
|
|
|
|
self.access_selected_points.append((proj_point1, proj_point2))
|
|
|
|
point1 = np.array(proj_point1)
|
|
point2 = np.array(proj_point2)
|
|
|
|
print(f"Line starts at: {point1}")
|
|
print(f"Line ends at: {point2}")
|
|
|
|
# Store this line for later use if needed
|
|
self.selected_edges.append((point1, point2))
|
|
|
|
# Create a new vtkLineSource for the picked edge
|
|
line_source = vtk.vtkLineSource()
|
|
line_source.SetPoint1(point1)
|
|
line_source.SetPoint2(point2)
|
|
|
|
# Create a mapper and actor for the picked edge
|
|
edge_mapper = vtk.vtkPolyDataMapper()
|
|
edge_mapper.SetInputConnection(line_source.GetOutputPort())
|
|
|
|
edge_actor = vtk.vtkActor()
|
|
edge_actor.SetMapper(edge_mapper)
|
|
edge_actor.GetProperty().SetColor(1.0, 0.0, 0.0) # Red color for picked edges
|
|
edge_actor.GetProperty().SetLineWidth(4) # Make the line thicker
|
|
|
|
# Add the actor to the renderer and store it
|
|
self.renderer.AddActor(edge_actor)
|
|
self.picked_edge_actors.append(edge_actor)
|
|
|
|
if len(self.selected_edges) == 2:
|
|
# Compute the normal from the two selected edges
|
|
edge1 = self.selected_edges[0][1] - self.selected_edges[0][0]
|
|
edge2 = self.selected_edges[1][1] - self.selected_edges[1][0]
|
|
self.selected_normal = np.cross(edge1, edge2)
|
|
self.selected_normal = self.selected_normal / np.linalg.norm(self.selected_normal)
|
|
print("Computed normal:", self.selected_normal)
|
|
|
|
centroid = np.mean([point for edge in self.selected_edges for point in edge], axis=0)
|
|
|
|
projected_polydata = self.project_mesh_to_plane(polydata, self.selected_normal, centroid)
|
|
projected_points = projected_polydata.GetPoints()
|
|
print("proj_points", projected_points)
|
|
|
|
# Extract 2D coordinates
|
|
self.project_tosketch_edge = self.compute_2d_coordinates(projected_polydata, self.selected_normal)
|
|
print("3d_points_proj", self.project_tosketch_edge)
|
|
|
|
# Create a mapper and actor for the projected data
|
|
mapper = vtk.vtkPolyDataMapper()
|
|
mapper.SetInputData(projected_polydata)
|
|
|
|
actor = vtk.vtkActor()
|
|
actor.SetMapper(mapper)
|
|
actor.GetProperty().SetColor(0.0, 1.0, 0.0) # Set color to green
|
|
actor.GetProperty().SetLineWidth(4) # Set line width
|
|
|
|
# Add the actor to the scene
|
|
self.renderer.AddActor(actor)
|
|
|
|
if len(self.selected_edges) > 2:
|
|
self.selected_edges = []
|
|
self.selected_normal = []
|
|
|
|
|
|
# Render the scene
|
|
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())
|