[1231] | 1 | # matlab-like funcs using VTK, to be used in matplotlib |
---|
| 2 | # Randy Heiland, Indiana University |
---|
| 3 | |
---|
| 4 | from pylab import zeros |
---|
| 5 | from vtk import * |
---|
| 6 | from vtkPipe import * |
---|
| 7 | |
---|
| 8 | meshActor = vtkActor() |
---|
| 9 | plot3Actor = vtkActor() |
---|
| 10 | |
---|
| 11 | def mesh(x,y,z, colored): |
---|
| 12 | # global zlo,zmax |
---|
| 13 | pts = vtkPoints() |
---|
| 14 | xlen = len(x[1]) |
---|
| 15 | ylen = len(y) |
---|
| 16 | print 'xlen,ylen=',xlen,ylen |
---|
| 17 | print 'zsize=',size(z) |
---|
| 18 | pts.SetNumberOfPoints(size(z)) |
---|
| 19 | idx = 0 |
---|
| 20 | zlo = z[0,0] |
---|
| 21 | zmax = z[0,0] |
---|
| 22 | for iy in range(ylen): |
---|
| 23 | for ix in range(xlen): |
---|
| 24 | pts.SetPoint(idx, x[iy,ix], y[iy,ix], z[iy,ix]) |
---|
| 25 | idx = idx + 1 |
---|
| 26 | if z[iy,ix] < zlo: |
---|
| 27 | zlo = z[iy,ix] |
---|
| 28 | if z[iy,ix] > zmax: |
---|
| 29 | zmax = z[iy,ix] |
---|
| 30 | print 'zlo= ',zlo |
---|
| 31 | |
---|
| 32 | meshSG = vtkStructuredGrid() |
---|
| 33 | meshSG.SetDimensions(xlen,ylen,1) |
---|
| 34 | meshSG.SetPoints(pts) |
---|
| 35 | |
---|
| 36 | meshGeom = vtkStructuredGridGeometryFilter() |
---|
| 37 | meshGeom.SetInput(meshSG) |
---|
| 38 | |
---|
| 39 | if colored > 0: |
---|
| 40 | elev = vtkElevationFilter() |
---|
| 41 | elev.SetInput(meshGeom.GetOutput()) |
---|
| 42 | pdnormals = vtkPolyDataNormals() |
---|
| 43 | pdnormals.SetInput(elev.GetPolyDataOutput()) |
---|
| 44 | |
---|
| 45 | meshMapper = vtkPolyDataMapper() |
---|
| 46 | |
---|
| 47 | if colored > 0: |
---|
| 48 | meshLUT = vtkLookupTable() |
---|
| 49 | meshLUT.SetHueRange(lutBlueRed.GetHueRange()) |
---|
| 50 | meshLUT.SetSaturationRange(lutBlueRed.GetSaturationRange()) |
---|
| 51 | meshLUT.SetValueRange(lutBlueRed.GetValueRange()) |
---|
| 52 | meshLUT.Build() |
---|
| 53 | |
---|
| 54 | meshMapper.SetInput(pdnormals.GetOutput()) |
---|
| 55 | meshMapper.ScalarVisibilityOn() |
---|
| 56 | meshMapper.SetScalarRange(zlo,zmax) |
---|
| 57 | meshMapper.SetLookupTable(meshLUT) |
---|
| 58 | meshLUT.SetTableRange(zlo,zmax) |
---|
| 59 | else: |
---|
| 60 | meshMapper.SetInput(meshGeom.GetOutput()) |
---|
| 61 | |
---|
| 62 | |
---|
| 63 | meshActor.SetMapper(meshMapper) |
---|
| 64 | vtkRen.AddActor(meshActor) |
---|
| 65 | |
---|
| 66 | |
---|
| 67 | def plot3(x,y,z): |
---|
| 68 | pts = vtkPoints() |
---|
| 69 | vlen = len(x) |
---|
| 70 | print 'vlen=',vlen |
---|
| 71 | pts.SetNumberOfPoints(vlen) |
---|
| 72 | for idx in range(vlen): |
---|
| 73 | pts.SetPoint(idx, x[idx], y[idx], z[idx]) |
---|
| 74 | print 'plot3 bounds= ',pts.GetBounds() |
---|
| 75 | |
---|
| 76 | # create a polyline (cellarray) |
---|
| 77 | pl = vtkCellArray() |
---|
| 78 | pl.InsertNextCell(vlen) |
---|
| 79 | for idx in range(vlen): |
---|
| 80 | pl.InsertCellPoint(idx) |
---|
| 81 | |
---|
| 82 | pd = vtkPolyData() |
---|
| 83 | pd.SetPoints(pts) |
---|
| 84 | pd.SetLines(pl) |
---|
| 85 | |
---|
| 86 | mapper = vtkPolyDataMapper() |
---|
| 87 | mapper.SetInput(pd) |
---|
| 88 | |
---|
| 89 | plot3Actor.SetMapper(mapper) |
---|
| 90 | plot3Actor.SetScale(1,1,0.1) |
---|
| 91 | vtkRen.AddActor(plot3Actor) |
---|
| 92 | |
---|
| 93 | # Create a text property for both cube axes |
---|
| 94 | tprop = vtkTextProperty() |
---|
| 95 | tprop.SetColor(1, 1, 1) |
---|
| 96 | tprop.ShadowOn() |
---|
| 97 | |
---|
| 98 | # rf. /home/heiland/VTK/Examples/Annotation/Python/cubeAxes.py |
---|
| 99 | axes = vtkCubeAxesActor2D() |
---|
| 100 | axes.SetInput(pd) |
---|
| 101 | axes.SetCamera(vtkRen.GetActiveCamera()) |
---|
| 102 | axes.SetLabelFormat("%6.4g") |
---|
| 103 | axes.SetFlyModeToOuterEdges() |
---|
| 104 | axes.SetFontFactor(0.8) |
---|
| 105 | axes.SetAxisTitleTextProperty(tprop) |
---|
| 106 | axes.SetAxisLabelTextProperty(tprop) |
---|
| 107 | |
---|
| 108 | #------------------------------------------------------------------ |
---|
| 109 | def vtkRender(): |
---|
| 110 | vtkRenWin.Render() |
---|
| 111 | |
---|
| 112 | def vtkClear(): |
---|
| 113 | # remove all actors hack |
---|
| 114 | vtkRen.RemoveActor(meshActor) |
---|
| 115 | vtkRen.RemoveActor(plot3Actor) |
---|
| 116 | return |
---|
| 117 | |
---|
| 118 | def vtkImage(): |
---|
| 119 | vtkRenWin.Render() |
---|
| 120 | vtkRenWin.GetPixelData(0,0,winWidth-1,winHeight-1, 1,vtkRGB) |
---|
| 121 | # vtkRGB.Squeeze() |
---|
| 122 | idx=0 |
---|
| 123 | for iy in range(winHeight-1,-1,-1): |
---|
| 124 | for ix in range(winWidth): |
---|
| 125 | im[iy,ix,0] = vtkRGB.GetValue(idx) / 255. |
---|
| 126 | im[iy,ix,1] = vtkRGB.GetValue(idx+1) / 255. |
---|
| 127 | im[iy,ix,2] = vtkRGB.GetValue(idx+2) / 255. |
---|
| 128 | idx += 3 |
---|
| 129 | |
---|
| 130 | def vtkRotX(degs): |
---|
| 131 | # vtkRen.GetActiveCamera().ComputeViewPlaneNormal() |
---|
| 132 | vtkRen.ResetCamera() |
---|
| 133 | vtkRen.GetActiveCamera().Elevation(degs) |
---|
| 134 | |
---|
| 135 | def vtkPerspective(flag): |
---|
| 136 | if flag > 0: |
---|
| 137 | vtkRen.GetActiveCamera().ParallelProjectionOff() |
---|
| 138 | else: |
---|
| 139 | vtkRen.GetActiveCamera().ParallelProjectionOn() |
---|
| 140 | |
---|
| 141 | def vtkWinsize(w,h): |
---|
| 142 | winWidth = w |
---|
| 143 | winHeight = h |
---|
| 144 | im = zeros((winWidth,winHeight,3), typecode=Float) |
---|
| 145 | vtkRenWin.SetSize(winWidth,winHeight) |
---|