Commit 0f0191ee authored by Patricia Fasel's avatar Patricia Fasel Committed by David E. DeMarle

use the new asynch writer in VTK to speed up writes

We see close to 2x speed up in ParaView with this. Before writing
superquadric took about 26 seconds, now it takes about 13.
parent d9f40ab3
......@@ -30,8 +30,6 @@ vtkEnabled = False
try:
import sys
if "paraview" in sys.modules:
# import paraview.vtk
# import paraview.vtk.vtkIOImage
from paraview.vtk.vtkIOImage import (vtkPNGReader,
vtkBMPReader,
vtkPNMReader,
......@@ -42,6 +40,7 @@ try:
vtkPNMWriter,
vtkTIFFWriter,
vtkJPEGWriter)
from paraview.vtk.vtkIOAsynchronous import vtkThreadedImageWriter
from paraview.vtk.vtkCommonDataModel import vtkImageData
from paraview import numpy_support as n2v
else:
......@@ -56,6 +55,7 @@ try:
vtkTIFFWriter,
vtkJPEGWriter,
vtkImageData)
from vtk.vtkIOAsynchronous import vtkThreadedImageWriter
from vtk.util import numpy_support as n2v
vtkEnabled = True
except ImportError:
......@@ -69,14 +69,19 @@ class RasterWrangler(object):
The choice of which is open to facilitate packaging in different
places, i.e. PIL for desktop and small packages, VTK for HPC contexts.
"""
def __init__(self):
self.threadedwriter = None
self.backends = set()
if exrEnabled:
self.backends.add("OpenEXR")
elif pilEnabled:
self.backends.add("PIL")
elif vtkEnabled:
self.backends.add("VTK")
self.threadedwriter = vtkThreadedImageWriter()
self.threadedwriter.Initialize()
elif pilEnabled:
self.backends.add("PIL")
# self.dontCompressFloatVals = False # don't expect we'll need this
# self.dontConvertValsToFloat = False # nor this
......@@ -160,8 +165,18 @@ class RasterWrangler(object):
def genericwriter(self, imageslice, fname):
"""write generic binary data dump"""
with open(fname, "w") as file:
file.write(imageslice)
if self.threadedwriter is not None:
height = imageslice.shape[1]
width = imageslice.shape[0]
contig = imageslice.reshape(height*width)
vtkarray = n2v.numpy_to_vtk(contig)
id = vtkImageData()
id.SetExtent(0, height-1, 0, width-1, 0, 0)
id.GetPointData().SetScalars(vtkarray)
self.threadedwriter.EncodeAndWrite(id, fname)
else:
with open(fname, "w") as file:
file.write(imageslice)
def rgbreader(self, fname):
"""opens a color image file and returns it as a color buffer"""
......@@ -202,10 +217,13 @@ class RasterWrangler(object):
id.SetExtent(0, height-1, 0, width-1, 0, 0)
id.GetPointData().SetScalars(vtkarray)
writer = self._make_writer(fname)
writer.SetInputData(id)
writer.SetFileName(fname)
writer.Write()
if self.threadedwriter is not None:
self.threadedwriter.EncodeAndWrite(id, fname)
else:
writer = self._make_writer(fname)
writer.SetInputData(id)
writer.SetFileName(fname)
writer.Write()
elif "PIL" in self.backends:
imageslice = numpy.flipud(imageslice)
......@@ -346,8 +364,19 @@ class RasterWrangler(object):
# Adjust the filename, replace .im with .npz
baseName, ext = os.path.splitext(fname)
adjustedName = baseName + ".Z"
with open(adjustedName, mode='wb') as file:
file.write(zlib.compress(numpy.array(imageslice)))
if self.threadedwriter is not None:
height = imageslice.shape[1]
width = imageslice.shape[0]
contig = imageslice.reshape(height*width)
vtkarray = n2v.numpy_to_vtk(contig)
id = vtkImageData()
id.SetExtent(0, height-1, 0, width-1, 0, 0)
id.GetPointData().SetScalars(vtkarray)
self.threadedwriter.EncodeAndWrite(id, adjustedName)
else:
with open(adjustedName, mode='wb') as file:
file.write(zlib.compress(numpy.array(imageslice)))
def assertvalidimage(self, filename):
"""tests that a given file is syntactically correct"""
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment