class MainWindow(QtGui.QMainWindow): def __init__(self, filename, parent=None): QtGui.QMainWindow.__init__(self, parent) # Initiate the UI as defined by Qt Designer self.ui = Ui_MainWindow() self.ui.setupUi(self) self.poly_data = None self.node_ids = None self.element_ids = None self.node_count = 0 self.read_data(filename) self.ui.txt_msg.appendPlainText("Model Loaded.") # initialize colors self.eidcolor = (0, 0.5, 0.5) self.edgecolor = (0, 0, 0) self.bgcolor1 = (0, 0, 1) self.bgcolor2 = (0.8, 0.8, 1) self.perspective = 0 self.solid = 1 self.idFilter = vtk.vtkIdFilter() self.idFilter.SetInputData(self.poly_data) self.idFilter.SetIdsArrayName("OriginalIds") self.idFilter.Update() self.surfaceFilter = vtk.vtkDataSetSurfaceFilter() self.surfaceFilter.SetInputConnection(self.idFilter.GetOutputPort()) self.surfaceFilter.Update() self.input = self.surfaceFilter.GetOutput() self.renderer = vtk.vtkRenderer() #self.renderer2 = vtk_widget.vtkRenderer() viewport = [0.0,0.0,0.15,0.15] #self.renderer2.SetViewport(viewport) #self.renderer2.Transparent() self.renderWindowInteractor = QVTKRenderWindowInteractor(self.ui.frame) self.renderWindowInteractor.GetRenderWindow().AddRenderer(self.renderer) #self.renderWindowInteractor.GetRenderWindow().AddRenderer(self.renderer2) self.renderWindowInteractor.GetRenderWindow().SetAlphaBitPlanes(1) self.axes = CoordinateAxes(self.renderWindowInteractor) self.ui.vl.addWidget(self.renderWindowInteractor) self.iren = vtk.vtkRenderWindowInteractor() self.iren = self.renderWindowInteractor.GetRenderWindow().GetInteractor() self.mapper = vtk.vtkPolyDataMapper() self.mapper.SetInputData(self.input) self.mapper.ScalarVisibilityOff() self.actor = vtk.vtkActor() self.actor.SetMapper(self.mapper) self.actor.GetProperty().SetPointSize(2) self.actor.GetProperty().EdgeVisibilityOn() self.actor.GetProperty().SetColor(self.eidcolor) self.actor.GetProperty().SetEdgeColor(self.edgecolor) self.camera = vtk.vtkCamera() #self.camera2 = vtk_widget.vtkCamera() # trial... add glyph pd = vtk.vtkPolyData() pts = vtk.vtkPoints() scalars = vtk.vtkFloatArray() vectors = vtk.vtkFloatArray() vectors.SetNumberOfComponents(3) pd.SetPoints(pts) pd.GetPointData().SetScalars(scalars) pd.GetPointData().SetVectors(vectors) pts.InsertNextPoint(30, 30, 0.0) scalars.InsertNextValue(1) vectors.InsertNextTuple3(1, 1, 0.0) # Create simple PolyData for glyph table cs = vtk.vtkCubeSource() cs.SetXLength(0.5) cs.SetYLength(1) cs.SetZLength(2) # Set up the glyph filter glyph = vtk.vtkGlyph3D() #glyph.SetInputConnection(elev.GetOutputPort()) point_list = vtk.vtkPoints() point_list.InsertNextPoint([30, 30, 0]) poly_data = vtk.vtkPolyData() poly_data.SetPoints(point_list) idFilter = vtk.vtkIdFilter() idFilter.SetInputData(poly_data) idFilter.SetIdsArrayName("OriginalIds") idFilter.Update() surfaceFilter = vtk.vtkDataSetSurfaceFilter() surfaceFilter.SetInputConnection(idFilter.GetOutputPort()) surfaceFilter.Update() # Here is where we build the glyph table # that will be indexed into according to the IndexMode glyph.SetSourceData(0,cs.GetOutput()) #glyph.SetInputConnection(surfaceFilter.GetOutputPort()) glyph.SetInputData(pd) glyph.SetIndexModeToScalar() glyph.SetRange(0, 1) glyph.SetScaleModeToDataScalingOff() glyph.OrientOn() mapper3 = vtk.vtkPolyDataMapper() mapper3.SetInputConnection(glyph.GetOutputPort()) mapper3.SetScalarModeToUsePointFieldData() mapper3.SetColorModeToMapScalars() mapper3.ScalarVisibilityOn() mapper3.SetScalarRange(0, 1) actor3 = vtk.vtkActor() actor3.SetMapper(mapper3) #actor3.GetProperty().SetBackgroundOpacity(0.5) gs = vtk.vtkGlyphSource2D() gs.SetGlyphTypeToCircle() gs.SetScale(25) gs.FilledOff() #gs.CrossOn() gs.Update() # Create a table of glyphs glypher = vtk.vtkGlyph2D() glypher.SetInputData(pd) glypher.SetSourceData(0, gs.GetOutput()) glypher.SetIndexModeToScalar() glypher.SetRange(0, 1) glypher.SetScaleModeToDataScalingOff() mapper = vtk.vtkPolyDataMapper2D() mapper.SetInputConnection(glypher.GetOutputPort()) mapper.SetScalarRange(0, 1) actor2D = vtk.vtkActor2D() actor2D.SetMapper(mapper) self.renderer.AddActor(self.actor) #self.renderer.AddActor(mapper) #self.renderer2.AddActor(actor3) self.renderer.SetBackground(self.bgcolor1) self.renderer.SetBackground2(self.bgcolor2) self.renderer.GradientBackgroundOn() self.renderer.SetActiveCamera(self.camera) self.renderer.ResetCamera() #self.camera.ZoomOff() #self.renderer2.SetActiveCamera(self.camera) #self.renderer2.ResetCamera() #self.renderer2.SetBackground(0,0,0) #self.renderer2.GetProperty().SetBackgroundOpacity(0.5) #self.renderer2.SetLayer(1) #self.renderer2.Clear() self.areaPicker = vtk.vtkAreaPicker() self.renderWindowInteractor.SetPicker(self.areaPicker) self.style = MyInteractorStyle() self.style.SetPoints(self.input) self.style.SetDefaultRenderer(self.renderer) self.style.Data = self.idFilter.GetOutput() self.style.camera = self.camera self.style.node_ids = self.node_ids self.style.element_ids = self.element_ids self.style.node_count = self.node_count self.style.window = self self.style.print_message = self.ui.txt_msg.appendPlainText self.renderWindowInteractor.SetInteractorStyle(self.style) self.renderWindowInteractor.Start() # screenshot code:e #self.w2if = vtk_widget.vtkWindowToImageFilter() #self.w2if.SetInput(self.renWin) #self.w2if.Update() self.show() self.iren.Initialize() #self.iren.Start() # Setup Connections self.ui.btn_bgcolor1.clicked.connect(self.on_color1) self.ui.btn_bgcolor2.clicked.connect(self.on_color2) self.ui.btn_edgecolor.clicked.connect(self.on_edgecolor) self.ui.btn_elementcolor.clicked.connect(self.on_elementcolor) self.ui.btn_nofillededge.clicked.connect(self.on_nofillededge) self.ui.btn_switch.clicked.connect(self.on_switch) self.ui.btn_perspectivetoggle.clicked.connect(self.on_toggleperspective) self.ui.btn_saveimg.clicked.connect(self.on_saveimg) self.ui.btn_togglewire.clicked.connect(self.on_togglewire) # Setup a shortcuts self.setusrstyle = QtGui.QShortcut(self) self.setusrstyle.setKey(("CTRL+c")) self.setusrstyle.activated.connect(self.on_copyimg) def read_data(self, filename): # Create source point_list = vtk.vtkPoints() cell_list = vtk.vtkCellArray() nodes = {} self.node_ids = array('i') from fem_reader.nastran.bdf.reader import BDFReader bdf = BDFReader() bdf.read_bdf(r'../data/wing.bdf') grids = bdf.nodes.keys() elements = bdf.elements.keys() for i in xrange(len(grids)): # grids data_type = bdf.nodes[grids[i]].card_name if data_type != 'GRID': continue node_id = str(grids[i]) self.node_ids.append(int(node_id)) nodes[node_id] = {} nodes[node_id]['order'] = i nodes[node_id]['id'] = node_id #x = bdf.nodes[grids[i]].X1 #float(line[2]) # grids[i].x #y = bdf.nodes[grids[i]].X2 # float(line[3]) #z = bdf.nodes[grids[i]].X3 # float(line[4]) x, y, z = bdf.nodes[grids[i]].to_global() nodes[node_id]['x'] = float(x) nodes[node_id]['y'] = float(y) nodes[node_id]['z'] = float(z) _id = point_list.InsertNextPoint([x, y, z]) cell_list.InsertNextCell(1) cell_list.InsertCellPoint(_id) self.node_count = len(self.node_ids) last_i = self.node_count self.element_ids = array('i') for i in xrange(len(elements)): #line = lines[i].split(',') #data_type = line[0] #print i data_type = bdf.elements[elements[i]].card_name if data_type == 'GRID': continue try: self.element_ids.append(int(bdf.elements[elements[i]].ID)) except IndexError: break if data_type == 'CBEAM': beam = Bar2() beam.id = bdf.elements[elements[i]].ID beam.order = last_i + i node1 = nodes[str(bdf.elements[elements[i]].G1)] node2 = nodes[str(bdf.elements[elements[i]].G2)] beam.set_nodes(node1, node2) cell_list.InsertNextCell(beam.vtk_cell) point_list.InsertNextPoint(beam.center) if data_type == 'CTRIA3': tria = Tria3() tria.id = bdf.elements[elements[i]].ID tria.order = last_i + i node1 = nodes[str(bdf.elements[elements[i]].G1)] #bdf.elements[elements[i]].G1 node2 = nodes[str(bdf.elements[elements[i]].G2)] node3 = nodes[str(bdf.elements[elements[i]].G3)] tria.set_nodes(node1, node2, node3) cell_list.InsertNextCell(tria.vtk_cell) point_list.InsertNextPoint(tria.center) if data_type == 'CQUAD4': quad = Quad4() quad.id = bdf.elements[elements[i]].ID quad.order = last_i + i node1 = nodes[str(bdf.elements[elements[i]].G1)] node2 = nodes[str(bdf.elements[elements[i]].G2)] node3 = nodes[str(bdf.elements[elements[i]].G3)] node4 = nodes[str(bdf.elements[elements[i]].G4)] quad.set_nodes(node1, node2, node3, node4) cell_list.InsertNextCell(quad.vtk_cell) point_list.InsertNextPoint(quad.center) self.poly_data = vtk.vtkPolyData() self.poly_data.SetPoints(point_list) self.poly_data.SetPolys(cell_list) self.poly_data.SetLines(cell_list) #self.poly_data.SetVerts(cell_list) def on_color1(self): color = QtGui.QColorDialog.getColor(QtCore.Qt.blue,self) red = color.red() / 255. blue = color.blue() / 255. green = color.green() / 255. self.bgcolor1 = (red, green, blue) self.renderer.SetBackground(self.bgcolor1) self.show() def on_color2(self): color = QtGui.QColorDialog.getColor(QtCore.Qt.blue,self) red = color.red() / 255. blue = color.blue() / 255. green = color.green() / 255. self.bgcolor2 = (red, green, blue) self.renderer.SetBackground2(self.bgcolor2) self.show() def on_edgecolor(self): color = QtGui.QColorDialog.getColor(QtCore.Qt.blue,self) red = color.red() / 255. blue = color.blue() / 255. green = color.green() / 255. self.edgecolor = (red, green, blue) self.actor.GetProperty().SetEdgeColor(self.edgecolor) self.actor.GetProperty().EdgeVisibilityOn() self.actor.GetProperty().SetPointSize(2) self.show() def on_elementcolor(self): color = QtGui.QColorDialog.getColor(QtCore.Qt.blue,self) red = color.red() / 255. blue = color.blue() / 255. green = color.green() / 255. self.eidcolor = (red, green, blue) self.actor.GetProperty().SetColor(self.eidcolor) self.show() def on_nofillededge(self): self.actor.GetProperty().EdgeVisibilityOff() self.actor.GetProperty().SetPointSize(0.001) self.show() def on_switch(self): if self.bgcolor1 == (1,1,1) and self.bgcolor2 == (1,1,1): self.bgcolor1 = (0, 0, 1) self.bgcolor2 = (0.8, 0.8, 1) self.renderer.SetBackground(self.bgcolor1) self.renderer.SetBackground2(self.bgcolor2) self.show() else: self.bgcolor1 = (1, 1, 1) self.bgcolor2 = (1, 1, 1) self.renderer.SetBackground(self.bgcolor1) self.renderer.SetBackground2(self.bgcolor2) self.show() def on_toggleperspective(self): if self.perspective == 0: #self.actor.GetProperty().ParallelProjectionOn() self.camera.ParallelProjectionOn() self.show() self.perspective = 1 else: #self.actor.GetProperty().ParallelProjectionOff() self.camera.ParallelProjectionOff() self.show() self.perspective = 0 def on_saveimg(self): try: picfile = QtGui.QFileDialog.getSaveFileName(self, "Save Image", '', "Image files (*.png)") except IOError as err: print err if picfile == []: return print picfile[0] # screenshot code: w2if = vtk.vtkWindowToImageFilter() w2if.SetInput(self.renderWindowInteractor.GetRenderWindow()) w2if.Update() writer = vtk.vtkPNGWriter() writer.SetFileName(picfile[0]) writer.SetInput(w2if.GetOutput()) writer.Write() def on_copyimg(self): # Take a screenshot: w2if = vtk.vtkWindowToImageFilter() w2if.SetInput(self.renderWindowInteractor.GetRenderWindow()) w2if.Update() # screenshot is a vtk_widget object image = w2if.GetOutput() # write a temp image file writer = vtk.vtkPNGWriter() writer.SetFileName("tempfile.png") writer.SetInput(image) writer.Write() # read the temp image file ### This works... copying image from file to clipboard self.clipboard = QtGui.QApplication.clipboard() data = QtCore.QMimeData() #data.setImageData(QtGui.QImage(r'D:\PGM\01_DEV\VTK\MVC\1.png')) data.setImageData(QtGui.QImage("tempfile.png")) self.clipboard.setMimeData(data) # remove the tempfile os.remove("tempfile.png") # how to covert vtkobject image to Qimage?? def on_togglewire(self): pass
class PickWidget(QWidget): def __init__(self, parent=None): QWidget.__init__(self, parent) self.setStyleSheet('background-color:#333333') self.vtk_layout = QtGui.QGridLayout(self) self.vtk_widget = QVTKRenderWindowInteractor(self) self.vtk_layout.addWidget(self.vtk_widget) self.ren = vtk.vtkRenderer() self.vtk_widget.GetRenderWindow().AddRenderer(self.ren) self.meshMapper = vtk.vtkPolyDataMapper() self.meshActor = vtk.vtkActor() self.meshActor.SetMapper(self.meshMapper) self.ren.AddActor(self.meshActor) self.interactionStyle = RubberBandFacePicker() self.interactionStyle.SetDefaultRenderer(self.ren) self.vtk_widget.SetInteractorStyle(self.interactionStyle) self.vtk_widget.SetPicker(vtk.vtkAreaPicker()) self.iren = self.vtk_widget.GetRenderWindow().GetInteractor() self.show() self.iren.Initialize() # name2actor map for calls to display_XXXX self.name2actor = {} def start(self): self.reader = vtk.vtkOBJReader() self.reader.SetFileName(self.filename) self.meshMapper.SetInputConnection(self.reader.GetOutputPort()) self.interactionStyle.Data = self.reader.GetOutput() self.interactionStyle.Widget = self def set_source(self, filename): self.filename = filename def set_selected_color(self, color): self.interactionStyle.set_selected_pt_color(color) def get_selected_vertices(self): V = self.interactionStyle.get_selected_vertices() return V def display_cloud(self, actorName, list_points, size, vec3_color): # extract inlier points vtk_points = vtk.vtkPoints() for p in list_points: vtk_points.InsertNextPoint(p[0], p[1], p[2]) temp_cloud = vtk.vtkPolyData() temp_cloud.SetPoints(vtk_points) # add 0D topology to every point glfilter = vtk.vtkVertexGlyphFilter() glfilter.SetInput(temp_cloud) glfilter.Update() cloud = glfilter.GetOutput() cloudMapper = vtk.vtkPolyDataMapper() cloudMapper.SetInput(cloud) actor = vtk.vtkActor() actor.SetMapper(cloudMapper) actor.GetProperty().SetColor(vec3_color[0], vec3_color[1], vec3_color[2]) actor.GetProperty().SetPointSize(size) self.ren.AddActor(actor) self.name2actor[actorName] = actor pass def display_sphere(self, actorName, vec3_center, radius, vec3_color): sphere = vtk.vtkSphereSource() sphere.SetRadius(radius) sphere.SetCenter(vec3_center) sphere.SetThetaResolution(40) sphere.SetPhiResolution(40) sphere.SetCenter(vec3_center[0], vec3_center[1], vec3_center[2]) sphere.Update() mapper = vtk.vtkPolyDataMapper() mapper.SetInputConnection(sphere.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(mapper) actor.GetProperty().SetColor(vec3_color[0], vec3_color[1], vec3_color[2]) self.ren.AddActor(actor) self.name2actor[actorName] = actor pass def display_cylinder(self, actorName, vec3_origin, vec3_direction, height, radius, vec3_color): cyl = vtk.vtkCylinderSource() cyl.SetRadius(radius) cyl.SetHeight(height) cyl.SetCapping(1) cyl.Update() # transform to center with orientation according to normal axis1 = [vec3_direction[0], vec3_direction[1], vec3_direction[2]] axis2 = np.cross(axis1, [1, 1, 1]) axis2 = axis2 / np.linalg.norm(axis2) axis3 = np.cross(axis2, axis1) # print axis1 # print axis2 # print axis3 trans = np.eye(4) trans[0, 0] = axis2[0] trans[0, 1] = axis1[0] trans[0, 2] = axis3[0] trans[0, 3] = vec3_origin[0] trans[1, 0] = axis2[1] trans[1, 1] = axis1[1] trans[1, 2] = axis3[1] trans[1, 3] = vec3_origin[1] trans[2, 0] = axis2[2] trans[2, 1] = axis1[2] trans[2, 2] = axis3[2] trans[2, 3] = vec3_origin[2] trans[3, 0] = 0 trans[3, 1] = 0 trans[3, 2] = 0 trans[3, 3] = 1 vtk_trans = vtk.vtkMatrix4x4() for i in range(0, 4): for j in range(0, 4): vtk_trans.SetElement(i, j, trans[i, j]) ar_trans = vtk.vtkTransform() ar_trans.SetMatrix(vtk_trans) ar_trans_filter = vtk.vtkTransformPolyDataFilter() ar_trans_filter.SetTransform(ar_trans) ar_trans_filter.SetInputConnection(cyl.GetOutputPort()) mapper = vtk.vtkPolyDataMapper() mapper.SetInputConnection(ar_trans_filter.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(mapper) actor.GetProperty().SetColor(vec3_color[0], vec3_color[1], vec3_color[2]) self.ren.AddActor(actor) self.name2actor[actorName] = actor return def display_disk(self, actorName, vec3_origin, vec3_direction, outer_radius, inner_radius, vec3_color): disk = vtk.vtkDiskSource() disk.SetInnerRadius(inner_radius) disk.SetOuterRadius(outer_radius) disk.SetRadialResolution(30) disk.SetCircumferentialResolution(30) disk.Update() # transform to center with orientation according to normal axis1 = [vec3_direction[0], vec3_direction[1], vec3_direction[2]] axis2 = np.cross(axis1, [1, 1, 1]) axis2 = axis2 / np.linalg.norm(axis2) axis3 = np.cross(axis1, axis2) # print axis1 # print axis2 # print axis3 trans = np.eye(4) trans[0, 0] = axis3[0] trans[0, 1] = axis2[0] trans[0, 2] = axis1[0] trans[0, 3] = vec3_origin[0] trans[1, 0] = axis3[1] trans[1, 1] = axis2[1] trans[1, 2] = axis1[1] trans[1, 3] = vec3_origin[1] trans[2, 0] = axis3[2] trans[2, 1] = axis2[2] trans[2, 2] = axis1[2] trans[2, 3] = vec3_origin[2] trans[3, 0] = 0 trans[3, 1] = 0 trans[3, 2] = 0 trans[3, 3] = 1 vtk_trans = vtk.vtkMatrix4x4() for i in range(0, 4): for j in range(0, 4): vtk_trans.SetElement(i, j, trans[i, j]) ar_trans = vtk.vtkTransform() ar_trans.SetMatrix(vtk_trans) ar_trans_filter = vtk.vtkTransformPolyDataFilter() ar_trans_filter.SetTransform(ar_trans) ar_trans_filter.SetInputConnection(disk.GetOutputPort()) diskMapper = vtk.vtkPolyDataMapper() diskMapper.SetInputConnection(ar_trans_filter.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(diskMapper) actor.GetProperty().SetColor(vec3_color[0], vec3_color[1], vec3_color[2]) self.ren.AddActor(actor) self.name2actor[actorName] = actor return def display_arrow(self, actorName, vec3_origin, vec3_direction, length_in_m, shaft_radius, vec3_color): # visualize direction arrow arrow = vtk.vtkArrowSource() arrow.SetTipLength(0.25) arrow.SetTipRadius((1.5 * shaft_radius) / length_in_m) arrow.SetTipResolution(20) arrow.SetShaftRadius(shaft_radius / length_in_m) # account for scaling arrow.SetShaftResolution(20) arrow.Update() # scale scale = vtk.vtkTransform() scale.Scale(length_in_m, length_in_m, length_in_m) scale_filter = vtk.vtkTransformPolyDataFilter() scale_filter.SetTransform(scale) scale_filter.SetInputConnection(arrow.GetOutputPort()) # transform to center with orientation according to normal axis1 = [vec3_direction[0], vec3_direction[1], vec3_direction[2]] axis2 = np.cross(axis1, [1, 1, 1]) axis2 = axis2 / np.linalg.norm(axis2) axis3 = np.cross(axis1, axis2) # print axis1 # print axis2 # print axis3 trans = np.eye(4) trans[0, 0] = axis1[0] trans[0, 1] = axis2[0] trans[0, 2] = axis3[0] trans[0, 3] = vec3_origin[0] trans[1, 0] = axis1[1] trans[1, 1] = axis2[1] trans[1, 2] = axis3[1] trans[1, 3] = vec3_origin[1] trans[2, 0] = axis1[2] trans[2, 1] = axis2[2] trans[2, 2] = axis3[2] trans[2, 3] = vec3_origin[2] trans[3, 0] = 0 trans[3, 1] = 0 trans[3, 2] = 0 trans[3, 3] = 1 vtk_trans = vtk.vtkMatrix4x4() for i in range(0, 4): for j in range(0, 4): vtk_trans.SetElement(i, j, trans[i, j]) ar_trans = vtk.vtkTransform() ar_trans.SetMatrix(vtk_trans) ar_trans_filter = vtk.vtkTransformPolyDataFilter() ar_trans_filter.SetTransform(ar_trans) ar_trans_filter.SetInputConnection(scale_filter.GetOutputPort()) arrowMapper = vtk.vtkPolyDataMapper() arrowMapper.SetInputConnection(ar_trans_filter.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(arrowMapper) actor.GetProperty().SetColor(vec3_color[0], vec3_color[1], vec3_color[2]) self.ren.AddActor(actor) self.name2actor[actorName] = actor pass def display_affordance_primitive(self, g, type_): self.remove_all_displays() if type_ == 'axis': self.display_arrow('p', g.base, g.direction, g.length, g.radius, [1, .2, 0]) elif type_ == 'plane': self.display_disk('p', g.center, g.normal, g.radius, 0, [1, .2, 0]) self.display_arrow('p', g.center, g.normal, .1, .002, [1, .3, 0]) def remove_display(self, actorName): actor = self.name2actor.get(actorName) if (actor != None): self.ren.RemoveActor(actor) self.name2actor.pop(actorName) pass def remove_all_displays(self): for actor in self.name2actor.viewvalues(): self.ren.RemoveActor(actor) self.name2actor.clear() pass