382 lines
14 KiB
Python
382 lines
14 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)
|
|
|
|
@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()
|
|
|
|
self.cell_normals = vtk_to_numpy(normalGenerator.GetOutput().GetCellData().GetNormals())
|
|
|
|
# Create a mapper and actor
|
|
mapper = vtk.vtkPolyDataMapper()
|
|
mapper.SetInputData(polydata)
|
|
|
|
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 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)
|
|
|
|
# Compute the centroid of the selected edges
|
|
centroid = (
|
|
0, 0, 0) # point1 #np.mean([point for edge in self.selected_edges for point in edge], axis=0)
|
|
|
|
# Create a transform for projection
|
|
transform = vtk.vtkTransform()
|
|
transform.Identity()
|
|
|
|
# Translate to center the transform on the centroid
|
|
transform.Translate(-centroid[0], -centroid[1], -centroid[2])
|
|
|
|
# Compute rotation to align normal with Z-axis
|
|
z_axis = np.array([0, 0, 1])
|
|
rotation_axis = np.cross(self.selected_normal, z_axis)
|
|
rotation_angle = np.arccos(np.dot(self.selected_normal, z_axis)) * 180 / np.pi
|
|
|
|
if np.linalg.norm(rotation_axis) > 1e-6: # Avoid division by zero
|
|
transform.RotateWXYZ(rotation_angle, rotation_axis[0], rotation_axis[1], rotation_axis[2])
|
|
|
|
# Apply scaling to flatten
|
|
transform.Scale(1, 1, 0)
|
|
|
|
# Apply the inverse rotation to bring it back to original orientation
|
|
transform.RotateWXYZ(-rotation_angle, rotation_axis[0], rotation_axis[1], rotation_axis[2])
|
|
|
|
# Translate back
|
|
transform.Translate(centroid[0], centroid[1], centroid[2])
|
|
|
|
# Apply the transform to the polydata
|
|
transformFilter = vtk.vtkTransformPolyDataFilter()
|
|
transformFilter.SetInputData(polydata)
|
|
transformFilter.SetTransform(transform)
|
|
transformFilter.Update()
|
|
|
|
# Get the projected polydata
|
|
projected_polydata = transformFilter.GetOutput()
|
|
print("Polydata", projected_polydata)
|
|
|
|
# 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)
|
|
|
|
# Add a plane to visualize the projection plane
|
|
plane_source = vtk.vtkPlaneSource()
|
|
plane_source.SetCenter(centroid)
|
|
plane_source.SetNormal(self.selected_normal)
|
|
plane_mapper = vtk.vtkPolyDataMapper()
|
|
plane_mapper.SetInputConnection(plane_source.GetOutputPort())
|
|
plane_actor = vtk.vtkActor()
|
|
plane_actor.SetMapper(plane_mapper)
|
|
plane_actor.GetProperty().SetColor(0.8, 0.8, 0.8) # Light gray
|
|
plane_actor.GetProperty().SetOpacity(0.5)
|
|
self.renderer.AddActor(plane_actor)
|
|
|
|
# Reset camera to show all actors
|
|
self.renderer.ResetCamera()
|
|
|
|
# Render and interact
|
|
self.vtk_widget.GetRenderWindow().Render()
|
|
|
|
self.project_tosketch_edge = self.get_points_and_edges_from_polydata(projected_polydata)
|
|
print("Edges", self.project_tosketch_edge[0])
|
|
|
|
|
|
elif len(self.access_selected_points) > 2:
|
|
self.access_selected_points = []
|
|
# Clear the picked edge actors
|
|
for actor in self.picked_edge_actors:
|
|
self.renderer.RemoveActor(actor)
|
|
self.picked_edge_actors.clear()
|
|
# Reset selected edges
|
|
self.selected_edges = []
|
|
else:
|
|
print("Selected cell is not a line")
|
|
|
|
# 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())
|