tools/vizTools/velVis.py

#!/usr/bin/env python

#
# $Id$
#
#######################################################
#
#           Copyright 2003-2007 by ACceSS MNRF
#       Copyright 2007 by University of Queensland
#
#                http://esscc.uq.edu.au
#        Primary Business: Queensland, Australia
#  Licensed under the Open Software License version 3.0
#     http://www.opensource.org/licenses/osl-3.0.php
#
#######################################################
#

import vtk
import sys, os, re
import getopt

(opts, args) = getopt.getopt(sys.argv[1:],
        "d:f:ph",
        ["dirname=", "fnamestem=", "preproc", "help"],
        )

def usage():
    print "Usage:"
    print "  velVis.py -d <dirname> -f <fnamestem> -p -i\n"
    print "  Arguments:"
    print "   -d/--dirname: directory of vtk files to process (req)"
    print "   -f/--fnamestem: filename stem of vtk files to process (req)"
    print "   -p/--preproc: preprocess vtk files to find global vector max?  (opt)"
    print "   -i/--interactive: interactive use of first frame? (opt)"

#dirName = '/data/raid2/matt/results/050516/m3dpData/vis'
dirName = None
fnameStem = None
interactive = False
preproc = False   # preprocess to find the global max of the data

for option, arg in opts:
    if option in ('-d', '--dirname'):
        dirName = arg
    elif option in ('-f', '--fnamestem'):
        fnameStem = arg
    elif option in ('-p', '--preproc'):
        preproc = True
    elif option in ('-i', '--interactive'):
        interactive = True
    elif option in ('-h', '--help'):
        usage()
        sys.exit(0)

if dirName is None:
    print "You must supply a directory of vtk files to process\n"
    usage()
    sys.exit(1)

if fnameStem is None:
    print "You must supply the filename stem of the files to process"
    print "The filename stem is the part of the filename without the"
    print "frame number or extension\n"
    usage()
    sys.exit(1)

# now need to determine the number of files to process
dirList = os.listdir(dirName)
r = re.compile( "%s\\d+\\.\\w+" % fnameStem )
count = 0  # counter for the number of files found
fnames = []
for fname in dirList:
    if r.match(fname):
        fnames.append(fname)
        count += 1

# determine the first frame number and number of frames
fnames.sort()
firstFile = fnames[0]

r = re.compile("([a-zA-Z-])(\\d+)(\\.)(\\w+)")
firstNum = r.findall(firstFile)
firstNum = int(firstNum[0][1])

# grab the filename extension as well, just in case
extn = r.findall(firstFile)
extn = extn[0][3]

# the number of frames should be equal to the count
if interactive:
    numFrames = 1
    mesa = False
else:
    numFrames = count
    mesa = True

# try and work out the format of the zero-padded frame number
if count < 10:
    formatString = '"%01d"'
elif count < 100:
    formatString = '"%02d"'
elif count < 1000:
    formatString = '"%03d"'
elif count < 10000:
    formatString = '"%04d"'
elif count < 100000:
    formatString = '"%05d"'
else:
    print "Sorry, can't handle greater than 100000 input frames"
    sys.exit(1)

# use mesa if required
if mesa:
    factGraphics = vtk.vtkGraphicsFactory()
    factGraphics.SetUseMesaClasses(1)

    factImage = vtk.vtkImagingFactory()
    factImage.SetUseMesaClasses(1)

# perform some preprocessing to find the global max of the vector data norm
maxNorm = 0
if not interactive and preproc:
    print "Preprocessing vector data"
    sys.stdout.write("Completed: ")
    for i in range(numFrames):
        frameNum = eval(formatString) % i
        reader = vtk.vtkXMLUnstructuredGridReader()
        reader.SetFileName("%s/%s%s.%s" % (dirName,fnameStem,frameNum,extn))
        reader.Update()

        grid = reader.GetOutput()
        norm = grid.GetPointData().GetVectors().GetMaxNorm()
        if norm > maxNorm:
            maxNorm = norm

    sys.stdout.write("%3.0f%%\b\b\b\b" % (float(i)/float(numFrames), ))
    print ""

