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())