def add_base(index, mirror, B):
"""
build pipeline, for given basis
"""
B = util.normalize(B.T).T
PP = np.dot(B, self.complex.geometry.decomposed.T).T #primal coords transformed into current frame
x,y,z = PP.T
FV = self.complex.topology.FV[:,::np.sign(np.linalg.det(B))] #reverse vertex order depending on orientation
source = self.scene.mlab.pipeline.triangular_mesh_source(x,y,z, FV)
#put these guys under a ui list
l=lut_manager.tvtk.LookupTable()
lut_manager.set_lut(l, lut_manager.pylab_luts['jet'])
## lut_manager.set_lut(l, lut_manager.pylab_luts.values()[0])
#add polydatamapper, to control color mapping and interpolation
mapper = tvtk.PolyDataMapper(lookup_table=l)
from tvtk.common import configure_input
configure_input(mapper, source.outputs[0])
## mapper = tvtk.PolyDataMapper(input=source.outputs[0])
mapper.interpolate_scalars_before_mapping = True
mapper.immediate_mode_rendering = False
return mirror, source, mapper
def pointcloud_as_vtk_actor(points, pt_colors=None, point_size=5.0, alpha=1.0):
pd = pointcloud_as_vtk_polydata(points, pt_colors)
mapper = tvtk.PolyDataMapper()
configure_input(mapper, pd)
actor = tvtk.Actor(mapper=mapper)
actor.property.set(point_size=point_size, opacity=alpha)
return actor, pd
def marching_cubes(volume, smooth=True, decimate=True, **kwargs):
from tvtk.api import tvtk
from tvtk.common import configure_input
imgdata = tvtk.ImageData(dimensions=volume.shape)
imgdata.point_data.scalars = volume.flatten('F')
contours = tvtk.ContourFilter(number_of_contours=1)
contours.set_value(0, 1)
configure_input(contours, imgdata)
if smooth:
smoothargs = dict(number_of_iterations=40,
feature_angle=90,
pass_band=.05)
smoothargs.update(kwargs)
contours = tvtk.WindowedSincPolyDataFilter(input=contours.output,
**smoothargs)
if decimate:
contours = tvtk.QuadricDecimation(input=contours.output,
target_reduction=.75)
contours.update()
pts = contours.output.points.to_array()
polys = contours.output.polys.to_array().reshape(-1, 4)[:, 1:]
return pts, polys
def __init__(self, renwin, **traits):
super(Picker, self).__init__(**traits)
self.pointpicker = tvtk.PointPicker()
self.cellpicker = tvtk.CellPicker()
self.worldpicker = tvtk.WorldPointPicker()
self.probe_data = tvtk.PolyData()
# Use a set of axis to show the picked point.
self.p_source = tvtk.Axes()
self.p_mapper = tvtk.PolyDataMapper()
self.p_actor = tvtk.Actor()
self.p_source.symmetric = 1
self.p_actor.pickable = 0
self.p_actor.visibility = 0
prop = self.p_actor.property
prop.line_width = 2
prop.ambient = 1.0
prop.diffuse = 0.0
configure_input(self.p_mapper, self.p_source)
self.p_actor.mapper = self.p_mapper
self.probe_data.points = [[0.0, 0.0, 0.0]]
self.text_rep = tvtk.TextRepresentation()
self.text_widget = tvtk.TextWidget()
self.data = PickedData(renwin=renwin)
self.data.text_actor = tvtk.TextActor()
self.text_setup()
self.widgets = False
def _volume_mapper_type_changed(self, value):
if self.module_manager is None:
return
old_vm = self._volume_mapper
if old_vm is not None:
old_vm.on_trait_change(self.render, remove=True)
ray_cast_functions = ['']
if value == 'RayCastMapper':
new_vm = self._get_mapper(tvtk.VolumeRayCastMapper)
vrc_func = tvtk.VolumeRayCastCompositeFunction()
new_vm.volume_ray_cast_function = vrc_func
ray_cast_functions = ['RayCastCompositeFunction',
'RayCastMIPFunction',
'RayCastIsosurfaceFunction']
elif value == 'TextureMapper2D':
new_vm = self._get_mapper(tvtk.VolumeTextureMapper2D)
elif value == 'VolumeTextureMapper3D':
new_vm = self._get_mapper(tvtk.VolumeTextureMapper3D)
elif value == 'VolumeProMapper':
new_vm = self._get_mapper(tvtk.VolumeProMapper)
elif value == 'FixedPointVolumeRayCastMapper':
new_vm = self._get_mapper(tvtk.FixedPointVolumeRayCastMapper)
else:
msg = "Unsupported volume mapper type: '{0}'".format(value)
raise ValueError(msg)
self._volume_mapper = new_vm
self._ray_cast_functions = ray_cast_functions
configure_input(new_vm, self.module_manager.source.outputs[0])
self.volume.mapper = new_vm
new_vm.on_trait_change(self.render)
def __init__(self, renwin, **traits):
super(Picker, self).__init__(**traits)
self.renwin = renwin
self.pointpicker = tvtk.PointPicker()
self.cellpicker = tvtk.CellPicker()
self.worldpicker = tvtk.WorldPointPicker()
self.probe_data = tvtk.PolyData()
self._tolerance_changed(self.tolerance)
# Use a set of axis to show the picked point.
self.p_source = tvtk.Axes()
self.p_mapper = tvtk.PolyDataMapper()
self.p_actor = tvtk.Actor ()
self.p_source.symmetric = 1
self.p_actor.pickable = 0
self.p_actor.visibility = 0
prop = self.p_actor.property
prop.line_width = 2
prop.ambient = 1.0
prop.diffuse = 0.0
configure_input(self.p_mapper, self.p_source)
self.p_actor.mapper = self.p_mapper
self.probe_data.points = [[0.0, 0.0, 0.0]]
self.ui = None
def _oct_glyph(glyph_source, transform):
from tvtk.api import tvtk
from tvtk.common import configure_input
from traits.api import Array
gs = tvtk.PlatonicSolidSource()
# Workaround for:
# File "mayavi/components/glyph_source.py", line 231, in _glyph_position_changed # noqa: E501
# g.center = 0.0, 0.0, 0.0
# traits.trait_errors.TraitError: Cannot set the undefined 'center' attribute of a 'TransformPolyDataFilter' object. # noqa: E501
class SafeTransformPolyDataFilter(tvtk.TransformPolyDataFilter):
center = Array(shape=(3,), value=np.zeros(3))
gs.solid_type = 'octahedron'
if transform is not None:
# glyph: mayavi.modules.vectors.Vectors
# glyph.glyph: vtkGlyph3D
# glyph.glyph.glyph: mayavi.components.glyph.Glyph
assert transform.shape == (4, 4)
tr = tvtk.Transform()
tr.set_matrix(transform.ravel())
trp = SafeTransformPolyDataFilter()
configure_input(trp, gs)
trp.transform = tr
trp.update()
gs = trp
glyph_source.glyph_source = gs
def convert_to_poly_data(obj):
"""Given a VTK dataset object, this returns the data as PolyData.
This is primarily used to convert the data suitably for filters
that only work for PolyData.
"""
if obj.is_a('vtkDataSet'):
data = obj
else:
# FIXME
data = obj.output
if obj.is_a('vtkPolyData') or data.is_a('vtkPolyData'):
return obj
conv = {'vtkStructuredPoints': tvtk.ImageDataGeometryFilter,
'vtkImageData': tvtk.ImageDataGeometryFilter,
'vtkRectilinearGrid': tvtk.RectilinearGridGeometryFilter,
'vtkStructuredGrid': tvtk.StructuredGridGeometryFilter,
'vtkUnstructuredGrid':tvtk.GeometryFilter}
fil = None
for name, fil_class in conv.items():
if data.is_a(name):
fil = fil_class()
break
if fil is not None:
configure_input(fil, obj)
fil.update()
return fil
else:
error('Given object is not a VTK dataset: %s'%data.__class__.__name__)
def main(hdf5_animation_file):
weights = None
with h5py.File(hdf5_animation_file, 'r') as f:
verts = f['verts'].value
tris = f['tris'].value
if 'weights' in f:
weights = f['weights'].value
pd = tvtk.PolyData(points=verts[0], polys=tris)
normals = tvtk.PolyDataNormals(splitting=False)
configure_input_data(normals, pd)
actor = tvtk.Actor(mapper=tvtk.PolyDataMapper())
configure_input(actor.mapper, normals)
actor.property.set(edge_color=(0.5, 0.5, 0.5), ambient=0.0,
specular=0.15, specular_power=128., shading=True, diffuse=0.8)
fig = mlab.figure(bgcolor=(1,1,1))
fig.scene.add_actor(actor)
@mlab.animate(delay=40, ui=False)
def animation():
for i in count():
if weights is not None:
w_str = ",".join(["%0.2f"] * weights.shape[1])
print ("Frame %d Weights = " + w_str) % tuple([i] + weights[i].tolist())
frame = i % len(verts)
pd.points = verts[frame]
fig.scene.render()
yield
a = animation()
fig.scene.z_minus_view()
mlab.show()
def __init__(self, renwin, **traits):
super(Picker, self).__init__(**traits)
self.renwin = renwin
self.pointpicker = tvtk.PointPicker()
self.cellpicker = tvtk.CellPicker()
self.worldpicker = tvtk.WorldPointPicker()
self.probe_data = tvtk.PolyData()
self._tolerance_changed(self.tolerance)
# Use a set of axis to show the picked point.
self.p_source = tvtk.Axes()
self.p_mapper = tvtk.PolyDataMapper()
self.p_actor = tvtk.Actor()
self.p_source.symmetric = 1
self.p_actor.pickable = 0
self.p_actor.visibility = 0
prop = self.p_actor.property
prop.line_width = 2
prop.ambient = 1.0
prop.diffuse = 0.0
configure_input(self.p_mapper, self.p_source)
self.p_actor.mapper = self.p_mapper
self.probe_data.points = [[0.0, 0.0, 0.0]]
self.ui = None
def concave_hull(polydata):
"""extract the concave hull of a vtk polydata object"""
# setup the extraction pipeline
extract = tvtk.FeatureEdges()
# we're only interested in the boundaries
extract.feature_edges = False
extract.boundary_edges = True
extract.manifold_edges = False
extract.non_manifold_edges = False
configure_input(extract, polydata)
# compute edges
extract.update()
# extract the points
points = extract.output.points.to_array()
# slice every 2nd and 3rd point
line_ids = np.c_[
extract.output.lines.to_array()[1::3],
extract.output.lines.to_array()[2::3]
]
# 1st points should all be 2
assert (extract.output.lines.to_array()[0::3] == 2).all(), "expected only lines"
# construct a directed graph
D = networkx.DiGraph(data=line_ids.tolist())
# find the first cycle
first_cycle = networkx.find_cycle(D)
# create the index to lookup points, including last point
cycle = [x[0] for x in first_cycle] + [first_cycle[0][0]]
# fill in index and return first 2 coordinates
return points[cycle, :2]
def write_flux(file_name, surface, surface_density, surface_va, surface_beta,
surface_cs, Fpar, Fperp, Fphi):
pd_density = tvtk.PointData(scalars=surface_density)
pd_density.scalars.name = "surface_density"
pd_va = tvtk.PointData(scalars=surface_va)
pd_va.scalars.name = "surface_va"
pd_beta = tvtk.PointData(scalars=surface_beta)
pd_beta.scalars.name = "surface_beta"
pd_cs = tvtk.PointData(scalars=surface_cs)
pd_cs.scalars.name = "surface_cs"
pd_Fpar = tvtk.PointData(scalars=Fpar)
pd_Fpar.scalars.name = "Fpar"
pd_Fperp = tvtk.PointData(scalars=Fperp)
pd_Fperp.scalars.name = "Fperp"
pd_Fphi = tvtk.PointData(scalars=Fphi)
pd_Fphi.scalars.name = "Fphi"
poly_out = surface
poly_out.point_data.add_array(pd_density.scalars)
poly_out.point_data.add_array(pd_va.scalars)
poly_out.point_data.add_array(pd_beta.scalars)
poly_out.point_data.add_array(pd_cs.scalars)
poly_out.point_data.add_array(pd_Fpar.scalars)
poly_out.point_data.add_array(pd_Fperp.scalars)
poly_out.point_data.add_array(pd_Fphi.scalars)
w = tvtk.XMLPolyDataWriter(file_name=file_name)
tvtk_common.configure_input(w, poly_out)
w.write()
def save_png(self, file_name):
"""Save to a PNG image file."""
if len(file_name) != 0:
w2if = self._get_window_to_image()
ex = tvtk.PNGWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def save_png(self, file_name):
"""Save to a PNG image file."""
if len(file_name) != 0:
w2if = self._get_window_to_image()
ex = tvtk.PNGWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def _mask_input_points_changed(self, value):
inputs = self.inputs
if len(inputs) == 0:
return
if value:
mask = self.mask_points
tvtk_common.configure_input(mask, inputs[0].outputs[0])
else:
self.configure_connection(self.glyph, inputs[0])
def _mask_input_points_changed(self, value):
inputs = self.inputs
if len(inputs) == 0:
return
if value:
mask = self.mask_points
tvtk_common.configure_input(mask, inputs[0].outputs[0])
else:
self.configure_connection(self.glyph, inputs[0])
def save_ps(self, file_name):
"""Saves the rendered scene to a rasterized PostScript image.
For vector graphics use the save_gl2ps method."""
if len(file_name) != 0:
w2if = self._get_window_to_image()
ex = tvtk.PostScriptWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def make_outline(input_obj):
from tvtk.api import tvtk
from tvtk.common import configure_input
outline = tvtk.StructuredGridOutlineFilter()
configure_input(outline, input_obj)
outline_mapper = tvtk.PolyDataMapper(input_connection = outline.output_port)
outline_actor = tvtk.Actor(mapper = outline_mapper)
outline_actor.property.color = 0.3, 0.3, 0.3
return outline_actor
def save_ps(self, file_name):
"""Saves the rendered scene to a rasterized PostScript image.
For vector graphics use the save_gl2ps method."""
if len(file_name) != 0:
w2if = self._get_window_to_image()
ex = tvtk.PostScriptWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def add_actors_to_scene(self, scene_model, volume_actor):
# An outline of the bounds of the Volume actor's data
outline = tvtk.OutlineFilter()
configure_input(outline, volume_actor.mapper.input)
outline_mapper = tvtk.PolyDataMapper()
configure_input(outline_mapper, outline.output)
outline_actor = tvtk.Actor(mapper=outline_mapper)
outline_actor.property.opacity = 0.3
scene_model.renderer.add_actor(outline_actor)
def polydata_actor(polydata, compute_normals=True):
""" create a vtk actor with given polydata as input """
if compute_normals:
normals = tvtk.PolyDataNormals(splitting=False)
configure_input_data(normals, polydata)
polydata = normals
actor = tvtk.Actor(mapper=tvtk.PolyDataMapper())
configure_input(actor.mapper, polydata)
actor.mapper.lookup_table.hue_range = (0.33, 0.)
return actor
def _mask_input_points_changed(self, value):
from tvtk.common import configure_input
inputs = self.inputs
if len(inputs) == 0:
return
if value:
mask = self.mask_points
configure_input(mask, inputs[0].outputs[0])
self.configure_connection(self.glyph, mask)
else:
self.configure_connection(self.glyph, inputs[0])
def get_line(a, b):
src = tvtk.LineSource(point1=a, point2=b)
if new_tvtk:
mapper = tvtk.PolyDataMapper()
configure_input(mapper, src)
else:
mapper = tvtk.PolyDataMapper(input=src.output)
actor = tvtk.Actor(mapper=mapper)
fig.scene.add_actor(actor)
return actor
def scene_to_png(scene):
w2if = tvtk.WindowToImageFilter()
w2if.input = scene.render_window
ex = tvtk.PNGWriter()
ex.write_to_memory = True
configure_input(ex, w2if)
ex.update()
ex.write()
data = base64.b64encode(ex.result.to_array()).decode("ascii")
html = '<img src="data:image/png;base64,%s" alt="PNG image"></img>'
return html % data
def scene_to_png(scene):
w2if = tvtk.WindowToImageFilter()
w2if.input = scene.render_window
ex = tvtk.PNGWriter()
ex.write_to_memory = True
configure_input(ex, w2if)
ex.update()
ex.write()
data = base64.b64encode(ex.result.to_array()).decode('ascii')
html = '<img src="data:image/png;base64,%s" alt="PNG image"></img>'
return html % data
def save_png(self, file_name):
"""Save to a PNG image file."""
if len(file_name) != 0:
w2if = tvtk.WindowToImageFilter(read_front_buffer=not self.off_screen_rendering)
w2if.magnification = self.magnification
self._lift()
w2if.input = self._renwin
ex = tvtk.PNGWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def _mask_input_points_changed(self, value):
from tvtk.common import configure_input
inputs = self.inputs
if len(inputs) == 0:
return
if value:
mask = self.mask_points
configure_input(mask, inputs[0].outputs[0])
self.configure_connection(self.glyph, mask)
else:
self.configure_connection(self.glyph, inputs[0])
def __init__(self, verts1, verts2, tris=None, lines=None, as_points=False,
scalars=None, scalars2=None, vmin=None, vmax=None,
actor_property=dict(specular=0.1, specular_power=128., diffuse=0.5),
):
if tris is None:
if lines is None:
rep = 'points'
else:
rep = 'wireframe'
else:
rep = 'surface'
super(Morpher, self).__init__()
self._verts1, self._verts2 = verts1, verts2
self._polydata = tvtk.PolyData(points=verts1)
if rep == 'points':
self._polydata.verts = np.r_[:len(verts1)].reshape(-1,1)
if tris is not None:
self._polydata.polys = tris
if lines is not None:
self._polydata.lines = lines
n = tvtk.PolyDataNormals(splitting=False)
configure_input_data(n, self._polydata)
self._actor = tvtk.Actor(mapper=tvtk.PolyDataMapper())
configure_input(self._actor.mapper, n)
self._actor.property.representation = rep
if rep == 'points':
self._actor.property.point_size = 5
if as_points and scalars is None:
self._polydata.point_data.scalars = \
np.random.uniform(0, 255, (len(verts1), 3)).astype(np.uint8)
self._scalars12 = None
if scalars is not None:
self._polydata.point_data.scalars = scalars
# automatically determine minimum/maximum from scalars if not given by user
if vmin is None:
vmin = scalars.min()
if scalars2 is not None:
vmin = min(vmin, scalars2.min())
if vmax is None:
vmax = scalars.max()
if scalars2 is not None:
vmax = max(vmax, scalars2.max())
if scalars.ndim == 1:
self._actor.mapper.use_lookup_table_scalar_range = False
self._actor.mapper.scalar_range = (vmin, vmax)
self._actor.mapper.lookup_table.hue_range = (0.33, 0)
# when scalars of second mesh given we need to store both scalars in order
# to interpolate between them during rendering
if scalars2 is not None:
self._scalars12 = (scalars, scalars2)
else:
self._actor.property.set(**actor_property)
mlab.gcf().scene.add_actor(self._actor)
def save_png(self, file_name):
"""Save to a PNG image file."""
if len(file_name) != 0:
w2if = tvtk.WindowToImageFilter(
read_front_buffer=not self.off_screen_rendering)
w2if.magnification = self.magnification
self._lift()
w2if.input = self._renwin
ex = tvtk.PNGWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def save_ps(self, file_name):
"""Saves the rendered scene to a rasterized PostScript image.
For vector graphics use the save_gl2ps method."""
if len(file_name) != 0:
w2if = tvtk.WindowToImageFilter(read_front_buffer=not self.off_screen_rendering)
w2if.magnification = self.magnification
self._lift()
w2if.input = self._renwin
ex = tvtk.PostScriptWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def save_ps(self, file_name):
"""Saves the rendered scene to a rasterized PostScript image.
For vector graphics use the save_gl2ps method."""
if len(file_name) != 0:
w2if = tvtk.WindowToImageFilter(
read_front_buffer=not self.off_screen_rendering)
w2if.magnification = self.magnification
self._lift()
w2if.input = self._renwin
ex = tvtk.PostScriptWriter()
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def vismesh(pts,
tris,
color=None,
edge_visibility=False,
shader=None,
triangle_scalars=None,
colors=None,
nan_color=None,
**kwargs):
if 'scalars' in kwargs and np.asarray(kwargs['scalars']).ndim == 2:
colors = kwargs['scalars']
del kwargs['scalars']
# VTK does not allow bool arrays as scalars normally, so convert to float
if 'scalars' in kwargs and np.asarray(kwargs['scalars']).dtype == np.bool:
kwargs['scalars'] = kwargs['scalars'].astype(np.float)
tm = mlab.triangular_mesh(pts[:, 0],
pts[:, 1],
pts[:, 2],
tris,
color=color,
**kwargs)
if shader is not None:
tm.actor.property.load_material(shader)
tm.actor.actor.property.shading = True
diffuse = 1.0 if colors is not None else 0.8
tm.actor.actor.property.set(edge_visibility=edge_visibility,
line_width=1,
specular=0.0,
specular_power=128.,
diffuse=diffuse)
if triangle_scalars is not None:
tm.actor.mapper.input.cell_data.scalars = triangle_scalars
tm.actor.mapper.set(scalar_mode='use_cell_data',
use_lookup_table_scalar_range=False,
scalar_visibility=True)
if "vmin" in kwargs and "vmax" in kwargs:
tm.actor.mapper.scalar_range = kwargs["vmin"], kwargs["vmax"]
if colors is not None:
# this basically is a hack which doesn't quite work,
# we have to completely replace the polydata behind the hands of mayavi
tm.mlab_source.dataset.point_data.scalars = colors.astype(np.uint8)
normals = tvtk.PolyDataNormals(splitting=False)
configure_input_data(normals, tm.mlab_source.dataset)
configure_input(tm.actor.mapper, normals)
if nan_color is not None:
if len(nan_color) == 3:
nan_color = list(nan_color) + [1]
tm.module_manager.scalar_lut_manager.lut.nan_color = nan_color
tm.update_pipeline()
return tm
def _poly_data_to_volume(data, voxel_size, margins):
"""
Render the given, closed <tvtk.PolyData> into an image volume and return it
as a Numpy array.
"""
# Following [*] for the necessary steps.
#
# References
# [*] http://www.paraview.org/Wiki/VTK/Examples/Cxx/PolyData/PolyDataToImageData (20150710)
# Calculate the necessary number of voxels, extent, and origin
bounds = np.array(data.bounds)
n_vox = np.ceil(
(bounds[1::2] - bounds[0::2] + np.sum(margins.reshape(2, 3), axis=0)) /
voxel_size).astype(np.int)
extent = [0, n_vox[0] - 1, 0, n_vox[1] - 1, 0, n_vox[2] - 1]
origin = bounds[0::2] - margins[:3] + 0.5 * voxel_size
# Create the image volume
volume = tvtk.ImageData(spacing=voxel_size,
dimensions=n_vox,
extent=extent,
origin=origin)
if is_old_pipeline():
volume.scalar_type = "unsigned_char"
voxel_data = (np.ones(n_vox, dtype=np.uint8) * 100).ravel()
volume.point_data.scalars = voxel_data
if is_old_pipeline():
volume.update()
# Actual transformation from polygon to image data
voxelizer = tvtk.PolyDataToImageStencil(output_origin=origin,
output_spacing=voxel_size,
output_whole_extent=volume.extent,
tolerance=0.0)
configure_input(voxelizer, data)
voxelizer.update()
# Draw the result to a new image, extract Numpy array and return
painter = tvtk.ImageStencil(reverse_stencil=False, background_value=0)
configure_input(painter, volume)
if is_old_pipeline():
painter.stencil = voxelizer.output
else:
painter.set_stencil_connection(voxelizer.output_port)
painter.update()
voxel_data = painter.output.point_data.scalars.to_array().reshape(
n_vox[::-1])
voxel_data = np.transpose(voxel_data, [2, 1, 0])
return voxel_data
def save_jpg(self, file_name, quality=None, progressive=None):
"""Arguments: file_name if passed will be used, quality is the
quality of the JPEG(10-100) are valid, the progressive
arguments toggles progressive jpegs."""
if len(file_name) != 0:
if not quality and not progressive:
quality, progressive = self.jpeg_quality, self.jpeg_progressive
w2if = self._get_window_to_image()
ex = tvtk.JPEGWriter()
ex.quality = quality
ex.progressive = progressive
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def save_jpg(self, file_name, quality=None, progressive=None):
"""Arguments: file_name if passed will be used, quality is the
quality of the JPEG(10-100) are valid, the progressive
arguments toggles progressive jpegs."""
if len(file_name) != 0:
if not quality and not progressive:
quality, progressive = self.jpeg_quality, self.jpeg_progressive
w2if = self._get_window_to_image()
ex = tvtk.JPEGWriter()
ex.quality = quality
ex.progressive = progressive
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def save_jpg(self, file_name, quality=None, progressive=None):
"""Arguments: file_name if passed will be used, quality is the
quality of the JPEG(10-100) are valid, the progressive
arguments toggles progressive jpegs."""
if len(file_name) != 0:
if not quality and not progressive:
quality, progressive = self.jpeg_quality, self.jpeg_progressive
w2if = tvtk.WindowToImageFilter(read_front_buffer=not self.off_screen_rendering)
w2if.magnification = self.magnification
self._lift()
w2if.input = self._renwin
ex = tvtk.JPEGWriter()
ex.quality = quality
ex.progressive = progressive
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def create_flux_surface(bfield, surf_seeds):
"""
Create a flux surface from an array of seeds and a tvtk vector field.
Parameters
----------
bfield: tvtk.ImageData
The vector field to use for streamline traceing
surf_seeds: numpy.ndarray
The array of seed points to start the fieldline tracing from
Returns
-------
surf_field_lines: tvtk.StreamTracer instance
The fieldline tracer with the fieldlines stored inside it.
surface: tvtk.RuledSurfaceFilter instance
The surface built from the StreamTracer instance
"""
#Make a streamline instance with the bfield
surf_field_lines = tvtk.StreamTracer()
# surf_field_lines.input_connection = bfield
tvtk_common.configure_input(surf_field_lines, bfield)
tvtk_common.configure_source_data(surf_field_lines, surf_seeds)
# surf_field_lines.source = surf_seeds
surf_field_lines.integrator = tvtk.RungeKutta4()
surf_field_lines.maximum_propagation = 1000
surf_field_lines.integration_direction = 'backward'
surf_field_lines.update()
#Create surface from 'parallel' lines
surface = tvtk.RuledSurfaceFilter()
tvtk_common.configure_connection(surface, surf_field_lines)
# surface.input = surf_field_lines.output
surface.close_surface = True
surface.pass_lines = True
surface.offset = 0
surface.distance_factor = 30
surface.ruled_mode = 'point_walk'
# surface.ruled_mode = 'resample'
# surface.resolution = (10,1)
surface.update()
return surf_field_lines, surface
def save(self, i):
self.mask_nonfluid_nodes()
os.environ['ETS_TOOLKIT'] = 'null'
from tvtk.api import tvtk
from tvtk.common import configure_input
idata = tvtk.ImageData(spacing=(1, 1, 1), origin=(0, 0, 0))
first = True
sample_field = None
for name, field in self._scalar_fields.items():
if first:
idata.point_data.scalars = field.flatten()
idata.point_data.scalars.name = name
first = False
sample_field = field
else:
t = idata.point_data.add_array(field.flatten())
idata.point_data.get_array(t).name = name
idata.update_traits()
dim = len(sample_field.shape)
for name, field in self._vector_fields.items():
if dim == 3:
tmp = idata.point_data.add_array(np.c_[field[0].flatten(),
field[1].flatten(),
field[2].flatten()])
else:
tmp = idata.point_data.add_array(
np.c_[field[0].flatten(), field[1].flatten(),
np.zeros_like(field[0].flatten())])
idata.point_data.get_array(tmp).name = name
if dim == 3:
idata.dimensions = list(reversed(sample_field.shape))
else:
idata.dimensions = list(reversed(sample_field.shape)) + [1]
fname = filename(self.basename,
self.digits,
self.subdomain_id,
i,
suffix='.vti')
w = tvtk.XMLImageDataWriter(file_name=fname)
configure_input(w, idata)
w.write()
def get_sphere(p, radius, res=8):
if type(p) != tuple:
p = tuple(p)
src = tvtk.SphereSource(center=p,
radius=radius,
phi_resolution=res,
theta_resolution=res)
if new_tvtk:
mapper = tvtk.PolyDataMapper()
configure_input(mapper, src)
else:
mapper = tvtk.PolyDataMapper(input=src.output)
actor = tvtk.Actor(mapper=mapper)
fig.scene.add_actor(actor)
return actor
def __init__(self, Xmean, tris, components):
HasTraits.__init__(self)
self._components = components
self._max_component_index = len(components)
self._Xmean = Xmean
self.pd = tvtk.PolyData(points=Xmean, polys=tris)
self.normals = tvtk.PolyDataNormals(splitting=False)
configure_input_data(self.normals, self.pd)
mapper = tvtk.PolyDataMapper(immediate_mode_rendering=True)
self.actor = tvtk.Actor(mapper=mapper)
configure_input(self.actor.mapper, self.normals)
self.actor.mapper.lookup_table = tvtk.LookupTable(
hue_range=(0.45, 0.6),
saturation_range=(0., 0.8),
value_range=(.6, 1.),
)
self.scene.add_actor(self.actor)
self.timer = Timer(40, self.animate().next)
def save_jpg(self, file_name, quality=None, progressive=None):
"""Arguments: file_name if passed will be used, quality is the
quality of the JPEG(10-100) are valid, the progressive
arguments toggles progressive jpegs."""
if len(file_name) != 0:
if not quality and not progressive:
quality, progressive = self.jpeg_quality, self.jpeg_progressive
w2if = tvtk.WindowToImageFilter(
read_front_buffer=not self.off_screen_rendering)
w2if.magnification = self.magnification
self._lift()
w2if.input = self._renwin
ex = tvtk.JPEGWriter()
ex.quality = quality
ex.progressive = progressive
ex.file_name = file_name
configure_input(ex, w2if)
self._exporter_write(ex)
def __init__(self, Xmean, tris, components):
HasTraits.__init__(self)
self._components = components
self._max_component_index = len(components)
self._Xmean = Xmean
self.pd = tvtk.PolyData(points=Xmean, polys=tris)
self.normals = tvtk.PolyDataNormals(splitting=False)
configure_input_data(self.normals, self.pd)
mapper = tvtk.PolyDataMapper(immediate_mode_rendering=True)
self.actor = tvtk.Actor(mapper=mapper)
configure_input(self.actor.mapper, self.normals)
self.actor.mapper.lookup_table = tvtk.LookupTable(
hue_range = (0.45, 0.6),
saturation_range = (0., 0.8),
value_range = (.6, 1.),
)
self.scene.add_actor(self.actor)
self.timer = Timer(40, self.animate().next)
def setup_pipeline(polydata, width=1000, height=1000):
""" VTk pipeline to resample data """
resample = tvtk.ResampleToImage()
cell2point = tvtk.CellDataToPointData()
# setup the pipeline
configure_input(cell2point, polydata)
configure_input(resample, cell2point)
resample.sampling_dimensions = np.array([width, height, 1])
resample.sampling_bounds = np.array([-180, 180, -90, 90, 0, 0.0])
resample.use_input_bounds = False
# TODO: check if this can be left out
resample.modified()
resample.update()
polydata.modified()
return resample
def test():
for module_name in [
'numpy', 'scipy', 'pandas', 'statsmodels', 'pytables', 'xarray',
'dask', 'PIL', 'sqlalchemy', 'psycopg2', 'pymongo', 'rasterio',
'shapely', 'fiona', 'rasterio', 'geopandas', 'mercantile',
'pyproj', 'hdf4', 'hdf5', 'h5py', 'netcdf4', 'matplotlib',
'seaborn', 'bokeh', 'holoviews', 'geoviews', 'notebook', 'flask',
'marshmallow', 'theano', 'tensorflow', 'keras'
]:
print("Importing {}:".format(module_name))
module = None
try:
module = importlib.import_module(module_name)
except (ImportError, AttributeError):
try:
module = importlib.import_module(module_name)
except (ImportError, AttributeError):
print(" Import failed")
if module is None:
continue
try:
print(" version={}".format(module.__version__))
except AttributeError:
print(" No version available")
print("Importing gdal:")
from osgeo import gdal
print(" version={}".format(gdal.__version__))
print("Importing tvtk:")
from tvtk.api import tvtk
from tvtk.common import configure_input
cs = tvtk.ConeSource()
cs.resolution = 36
m = tvtk.PolyDataMapper()
configure_input(m, cs.output)
a = tvtk.Actor()
a.mapper = m
p = a.property
p.representation = 'w'
print(" OK")
def get_trunkcone(b, a):
phi_base, theta_base = sph.to(a, b)[1:]
quads = tvtk.CellArray() #vtk.vtkCellArray()
points = tvtk.Points() #vtk.vtkPoints()
Nface = 3
for i in range(Nface + 1):
# rotate
phi, theta = convdir((i % Nface) * 2 * pi / Nface, pi * 0.5, phi_base,
theta_base)
# generate new points
p = tuple(sph.xyz(a[3] * 0.5 * radius_factor, phi, theta, a[:3]))
q = tuple(sph.xyz(b[3] * 0.5 * radius_factor, phi, theta, b[:3]))
# insert points
points.append(p)
points.append(q)
if i >= 1:
# create a face
quad = tvtk.Quad()
n = points.number_of_points - 1
quad.point_ids.set_id(0, n - 3) # p
quad.point_ids.set_id(1, n - 2) # q
quad.point_ids.set_id(2, n) # q
quad.point_ids.set_id(3, n - 1) # p
# insert the new face
quads.insert_next_cell(quad)
# create the actor
polydata = tvtk.PolyData(points=points, polys=quads)
if new_tvtk:
mapper = tvtk.PolyDataMapper()
configure_input(mapper, polydata)
else:
mapper = tvtk.PolyDataMapper(input=polydata)
actor = tvtk.Actor(mapper=mapper)
fig.scene.add_actor(actor)
return actor
def get_cone(base, radius, v):
if type(base) != tuple:
base = tuple(base)
if type(v) != tuple:
v = tuple(v)
src = tvtk.ConeSource(center=base,
radius=radius * radius_factor,
height=radius,
direction=v,
resolution=20)
if new_tvtk:
mapper = tvtk.PolyDataMapper()
configure_input(mapper, src)
else:
mapper = tvtk.PolyDataMapper(input=src.output)
actor = tvtk.Actor(mapper=mapper)
fig.scene.add_actor(actor)
return actor
def marching_cubes(volume, smooth=True, decimate=True, **kwargs):
from tvtk.api import tvtk
from tvtk.common import configure_input
imgdata = tvtk.ImageData(dimensions=volume.shape)
imgdata.point_data.scalars = volume.flatten('F')
contours = tvtk.ContourFilter(number_of_contours=1)
contours.set_value(0, 1)
configure_input(contours, imgdata)
if smooth:
smoothargs = dict(number_of_iterations=40, feature_angle = 90, pass_band=.05)
smoothargs.update(kwargs)
contours = tvtk.WindowedSincPolyDataFilter(input=contours.output, **smoothargs)
if decimate:
contours = tvtk.QuadricDecimation(input=contours.output, target_reduction=.75)
contours.update()
pts = contours.output.points.to_array()
polys = contours.output.polys.to_array().reshape(-1, 4)[:,1:]
return pts, polys
def _decimate_surface(points, triangles, reduction):
"""Aux function."""
if 'DISPLAY' not in os.environ and sys.platform != 'win32':
os.environ['ETS_TOOLKIT'] = 'null'
try:
from tvtk.api import tvtk
from tvtk.common import configure_input
except ImportError:
raise ValueError('This function requires the TVTK package to be '
'installed')
if triangles.max() > len(points) - 1:
raise ValueError('The triangles refer to undefined points. '
'Please check your mesh.')
src = tvtk.PolyData(points=points, polys=triangles)
decimate = tvtk.QuadricDecimation(target_reduction=reduction)
configure_input(decimate, src)
decimate.update()
out = decimate.output
tris = out.polys.to_array()
# n-tuples + interleaved n-next -- reshape trick
return out.points.to_array(), tris.reshape(tris.size // 4, 4)[:, 1:]
def _decimate_surface(points, triangles, reduction):
"""Aux function."""
if 'DISPLAY' not in os.environ and sys.platform != 'win32':
os.environ['ETS_TOOLKIT'] = 'null'
try:
from tvtk.api import tvtk
from tvtk.common import configure_input
except ImportError:
raise ValueError('This function requires the TVTK package to be '
'installed')
if triangles.max() > len(points) - 1:
raise ValueError('The triangles refer to undefined points. '
'Please check your mesh.')
src = tvtk.PolyData(points=points, polys=triangles)
decimate = tvtk.QuadricDecimation(target_reduction=reduction)
configure_input(decimate, src)
decimate.update()
out = decimate.output
tris = out.polys.to_array()
# n-tuples + interleaved n-next -- reshape trick
return out.points.to_array(), tris.reshape(tris.size // 4, 4)[:, 1:]
def _render_animation_fired(self):
self.stop = True
n_frames_render = self.render_to_frame - self.render_from_frame
# prepare the render window
renwin = self._figure.scene.render_window
aa_frames = renwin.aa_frames
renwin.aa_frames = 8
renwin.alpha_bit_planes = 1
# turn on off screen rendering
#renwin.off_screen_rendering = True
# set size of window
if self.fix_image_size:
orig_size = renwin.size
renwin.size = self.image_size
# render the frames
progress = ProgressDialog(title="Rendering", max=n_frames_render,
show_time=True, can_cancel=True)
progress.open()
self.is_rendering_animation = True
for frame in range(self.render_from_frame, self.render_to_frame + 1):
# move animation to desired frame, this will also render the scene
self.current_frame = frame
# prepare window to image writer
render = tvtk.WindowToImageFilter(input=renwin, magnification=1)#, input_buffer_type='rgba')
if not self.fix_image_size:
render.magnification = self.magnification
exporter = tvtk.PNGWriter(file_name=path.join(self.render_directory, self.render_name_pattern % frame))
configure_input(exporter,render)
exporter.write()
do_continue, skip = progress.update(frame - self.render_from_frame)
if not do_continue:
break
# reset the render window to old values
renwin.aa_frames = aa_frames
if self.fix_image_size:
renwin.size = orig_size
#renwin.off_screen_rendering = False
self.is_rendering_animation = False
progress.close()
def write_wave_flux(file_name, surface_poly, parallels, normals, torsionals,
Fwpar, Fwperp, Fwphi):
pd_Fwpar = tvtk.PointData(scalars=Fwpar, vectors=parallels)
pd_Fwpar.scalars.name = "Fwpar"
pd_Fwpar.vectors.name = "par"
pd_Fwperp = tvtk.PointData(scalars=Fwperp, vectors=normals)
pd_Fwperp.scalars.name = "Fwperp"
pd_Fwperp.vectors.name = "perp"
pd_Fwphi = tvtk.PointData(scalars=Fwphi, vectors=torsionals)
pd_Fwphi.scalars.name = "Fwphi"
pd_Fwphi.vectors.name = "phi"
poly_out = surface_poly
poly_out.point_data.add_array(pd_Fwpar.scalars)
poly_out.point_data.add_array(pd_Fwperp.scalars)
poly_out.point_data.add_array(pd_Fwphi.scalars)
w = tvtk.XMLPolyDataWriter(file_name=file_name)
tvtk_common.configure_input(w, poly_out)
w.write()
def __init__(self):
self.el = 0.0
self.az = 0.0
# Create an arrow.
arrow = tvtk.ArrowSource()
# Transform it suitably so it is oriented correctly.
t = tvtk.Transform()
tf = tvtk.TransformFilter()
tf.transform = t
t.rotate_y(90.0)
t.translate((-2, 0, 0))
configure_input_data(tf, arrow.output)
mapper = tvtk.PolyDataMapper()
configure_input(mapper, tf)
self.actor = actor = tvtk.Follower()
actor.mapper = mapper
prop = actor.property
prop.color = 0, 1, 1
prop.ambient = 0.5
prop.diffuse = 0.5
def configure_input(self, inp, op):
""" Configure the inp using op."""
tvtk_common.configure_input(inp, op)
def render(self, filename=None, crop=True, add_actors=[]):
renderer = tvtk.Renderer(background = (1, 1, 1))
subset = tvtk.ExtractVOI(
sample_rate = (self.sample_rate, self.sample_rate, 1),
voi = (self._x_min, self._x_max, self._y_min, self._y_max, 0, 0))
configure_input(subset, self._dtm_reader)
cropped = tvtk.Threshold()
configure_input(cropped, subset)
cropped.threshold_by_upper(-9998)
geom = tvtk.GeometryFilter()
configure_input(geom, cropped)
surface1 = tvtk.WarpScalar(scale_factor = self.z_exaggeration)
configure_input(surface1, geom)
triangles = tvtk.TriangleFilter()
configure_input(triangles, surface1)
reduced = tvtk.DecimatePro(target_reduction = 0.5, preserve_topology = True)
configure_input(reduced, triangles)
strips = tvtk.Stripper()
configure_input(strips, reduced)
texture_plane = tvtk.TextureMapToPlane(
origin = (0, self._y_max, 0),
point1 = (self._x_max, self._y_max, 0),
point2 = (0, 0, 0))
configure_input(texture_plane, strips)
map_normals = tvtk.PolyDataNormals()
configure_input(map_normals, texture_plane)
surface = tvtk.PolyDataMapper(scalar_visibility = False)
configure_input(surface, map_normals)
#map_image = tvtk.PNGReader(file_name = overlay_file)
map_texture = tvtk.Texture(interpolate = True)
configure_input(map_texture, self._overlay)
geomap_actor = tvtk.Actor(mapper = surface, texture = map_texture, visibility = True)
renderer.add_actor(geomap_actor)
overhead_light = tvtk.Light()
overhead_light.set_color(1, 1, 1)
overhead_light.focal_point = ((self._x_max - self._x_min) / 2, (self._y_max - self._y_min) / 2, 0)
overhead_light.position = ((self._x_max - self._x_min) / 2, (self._y_max - self._y_min) / 2, 2000)
overhead_light.intensity = 0.5
renderer.add_light(overhead_light)
headlight = tvtk.Light(light_type = 'headlight')
renderer.add_light(headlight)
renderer.active_camera.set(**self.camera_settings)
## add axes indicator:
#axes = tvtk.AxesActor()
#axes.total_length = (300,) * 3
#text_property = tvtk.TextProperty()
#text_property.font_family = 'arial'
#text_property.font_size = 40
#text_property.color = (0, 0, 0)
#axes.x_axis_caption_actor2d.text_actor.text_scale_mode = False
#axes.y_axis_caption_actor2d.text_actor.text_scale_mode = False
#axes.z_axis_caption_actor2d.text_actor.text_scale_mode = False
#axes.x_axis_caption_actor2d.caption_text_property = text_property
#axes.y_axis_caption_actor2d.caption_text_property = text_property
#axes.z_axis_caption_actor2d.caption_text_property = text_property
#axes.x_axis_caption_actor2d.position = [10, 0]
#axes.y_axis_caption_actor2d.position = [20, -20]
#axes.z_axis_caption_actor2d.position = [20, 0]
#axes.x_axis_shaft_property.line_width = 3
#axes.y_axis_shaft_property.line_width = 3
#axes.z_axis_shaft_property.line_width = 3
#renderer.add_actor(axes)
## shift it a little bit:
#transform = tvtk.Transform()
#transform.translate((400, 400, 0))
#axes.user_transform = transform
# add additional actors:
for actor in add_actors:
renderer.add_actor(actor)
if filename is None:
v = ivtk.viewer()
v.scene.add_actors(renderer.actors)
cam = v.scene.camera
else:
render_window = tvtk.RenderWindow(size = (840, 700), position = (100, 100))
render_window.add_renderer(renderer)
cam = renderer.active_camera
if filename is not None:
render_window.render()
render_large = tvtk.RenderLargeImage(input = renderer, magnification = 2)
writer = tvtk.PNGWriter(file_name = filename)
configure_input(writer, render_large)
writer.write()
if crop:
# crop image to non-white pixels:
image = Image.open(filename)
bg = Image.new(image.mode, image.size, (255, 255, 255))
diff = ImageChops.difference(image, bg)
bbox = diff.getbbox()
image = image.crop(bbox)
image.save(filename)
# where X coords vary fastest, Y next and then Z.
try:
import scipy.special
z = numpy.reshape(numpy.transpose(5.0*scipy.special.j0(r)), (-1,) )
except ImportError:
z = numpy.reshape(numpy.transpose(5.0*numpy.sin(r)/r), (-1,) )
# Now set the scalar data for the StructuredPoints object. The
# scalars of the structured points object will be a view into our
# Numeric array. Thus, if we change `z` in-place, the changes will
# automatically affect the VTK arrays.
sp.point_data.scalars = z
# Convert this to a PolyData object.
geom_filter = tvtk.ImageDataGeometryFilter()
configure_input(geom_filter, sp)
# Now warp this using the scalar value to generate a carpet plot.
warp = tvtk.WarpScalar()
configure_input(warp, geom_filter)
# Smooth the resulting data so it looks good.
normals = tvtk.PolyDataNormals()
configure_input(normals, warp)
# The rest of the VTK pipeline.
m = tvtk.PolyDataMapper(scalar_range=(min(z), max(z)))
configure_input(m, normals)
a = tvtk.Actor(mapper=m)
def _handle_failed_image(idiff, src_img, pngr, img_fname):
"""Writes all the necessary images when an image comparison
failed."""
f_base, f_ext = os.path.splitext(img_fname)
# write out the difference file in full.
pngw = tvtk.PNGWriter(file_name=f_base + ".diff.png")
configure_input(pngw, idiff)
pngw.write()
# write the difference image scaled and gamma adjusted for the
# dashboard.
sz = pngr.output.dimensions
if sz[1] <= 250.0:
mag = 1.0
else:
mag = 250.0/sz[1]
shrink = tvtk.ImageResample(interpolate=1)
configure_input(shrink, idiff.output)
shrink.set_axis_magnification_factor(0, mag)
shrink.set_axis_magnification_factor(1, mag)
gamma = tvtk.ImageShiftScale(shift=0, scale=10)
configure_input(gamma, shrink)
jpegw = tvtk.JPEGWriter(file_name=f_base + ".diff.small.jpg",
quality=85)
configure_input(jpegw, gamma)
jpegw.write()
# write out the image that was generated.
pngw.set(file_name=f_base + ".test.png")
configure_input(pngw, src_img)
pngw.write()
# write out a smaller version of the image that was generated.
configure_input(shrink, idiff.input)
jpegw.set(file_name=f_base + ".test.small.jpg")
configure_input(jpegw, shrink)
jpegw.write()
# write out the valid image that matched.
configure_input(shrink, idiff.image)
jpegw.set(file_name=f_base + ".small.jpg")
configure_input(jpegw, shrink)
jpegw.write()
# The angular par of the spherical harmonic (3, 2) x, y, z = np.mgrid[-.5:.5:100j, -.5:.5:100j, -.5:.5:100j] Phi = np.angle((x+y*1j)**2*z) field = mlab.pipeline.scalar_field(x, y, z, Phi) ipw = mlab.pipeline.image_plane_widget(field) mlab.outline(field) surface = mlab.pipeline.builtin_surface() surface.source = 'sphere' surface.data_source.radius = .4 surface.data_source.phi_resolution = 200 surface.data_source.theta_resolution = 200 probe_filter = tvtk.ProbeFilter() configure_input(probe_filter, surface.data_source) configure_source_data(probe_filter, field.outputs[0]) probe = mlab.pipeline.user_defined(surface, filter=probe_filter) surf = mlab.pipeline.surface(probe) fig = mlab.gcf() ################################################################################ # Finally, to inspect the VTK Pipeline (and not the Mayavi one, we # use the TVTK pipeline browser) # Note that for Mayavi version < 3.4.1, there is a bug in the # PipelineBrowser preventing a good display of this pipeline. from tvtk.pipeline.browser import PipelineBrowser browser = PipelineBrowser(fig.scene) browser.show()
def compare_image_with_saved_image(src_img, img_fname, threshold=10,
allow_resize=True):
"""Compares a source image (src_img, which is a tvtk.ImageData)
with the saved image file whose name is given in the second
argument. If the image file does not exist the image is generated
and stored. If not the source image is compared to that of the
figure. This function also handles multiple images and finds the
best matching image. If `allow_resize` is True then the images
are rescaled if they are not of the same size.
"""
f_base, f_ext = os.path.splitext(os.path.abspath(img_fname))
if not os.path.isfile(img_fname):
# generate the image
pngw = tvtk.PNGWriter(file_name=img_fname, input=src_img)
pngw.write()
if VERBOSE:
print("Creating baseline image '%s'."%img_fname)
return
pngr = tvtk.PNGReader(file_name=img_fname)
pngr.update()
if allow_resize:
src_resample = tvtk.ImageResample(interpolate=1,
interpolation_mode='cubic')
configure_input(src_resample, src_img)
img_resample = tvtk.ImageResample(interpolate=1,
interpolation_mode='cubic')
configure_input(img_resample, pngr)
_set_scale(src_resample, img_resample)
idiff = tvtk.ImageDifference()
configure_input(idiff, src_resample.output)
if hasattr(idiff, 'set_image_data'):
idiff.set_image_data(img_resample.output)
else:
idiff.image = img_resample.output
else:
idiff = tvtk.ImageDifference()
configure_input(idiff, src_img)
if hasattr(idiff, 'set_image_data'):
idiff.set_image_data(pngr.output)
else:
idiff.image = pngr.output
idiff.update()
min_err = idiff.thresholded_error
img_err = min_err
best_img = img_fname
err_index = 0
count = 0
if min_err > threshold:
count = 1
test_failed = 1
err_index = -1
while 1: # keep trying images till we get the best match.
new_fname = f_base + "_%d.png"%count
if not os.path.exists(new_fname):
# no other image exists.
break
# since file exists check if it matches.
pngr.file_name = new_fname
pngr.update()
if allow_resize:
_set_scale(src_resample, img_resample)
idiff.update()
alt_err = idiff.thresholded_error
if alt_err < threshold:
# matched,
err_index = count
test_failed = 0
min_err = alt_err
img_err = alt_err
best_img = new_fname
break
else:
if alt_err < min_err:
# image is a better match.
err_index = count
min_err = alt_err
img_err = alt_err
best_img = new_fname
count = count + 1
# closes while loop.
if test_failed:
_handle_failed_image(idiff, src_img, pngr, best_img)
_print_image_error(img_err, err_index, f_base)
msg = "Failed image test: %f\n"%idiff.thresholded_error
raise AssertionError(msg)
# output the image error even if a test passed
_print_image_success(img_err, err_index)
