Commit 1d760de4 authored by David E. DeMarle's avatar David E. DeMarle Committed by Kitware Robot

Merge topic 'optimize-filewrites'

0f0191ee use the new asynch writer in VTK to speed up writes
Acked-by: Kitware Robot's avatarKitware Robot <kwrobot@kitware.com>
Merge-request: !118
parents d9f40ab3 0f0191ee
Pipeline #100710 failed with stage
in 0 seconds
......@@ -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