for i in range(firstNum, numFrames+firstNum):

    frameNum = eval(formatString) % i
    reader = vtk.vtkXMLUnstructuredGridReader()
    reader.SetFileName("%s/%s%s.%s" % (dirName,fnameStem,frameNum,extn))
    reader.Update()

    print "Processing %s" % fnames[i]

    grid = reader.GetOutput()

    # grab the model centre and bounds
    centre = grid.GetCenter()
    bounds = grid.GetBounds()

    # grab the norm of the vectors
    norm = vtk.vtkVectorNorm()
    norm.SetInput(grid)

    if not preproc:
        maxNorm = grid.GetPointData().GetVectors().GetMaxNorm()

    # to make arrow glyphs need an arrow source
    arrow = vtk.vtkArrowSource()
    
    # the arrows are 3D glyphs so set that up now
    glyph = vtk.vtkGlyph3D()
    glyph.ScalingOn()
    glyph.SetScaleModeToScaleByScalar()
    glyph.SetColorModeToColorByScalar()
    glyph.SetVectorModeToUseVector()
    glyph.SetScaleFactor(0.1/maxNorm)
    glyph.SetInput(norm.GetOutput())
    glyph.SetSource(arrow.GetOutput())
    glyph.ClampingOff()

    # set up a stripper to speed up rendening
    stripper = vtk.vtkStripper()
    stripper.SetInput(glyph.GetOutput())
    
    # make a lookup table for the colour map and invert it (colours look
    # better when it's inverted)
    lut = vtk.vtkLookupTable()
    refLut = vtk.vtkLookupTable()
    lut.Build()
    refLut.Build()
    for j in range(256):
        lut.SetTableValue(j, refLut.GetTableValue(255-j))
    
    # set up the mapper
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInput(stripper.GetOutput())
    mapper.SetScalarRange(0,maxNorm)
    mapper.SetLookupTable(lut)
    
    # set up the actor
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
   
    # set up the text properties for nice text
    font_size = 20
    textProp = vtk.vtkTextProperty()
    textProp.SetFontSize(font_size)
    textProp.SetFontFamilyToArial()
    textProp.BoldOff()
    textProp.ItalicOff()
    textProp.ShadowOff()
    textProp.SetColor(0.0, 0.0, 0.0)
 
    # make a title
    title = vtk.vtkTextMapper()
    title.SetInput("Velocity flow vectors in mantle convection")

    # make the title text use the text properties
    titleProp = title.GetTextProperty()
    titleProp.ShallowCopy(textProp)
    titleProp.SetJustificationToCentered()
    titleProp.SetVerticalJustificationToTop()
    titleProp.BoldOn()

    # make the actor for the title
    titleActor = vtk.vtkTextActor()
    titleActor.SetMapper(title)
    titleActor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay()
    titleActor.GetPositionCoordinate().SetValue(0.5, 0.95)

    # make a frame counter
    counter = vtk.vtkTextMapper()
    counter.SetInput("frame: %d" % i)

    # make the counter use the text properties
    counterProp = counter.GetTextProperty()
    counterProp.ShallowCopy(textProp)
    counterProp.SetJustificationToLeft()
    counterProp.SetVerticalJustificationToTop()
    counterProp.SetFontSize(14)

    # make the actor for the frame counter
    counterActor = vtk.vtkTextActor()
    counterActor.SetMapper(counter)
    counterActor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay()
    counterActor.GetPositionCoordinate().SetValue(0.05, 0.8)

    # use a scalar bar
    scalarBar = vtk.vtkScalarBarActor()
    scalarBar.SetLookupTable(lut)
    scalarBar.SetWidth(0.1)
    scalarBar.SetHeight(0.7)
    scalarBar.SetPosition(0.9, 0.2)
    scalarBar.SetTitle("|v|")

    # set up its title text properties
    scalarBarProp = scalarBar.GetTitleTextProperty()
    scalarBarProp.ShallowCopy(textProp)
    scalarBarProp.SetFontSize(10)

    # set up the label text properties 
    scalarBarTextProp = scalarBar.GetLabelTextProperty()
    scalarBarTextProp.ShallowCopy(textProp)
    scalarBarTextProp.SetFontSize(10)

    # put an outline around the data
    outline = vtk.vtkOutlineSource()
    outline.SetBounds(bounds)

    # make its mapper
    outlineMapper = vtk.vtkPolyDataMapper()
    outlineMapper.SetInput(outline.GetOutput())

    # make its actor
    outlineActor = vtk.vtkActor()
    outlineActor.SetMapper(outlineMapper)
    outlineActor.GetProperty().SetColor(0,0,0)

    # set up the renderer and render window
    ren = vtk.vtkRenderer()
    renWin = vtk.vtkRenderWindow()
    
    renWin.SetSize(800,600)
    renWin.AddRenderer(ren)
    ren.SetBackground(1,1,1)
    
    # add the relevant actors
    ren.AddActor(actor)
    ren.AddActor(titleActor)
    ren.AddActor(counterActor)
    ren.AddActor(scalarBar)
    ren.AddActor(outlineActor)
   
    cam = ren.GetActiveCamera()
    cam.Azimuth(0)
    cam.Elevation(-90)
    cam.Zoom(1.2)
    ren.SetActiveCamera(cam)
    ren.ResetCameraClippingRange()

    # add some axes
    axes = vtk.vtkCubeAxesActor2D()
    axes.SetInput(grid)
    axes.SetCamera(ren.GetActiveCamera())
    axes.SetLabelFormat("%6.4g")
    axes.SetFlyModeToOuterEdges()
    axes.SetFontFactor(0.8)
    axes.SetAxisTitleTextProperty(textProp)
    axes.SetAxisLabelTextProperty(textProp)
    axes.SetXLabel("x")
    axes.SetYLabel("y")
    axes.SetZLabel("z")
    axes.SetNumberOfLabels(5)
    axes.GetProperty().SetColor(0,0,0)
    ren.AddProp(axes)

    if interactive:
        # set up stuff for interactive viewing
        iren = vtk.vtkRenderWindowInteractor()
        iren.SetRenderWindow(renWin)
        
        iren.Initialize()
        renWin.Render()
        iren.Start()
    else:
        renWin.OffScreenRenderingOn()
        renWin.Render()

        # the WindowToImageFilter is what one uses to save the window to an 
        # image file
        win2img = vtk.vtkWindowToImageFilter()
        win2img.SetInput(renWin)
        
        # set up the PNMWriter as we're saving to pnm
        writer = vtk.vtkPNMWriter()
        writer.SetFileName("%s%04d.pnm" % (fnameStem,i))
        writer.SetInput(win2img.GetOutput())
        writer.Write()
    
        print "Wrote %s%04d.pnm" % (fnameStem,i)

# vim: expandtab shiftwidth=4:
1	jgs	123	#!/usr/bin/env python
2	ksteube	1312
3			#
4	jgs	123	# $Id$
5	ksteube	1312	#
6			#######################################################
7			#
8			# Copyright 2003-2007 by ACceSS MNRF
9			# Copyright 2007 by University of Queensland
10			#
11			# http://esscc.uq.edu.au
12			# Primary Business: Queensland, Australia
13			# Licensed under the Open Software License version 3.0
14			# http://www.opensource.org/licenses/osl-3.0.php
15			#
16			#######################################################
17			#
18	jgs	123
19			import vtk
20			import sys, os, re
21			import getopt
22
23			(opts, args) = getopt.getopt(sys.argv[1:],
24			"d:f:ph",
25			["dirname=", "fnamestem=", "preproc", "help"],
26			)
27
28			def usage():
29			print "Usage:"
30			print " velVis.py -d <dirname> -f <fnamestem> -p -i\n"
31			print " Arguments:"
32			print " -d/--dirname: directory of vtk files to process (req)"
33			print " -f/--fnamestem: filename stem of vtk files to process (req)"
34			print " -p/--preproc: preprocess vtk files to find global vector max? (opt)"
35			print " -i/--interactive: interactive use of first frame? (opt)"
36
37			#dirName = '/data/raid2/matt/results/050516/m3dpData/vis'
38			dirName = None
39			fnameStem = None
40			interactive = False
41			preproc = False # preprocess to find the global max of the data
42
43			for option, arg in opts:
44			if option in ('-d', '--dirname'):
45			dirName = arg
46			elif option in ('-f', '--fnamestem'):
47			fnameStem = arg
48			elif option in ('-p', '--preproc'):
49			preproc = True
50			elif option in ('-i', '--interactive'):
51			interactive = True
52			elif option in ('-h', '--help'):
53			usage()
54			sys.exit(0)
55
56			if dirName is None:
57			print "You must supply a directory of vtk files to process\n"
58			usage()
59			sys.exit(1)
60
61			if fnameStem is None:
62			print "You must supply the filename stem of the files to process"
63			print "The filename stem is the part of the filename without the"
64			print "frame number or extension\n"
65			usage()
66			sys.exit(1)
67
68			# now need to determine the number of files to process
69			dirList = os.listdir(dirName)
70			r = re.compile( "%s\\d+\\.\\w+" % fnameStem )
71			count = 0 # counter for the number of files found
72			fnames = []
73			for fname in dirList:
74			if r.match(fname):
75			fnames.append(fname)
76			count += 1
77
78			# determine the first frame number and number of frames
79			fnames.sort()
80			firstFile = fnames[0]
81
82			r = re.compile("([a-zA-Z-])(\\d+)(\\.)(\\w+)")
83			firstNum = r.findall(firstFile)
84			firstNum = int(firstNum[0][1])
85
86			# grab the filename extension as well, just in case
87			extn = r.findall(firstFile)
88			extn = extn[0][3]
89
90			# the number of frames should be equal to the count
91			if interactive:
92			numFrames = 1
93			mesa = False
94			else:
95			numFrames = count
96			mesa = True
97
98			# try and work out the format of the zero-padded frame number
99			if count < 10:
100			formatString = '"%01d"'
101			elif count < 100:
102			formatString = '"%02d"'
103			elif count < 1000:
104			formatString = '"%03d"'
105			elif count < 10000:
106			formatString = '"%04d"'
107			elif count < 100000:
108			formatString = '"%05d"'
109			else:
110			print "Sorry, can't handle greater than 100000 input frames"
111			sys.exit(1)
112
113			# use mesa if required
114			if mesa:
115			factGraphics = vtk.vtkGraphicsFactory()
116			factGraphics.SetUseMesaClasses(1)
117
118			factImage = vtk.vtkImagingFactory()
119			factImage.SetUseMesaClasses(1)
120
121			# perform some preprocessing to find the global max of the vector data norm
122			maxNorm = 0
123			if not interactive and preproc:
124			print "Preprocessing vector data"
125			sys.stdout.write("Completed: ")
126			for i in range(numFrames):
127			frameNum = eval(formatString) % i
128			reader = vtk.vtkXMLUnstructuredGridReader()
129			reader.SetFileName("%s/%s%s.%s" % (dirName,fnameStem,frameNum,extn))
130			reader.Update()
131
132			grid = reader.GetOutput()
133			norm = grid.GetPointData().GetVectors().GetMaxNorm()
134			if norm > maxNorm:
135			maxNorm = norm
136
137			sys.stdout.write("%3.0f%%\b\b\b\b" % (float(i)/float(numFrames), ))
138			print ""
139
140			for i in range(firstNum, numFrames+firstNum):
141
142			frameNum = eval(formatString) % i
143			reader = vtk.vtkXMLUnstructuredGridReader()
144			reader.SetFileName("%s/%s%s.%s" % (dirName,fnameStem,frameNum,extn))
145			reader.Update()
146
147			print "Processing %s" % fnames[i]
148
149			grid = reader.GetOutput()
150
151			# grab the model centre and bounds
152			centre = grid.GetCenter()
153			bounds = grid.GetBounds()
154
155			# grab the norm of the vectors
156			norm = vtk.vtkVectorNorm()
157			norm.SetInput(grid)
158
159			if not preproc:
160			maxNorm = grid.GetPointData().GetVectors().GetMaxNorm()
161
162			# to make arrow glyphs need an arrow source
163			arrow = vtk.vtkArrowSource()
164
165			# the arrows are 3D glyphs so set that up now
166			glyph = vtk.vtkGlyph3D()
167			glyph.ScalingOn()
168			glyph.SetScaleModeToScaleByScalar()
169			glyph.SetColorModeToColorByScalar()
170			glyph.SetVectorModeToUseVector()
171			glyph.SetScaleFactor(0.1/maxNorm)
172			glyph.SetInput(norm.GetOutput())
173			glyph.SetSource(arrow.GetOutput())
174			glyph.ClampingOff()
175
176			# set up a stripper to speed up rendening
177			stripper = vtk.vtkStripper()
178			stripper.SetInput(glyph.GetOutput())
179
180			# make a lookup table for the colour map and invert it (colours look
181			# better when it's inverted)
182			lut = vtk.vtkLookupTable()
183			refLut = vtk.vtkLookupTable()
184			lut.Build()
185			refLut.Build()
186			for j in range(256):
187			lut.SetTableValue(j, refLut.GetTableValue(255-j))
188
189			# set up the mapper
190			mapper = vtk.vtkPolyDataMapper()
191			mapper.SetInput(stripper.GetOutput())
192			mapper.SetScalarRange(0,maxNorm)
193			mapper.SetLookupTable(lut)
194
195			# set up the actor
196			actor = vtk.vtkActor()
197			actor.SetMapper(mapper)
198
199			# set up the text properties for nice text
200			font_size = 20
201			textProp = vtk.vtkTextProperty()
202			textProp.SetFontSize(font_size)
203			textProp.SetFontFamilyToArial()
204			textProp.BoldOff()
205			textProp.ItalicOff()
206			textProp.ShadowOff()
207			textProp.SetColor(0.0, 0.0, 0.0)
208
209			# make a title
210			title = vtk.vtkTextMapper()
211			title.SetInput("Velocity flow vectors in mantle convection")
212
213			# make the title text use the text properties
214			titleProp = title.GetTextProperty()
215			titleProp.ShallowCopy(textProp)
216			titleProp.SetJustificationToCentered()
217			titleProp.SetVerticalJustificationToTop()
218			titleProp.BoldOn()
219
220			# make the actor for the title
221			titleActor = vtk.vtkTextActor()
222			titleActor.SetMapper(title)
223			titleActor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay()
224			titleActor.GetPositionCoordinate().SetValue(0.5, 0.95)
225
226			# make a frame counter
227			counter = vtk.vtkTextMapper()
228			counter.SetInput("frame: %d" % i)
229
230			# make the counter use the text properties
231			counterProp = counter.GetTextProperty()
232			counterProp.ShallowCopy(textProp)
233			counterProp.SetJustificationToLeft()
234			counterProp.SetVerticalJustificationToTop()
235			counterProp.SetFontSize(14)
236
237			# make the actor for the frame counter
238			counterActor = vtk.vtkTextActor()
239			counterActor.SetMapper(counter)
240			counterActor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay()
241			counterActor.GetPositionCoordinate().SetValue(0.05, 0.8)
242
243			# use a scalar bar
244			scalarBar = vtk.vtkScalarBarActor()
245			scalarBar.SetLookupTable(lut)
246			scalarBar.SetWidth(0.1)
247			scalarBar.SetHeight(0.7)
248			scalarBar.SetPosition(0.9, 0.2)
249			scalarBar.SetTitle("\|v\|")
250
251			# set up its title text properties
252			scalarBarProp = scalarBar.GetTitleTextProperty()
253			scalarBarProp.ShallowCopy(textProp)
254			scalarBarProp.SetFontSize(10)
255
256			# set up the label text properties
257			scalarBarTextProp = scalarBar.GetLabelTextProperty()
258			scalarBarTextProp.ShallowCopy(textProp)
259			scalarBarTextProp.SetFontSize(10)
260
261			# put an outline around the data
262			outline = vtk.vtkOutlineSource()
263			outline.SetBounds(bounds)
264
265			# make its mapper
266			outlineMapper = vtk.vtkPolyDataMapper()
267			outlineMapper.SetInput(outline.GetOutput())
268
269			# make its actor
270			outlineActor = vtk.vtkActor()
271			outlineActor.SetMapper(outlineMapper)
272			outlineActor.GetProperty().SetColor(0,0,0)
273
274			# set up the renderer and render window
275			ren = vtk.vtkRenderer()
276			renWin = vtk.vtkRenderWindow()
277
278			renWin.SetSize(800,600)
279			renWin.AddRenderer(ren)
280			ren.SetBackground(1,1,1)
281
282			# add the relevant actors
283			ren.AddActor(actor)
284			ren.AddActor(titleActor)
285			ren.AddActor(counterActor)
286			ren.AddActor(scalarBar)
287			ren.AddActor(outlineActor)
288
289			cam = ren.GetActiveCamera()
290			cam.Azimuth(0)
291			cam.Elevation(-90)
292			cam.Zoom(1.2)
293			ren.SetActiveCamera(cam)
294			ren.ResetCameraClippingRange()
295
296			# add some axes
297			axes = vtk.vtkCubeAxesActor2D()
298			axes.SetInput(grid)
299			axes.SetCamera(ren.GetActiveCamera())
300			axes.SetLabelFormat("%6.4g")
301			axes.SetFlyModeToOuterEdges()
302			axes.SetFontFactor(0.8)
303			axes.SetAxisTitleTextProperty(textProp)
304			axes.SetAxisLabelTextProperty(textProp)
305			axes.SetXLabel("x")
306			axes.SetYLabel("y")
307			axes.SetZLabel("z")
308			axes.SetNumberOfLabels(5)
309			axes.GetProperty().SetColor(0,0,0)
310			ren.AddProp(axes)
311
312			if interactive:
313			# set up stuff for interactive viewing
314			iren = vtk.vtkRenderWindowInteractor()
315			iren.SetRenderWindow(renWin)
316
317			iren.Initialize()
318			renWin.Render()
319			iren.Start()
320			else:
321			renWin.OffScreenRenderingOn()
322			renWin.Render()
323
324			# the WindowToImageFilter is what one uses to save the window to an
325			# image file
326			win2img = vtk.vtkWindowToImageFilter()
327			win2img.SetInput(renWin)
328
329			# set up the PNMWriter as we're saving to pnm
330			writer = vtk.vtkPNMWriter()
331			writer.SetFileName("%s%04d.pnm" % (fnameStem,i))
332			writer.SetInput(win2img.GetOutput())
333			writer.Write()
334
335			print "Wrote %s%04d.pnm" % (fnameStem,i)
336
337			# vim: expandtab shiftwidth=4:
Name	Value
svn:eol-style	native
svn:keywords	Author Date Id Revision