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.selected_vtk_line = [] self.access_selected_points = [] self.selected_normal = None self.centroid = None self.selected_edges = [] self.cell_normals = None self.local_matrix = None self.project_tosketch_points = [] self.project_tosketch_lines = [] self.vtk_widget = QVTKRenderWindowInteractor(self) self.picked_edge_actors = [] self.displayed_normal_actors = [] self.body_actors_orig = [] self.projected_mesh_actors = [] self.flip_toggle = False # 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.renderer_projections = vtk.vtkRenderer() self.renderer_indicators = vtk.vtkRenderer() self.renderer.SetViewport(0, 0, 1, 1) # Full viewport self.renderer_projections.SetViewport(0, 0, 1, 1) # Full viewport, overlays the first self.renderer_indicators.SetViewport(0, 0, 1, 1) # Full viewport, overlays the first self.renderer.SetLayer(0) self.renderer_projections.SetLayer(1) self.renderer_indicators.SetLayer(2) # This will be on top # Add renderers to the render window render_window = self.vtk_widget.GetRenderWindow() render_window.SetNumberOfLayers(3) render_window.AddRenderer(self.renderer) render_window.AddRenderer(self.renderer_projections) render_window.AddRenderer(self.renderer_indicators) # Set up shared camera self.camera = vtk.vtkCamera() self.camera.SetPosition(5, 5, 1000) self.camera.SetFocalPoint(0, 0, 0) self.camera.SetClippingRange(0.1, 100000) self.renderer.SetActiveCamera(self.camera) self.renderer_projections.SetActiveCamera(self.camera) self.renderer_indicators.SetActiveCamera(self.camera) self.interactor = self.vtk_widget.GetRenderWindow().GetInteractor() # Light Setup def add_light(renderer, position, color=(1, 1, 1), intensity=1.0): light = vtk.vtkLight() light.SetPosition(position) light.SetColor(color) light.SetIntensity(intensity) renderer.AddLight(light) # Add lights from multiple directions add_light(self.renderer, (1000, 0, 0), intensity=1.5) add_light(self.renderer, (-1000, 0, 0), intensity=1.5) add_light(self.renderer, (0, 1000, 0), intensity=1.5) add_light(self.renderer, (0, -1000, 0), intensity=1.5) add_light(self.renderer, (0, 0, 1000), intensity=1.5) add_light(self.renderer, (0, 0, -1000), intensity=1.5) # 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() # Create the grid grid = self.create_grid(size=100, spacing=10) # Setup actor and mapper mapper = vtk.vtkPolyDataMapper() mapper.SetInputData(grid) actor = vtk.vtkActor() actor.SetPickable(False) actor.SetMapper(mapper) actor.GetProperty().SetColor(0.5, 0.5, 0.5) # Set grid color to gray self.renderer.AddActor(actor) def update_render(self): self.renderer.ResetCameraClippingRange() self.renderer_projections.ResetCameraClippingRange() self.renderer_indicators.ResetCameraClippingRange() self.vtk_widget.GetRenderWindow().Render() def create_grid(self, size=100, spacing=10): # Create a vtkPoints object and store the points in it points = vtk.vtkPoints() # Create lines lines = vtk.vtkCellArray() # Create the grid for i in range(-size, size + 1, spacing): # X-direction line points.InsertNextPoint(i, -size, 0) points.InsertNextPoint(i, size, 0) line = vtk.vtkLine() line.GetPointIds().SetId(0, points.GetNumberOfPoints() - 2) line.GetPointIds().SetId(1, points.GetNumberOfPoints() - 1) lines.InsertNextCell(line) # Y-direction line points.InsertNextPoint(-size, i, 0) points.InsertNextPoint(size, i, 0) line = vtk.vtkLine() line.GetPointIds().SetId(0, points.GetNumberOfPoints() - 2) line.GetPointIds().SetId(1, points.GetNumberOfPoints() - 1) lines.InsertNextCell(line) # Create a polydata to store everything in grid = vtk.vtkPolyData() # Add the points to the dataset grid.SetPoints(points) # Add the lines to the dataset grid.SetLines(lines) return grid def on_receive_command(self, command): """Calls the individual commands pressed in main""" print("Receive command: ", command) if command == "flip": self.clear_actors_projection() self.flip_toggle = not self.flip_toggle # Toggle the flag self.on_invert_normal() @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, off_vector): # 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 transform for mirroring across the y-axis matrix_transform = vtk.vtkTransform() if self.local_matrix: print(self.local_matrix) matrix = vtk.vtkMatrix4x4() matrix.DeepCopy(self.local_matrix) matrix.Invert() matrix_transform.SetMatrix(matrix) #matrix_transform.Scale(1, 1, 1) # This mirrors across the y-axis # Apply the matrix transform transformFilter = vtk.vtkTransformPolyDataFilter() transformFilter.SetInputData(polydata) transformFilter.SetTransform(matrix_transform) transformFilter.Update() # Create and apply the offset transform offset_transform = vtk.vtkTransform() offset_transform.Translate(off_vector[0], off_vector[1], off_vector[2]) offsetFilter = vtk.vtkTransformPolyDataFilter() offsetFilter.SetInputConnection(transformFilter.GetOutputPort()) offsetFilter.SetTransform(offset_transform) offsetFilter.Update() # Create a mapper and actor mapper = vtk.vtkPolyDataMapper() mapper.SetInputConnection(offsetFilter.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(mapper) actor.GetProperty().SetColor(1.0, 1.0, 1.0) actor.GetProperty().SetLineWidth(4) # Set line width # Add the actor to the scene self.renderer.AddActor(actor) mapper.Update() self.vtk_widget.GetRenderWindow().Render() def render_from_points_direct_with_faces(self, vertices, faces, color=(0.1, 0.2, 0.8), line_width=2, point_size=5): """Sketch Widget has inverted Y axiis therefore we invert y via scale here until fix""" points = vtk.vtkPoints() # Use SetData with numpy array vtk_array = numpy_to_vtk(vertices, deep=True) points.SetData(vtk_array) # Create a vtkCellArray to store the triangles triangles = 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]) 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(color) actor.GetProperty().EdgeVisibilityOff() actor.GetProperty().SetLineWidth(line_width) actor.GetProperty().SetMetallic(1) actor.GetProperty().SetOpacity(0.8) actor.SetPickable(False) self.renderer.AddActor(actor) self.body_actors_orig.append(actor) self.vtk_widget.GetRenderWindow().Render() def clear_body_actors(self): for actor in self.body_actors_orig: self.renderer.RemoveActor(actor) def visualize_matrix(self, matrix): points = vtk.vtkPoints() for i in range(4): for j in range(4): points.InsertNextPoint(matrix.GetElement(0, j), matrix.GetElement(1, j), matrix.GetElement(2, j)) polydata = vtk.vtkPolyData() polydata.SetPoints(points) mapper = vtk.vtkPolyDataMapper() mapper.SetInputData(polydata) actor = vtk.vtkActor() actor.SetMapper(mapper) actor.GetProperty().SetPointSize(5) self.renderer.AddActor(actor) def numpy_to_vtk(self, array, deep=True): """Convert a numpy array to a vtk array.""" vtk_array = vtk.vtkDoubleArray() vtk_array.SetNumberOfComponents(array.shape[1]) vtk_array.SetNumberOfTuples(array.shape[0]) for i in range(array.shape[0]): for j in range(array.shape[1]): vtk_array.SetComponent(i, j, array[i, j]) return vtk_array 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): # Normalize the normal vector normal = np.array(normal) normal = normal / np.linalg.norm(normal) # Create a vtkTransform transform = vtk.vtkTransform() transform.PostMultiply() # This ensures transforms are applied in the order we specify # Rotate so that the normal aligns with the Z-axis rotation_axis = np.cross(normal, [0, 0, 1]) angle = np.arccos(np.dot(normal, [0, 0, 1])) * 180 / np.pi # Convert to degrees if np.linalg.norm(rotation_axis) > 1e-6: # Check if rotation is needed transform.RotateWXYZ(angle, rotation_axis[0], rotation_axis[1], rotation_axis[2]) # Get the transformation matrix matrix = transform.GetMatrix() self.local_matrix = [matrix.GetElement(i, j) for i in range(4) for j in range(4)] # Apply the transform to the polydata transformFilter = vtk.vtkTransformPolyDataFilter() transformFilter.SetInputData(projected_mesh) transformFilter.SetTransform(transform) transformFilter.Update() # Get the transformed points transformed_polydata = transformFilter.GetOutput() points = transformed_polydata.GetPoints() # Extract 2D coordinates xy_coordinates = [] for i in range(points.GetNumberOfPoints()): point = points.GetPoint(i) xy_coordinates.append((point[0], point[1])) return xy_coordinates def compute_2d_coordinates_line(self, line_source, normal): # Ensure the input is a vtkLineSource if not isinstance(line_source, vtk.vtkLineSource): raise ValueError("Input must be a vtkLineSource") # Normalize the normal vector normal = np.array(normal) normal = normal / np.linalg.norm(normal) # Create a vtkTransform transform = vtk.vtkTransform() transform.PostMultiply() # This ensures transforms are applied in the order we specify # Rotate so that the normal aligns with the Z-axis rotation_axis = np.cross(normal, [0, 0, 1]) angle = np.arccos(np.dot(normal, [0, 0, 1])) * 180 / np.pi # Convert to degrees if np.linalg.norm(rotation_axis) > 1e-6: # Check if rotation is needed transform.RotateWXYZ(angle, rotation_axis[0], rotation_axis[1], rotation_axis[2]) # Get the transformation matrix matrix = transform.GetMatrix() local_matrix = [matrix.GetElement(i, j) for i in range(4) for j in range(4)] # Get the polydata from the line source line_source.Update() polydata = line_source.GetOutput() # Apply the transform to the polydata transform_filter = vtk.vtkTransformPolyDataFilter() transform_filter.SetInputData(polydata) transform_filter.SetTransform(transform) transform_filter.Update() # Get the transformed points transformed_polydata = transform_filter.GetOutput() transformed_points = transformed_polydata.GetPoints() # Extract 2D coordinates xy_coordinates = [] for i in range(transformed_points.GetNumberOfPoints()): point = transformed_points.GetPoint(i) xy_coordinates.append((point[0], point[1])) return xy_coordinates def project_2d_to_3d(self, xy_coordinates, normal): # Normalize the normal vector normal = np.array(normal) normal = normal / np.linalg.norm(normal) # Create a vtkTransform for the reverse transformation reverse_transform = vtk.vtkTransform() reverse_transform.PostMultiply() # This ensures transforms are applied in the order we specify # Compute the rotation axis and angle (same as in compute_2d_coordinates) rotation_axis = np.cross(normal, [0, 0, 1]) angle = np.arccos(np.dot(normal, [0, 0, 1])) * 180 / np.pi # Convert to degrees if np.linalg.norm(rotation_axis) > 1e-6: # Check if rotation is needed # Apply the inverse rotation reverse_transform.RotateWXYZ(-angle, rotation_axis[0], rotation_axis[1], rotation_axis[2]) # Create vtkPoints to store the 2D points points_2d = vtk.vtkPoints() for x, y in xy_coordinates: points_2d.InsertNextPoint(x, y, 0) # Z-coordinate is 0 for 2D points # Create a polydata with the 2D points polydata_2d = vtk.vtkPolyData() polydata_2d.SetPoints(points_2d) # Apply the reverse transform to the polydata transform_filter = vtk.vtkTransformPolyDataFilter() transform_filter.SetInputData(polydata_2d) transform_filter.SetTransform(reverse_transform) transform_filter.Update() # Get the transformed points (now in 3D) transformed_polydata = transform_filter.GetOutput() transformed_points = transformed_polydata.GetPoints() # Extract 3D coordinates xyz_coordinates = [] for i in range(transformed_points.GetNumberOfPoints()): point = transformed_points.GetPoint(i) xyz_coordinates.append((point[0], point[1], point[2])) return xyz_coordinates def add_normal_line(self, origin, normal, length=10.0, color=(1, 0, 0)): # Normalize the normal vector normal = np.array(normal) normal = normal / np.linalg.norm(normal) # Calculate the end point end_point = origin + normal * length # Create vtkPoints points = vtk.vtkPoints() points.InsertNextPoint(origin) points.InsertNextPoint(end_point) # Create a line line = vtk.vtkLine() line.GetPointIds().SetId(0, 0) line.GetPointIds().SetId(1, 1) # Create a cell array to store the line lines = vtk.vtkCellArray() lines.InsertNextCell(line) # Create a polydata to store everything in polyData = vtk.vtkPolyData() polyData.SetPoints(points) polyData.SetLines(lines) # Create mapper and actor mapper = vtk.vtkPolyDataMapper() mapper.SetInputData(polyData) actor = vtk.vtkActor() actor.SetMapper(mapper) actor.GetProperty().SetColor(color) actor.GetProperty().SetLineWidth(2) # Adjust line width as needed # Add to renderer self.renderer.AddActor(actor) self.vtk_widget.GetRenderWindow().Render() return actor # Return the actor in case you need to remove or modify it later def on_invert_normal(self): # Kippstufe für Normal flip if self.selected_normal is not None: self.clear_actors_normals() self.compute_projection(self.flip_toggle) 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) self.selected_vtk_line.append(line_source) # 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(5) # Make the line thicker # Add the actor to the renderer and store it self.renderer_indicators.AddActor(edge_actor) self.picked_edge_actors.append(edge_actor) if len(self.selected_edges) == 2: self.compute_projection(False) if len(self.selected_edges) > 2: # Clear lists for selection self.selected_vtk_line.clear() self.selected_edges.clear() self.clear_edge_select() # Clear Actors from view self.clear_actors_projection() self.clear_actors_sel_edges() self.clear_actors_normals() def find_origin_vertex(self, edge1, edge2): if edge1[0] == edge2[0]or edge1[0] == edge2[1]: return edge1[0] elif edge1[1] == edge2[0] or edge1[1] == edge2[1]: return edge1[1] else: return None # The edges don't share a vertex def clear_edge_select(self ): # Clear selection after projection was succesful self.selected_edges = [] self.selected_normal = [] def clear_actors_projection(self): """Removes all actors that were used for projection""" for flat_mesh in self.projected_mesh_actors: self.renderer_projections.RemoveActor(flat_mesh) def clear_actors_normals(self): for normals in self.displayed_normal_actors: self.renderer_indicators.RemoveActor(normals) def clear_actors_sel_edges(self): for edge_line in self.picked_edge_actors: self.renderer_indicators.RemoveActor(edge_line) def compute_projection(self, direction_invert: bool = False): # 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] selected_normal = np.cross(edge1, edge2) selected_normal = selected_normal / np.linalg.norm(selected_normal) #print("Computed normal:", self.selected_normal) # Invert the normal in local z if direction_invert is True if direction_invert: self.selected_normal = -selected_normal else: self.selected_normal = selected_normal self.centroid = np.mean([point for edge in self.selected_edges for point in edge], axis=0) #self.centroid = self.find_origin_vertex(edge1, edge2) # Draw the normal line normal_length = 50 # Adjust this value to change the length of the normal line normal_actor = self.add_normal_line(self.centroid, self.selected_normal, length=normal_length, color=(1, 0, 0)) polydata = self.picker.GetActor().GetMapper().GetInput() projected_polydata = self.project_mesh_to_plane(polydata, self.selected_normal, self.centroid) # Extract 2D coordinates self.project_tosketch_points = self.compute_2d_coordinates(projected_polydata, self.selected_normal) # Seperately rotate selected edges for drawing self.project_tosketch_lines.clear() for vtk_line in self.selected_vtk_line: proj_vtk_line = self.compute_2d_coordinates_line(vtk_line, self.selected_normal) self.project_tosketch_lines.append(proj_vtk_line) print("outgoing lines", self.project_tosketch_lines) # Create a mapper and actor for the projected data mapper = vtk.vtkPolyDataMapper() mapper.SetInputData(projected_polydata) # Projected mesh in green actor = vtk.vtkActor() actor.SetMapper(mapper) #actor.GetProperty().SetRenderLinesAsTubes(True) actor.GetProperty().SetColor(0.0, 1.0, 0.0) # Set color to green actor.GetProperty().SetLineWidth(4) # Set line width self.renderer_indicators.AddActor(normal_actor) self.displayed_normal_actors.append(normal_actor) self.renderer_projections.AddActor(actor) self.projected_mesh_actors.append(actor) # Render the scene self.update_render() 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())