/usr/share/pyshared/mayavi/tools/figure.py is in mayavi2 4.1.0-1.
This file is owned by root:root, with mode 0o644.
The actual contents of the file can be viewed below.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 | """
Functions related to creating the engine or the figures.
"""
# Author: Gael Varoquaux <gael.varoquaux@normalesup.org>
# Copyright (c) 2007, Enthought, Inc.
# License: BSD Style.
# Standard library imports.
from types import IntType
import gc
import warnings
import copy
import numpy as np
# Enthought imports
from pyface.timer.api import do_later
# imports
from tvtk.api import tvtk
from mayavi.core.scene import Scene
from mayavi.core.registry import registry
from .camera import view
from .engine_manager import get_engine, options, set_engine
######################################################################
# A list to store the allocated scene numbers
__scene_number_list = set((0,))
def figure(figure=None, bgcolor=None, fgcolor=None, engine=None,
size=(400, 350)):
""" Creates a new scene or retrieves an existing scene. If the mayavi
engine is not running this also starts it.
**Keyword arguments**
:figure: The name of the figure, or handle to it.
:bgcolor: The color of the background (None is default).
:fgcolor: The color of the foreground, that is the color of all text
annotation labels (axes, orientation axes, scalar bar
labels). It should be sufficiently far from `bgcolor`
to see the annotation texts. (None is default).
:engine: The mayavi engine that controls the figure.
:size: The size of the scene created, in pixels. May not apply
for certain scene viewer.
"""
if isinstance(figure, Scene):
if figure.scene is None:
engine = registry.find_scene_engine(figure)
else:
engine = registry.find_scene_engine(figure.scene)
set_engine(engine)
engine.current_scene = figure
else:
if engine is None:
engine = get_engine()
if figure is None:
name = max(__scene_number_list) + 1
__scene_number_list.update((name,))
name = 'Mayavi Scene %d' % name
engine.new_scene(name=name, size=size)
engine.current_scene.name = name
else:
if type(figure) in (IntType, np.int, np.int0, np.int8,
np.int16, np.int32, np.int64):
name = int(figure)
__scene_number_list.update((name,))
name = 'Mayavi Scene %d' % name
else:
name = str(figure)
# Go looking in the engine see if the scene is not already
# running
for scene in engine.scenes:
if scene.name == name:
engine.current_scene = scene
return scene
else:
engine.new_scene(name=name, size=size)
engine.current_scene.name = name
figure = engine.current_scene
scene = figure.scene
if scene is not None:
if hasattr(scene, 'isometric_view'):
scene.isometric_view()
else:
# Not every viewer might implement this method
view(40, 50)
scene = figure.scene
if scene is not None:
if bgcolor is None:
bgcolor = options.background_color
scene.background = bgcolor
if fgcolor is None:
fgcolor = options.foreground_color
scene.foreground = fgcolor
return figure
def gcf(engine=None):
"""Return a handle to the current figure.
You can supply the engine from which you want to retrieve the
current figure, if you have several mayavi engines.
"""
if engine is None:
engine = get_engine()
scene = engine.current_scene
if scene is None:
return figure(engine=engine)
return scene
def clf(figure=None):
"""Clear the current figure.
You can also supply the figure that you want to clear.
"""
try:
if figure is None:
scene = gcf()
else:
scene = figure
disable_render = scene.scene.disable_render
scene.scene.disable_render = True
scene.children[:] = []
scene._mouse_pick_dispatcher.clear_callbacks()
scene.scene.disable_render = disable_render
except AttributeError:
pass
gc.collect()
def close(scene=None, all=False):
""" Close a figure window
close() by itself closes the current figure.
close(num) closes figure number num.
close(name) closes figure named name.
close(figure), where figure is a scene instance, closes that
figure.
close(all=True) closes all figures controlled by mlab
"""
if all is True:
engine = get_engine()
# We need the copy, as the list gets pruned as we close scenes
for scene in copy.copy(engine.scenes):
engine.close_scene(scene)
return
if not isinstance(scene, Scene):
engine = get_engine()
if scene is None:
scene = engine.current_scene
else:
if type(scene) in (IntType, np.int, np.int0, np.int8,
np.int16, np.int32, np.int64):
scene = int(scene)
name = 'Mayavi Scene %d' % scene
else:
name = str(scene)
# Go looking in the engine see if the scene is not already
# running
for scene in engine.scenes:
if scene.name == name:
break
else:
warnings.warn('Scene %s not managed by mlab' % name)
return
else:
if scene.scene is None:
engine = registry.find_scene_engine(scene)
else:
engine = registry.find_scene_engine(scene.scene)
engine.close_scene(scene)
def draw(figure=None):
""" Forces a redraw of the current figure.
"""
if figure is None:
figure = gcf()
figure.render()
def savefig(filename, size=None, figure=None, magnification='auto',
**kwargs):
""" Save the current scene.
The output format are deduced by the extension to filename.
Possibilities are png, jpg, bmp, tiff, ps, eps, pdf, rib (renderman),
oogl (geomview), iv (OpenInventor), vrml, obj (wavefront)
**Parameters**
:size: the size of the image created (unless magnification is
set, in which case it is the size of the window used
for rendering).
:figure: the figure instance to save to a file.
:magnification: the magnification is the scaling between the
pixels on the screen, and the pixels in the
file saved. If you do not specify it, it will be
calculated so that the file is saved with the
specified size. If you specify a magnification,
Mayavi will use the given size as a screen size,
and the file size will be 'magnification * size'.
**Notes**
If the size specified is larger than the window size, and no
magnification parameter is passed, the magnification of the scene
is changed so that the image created has the requested size.
Please note that if you are trying to save images with sizes
larger than the window size, there will be additional computation
cost.
Any extra keyword arguments are passed along to the respective
image format's save method.
"""
if figure is None:
figure = gcf()
current_mag = figure.scene.magnification
try:
if size is not None:
current_x, current_y = tuple(figure.scene.get_size())
target_x, target_y = size
if magnification is 'auto':
magnification = max(target_x//current_x,
target_y//current_y) + 1
target_x = int(target_x/magnification)
target_y = int(target_y/magnification)
size = target_x, target_y
elif magnification is 'auto':
magnification = 1
figure.scene.magnification = int(magnification)
figure.scene.save(filename,
size=size,
**kwargs)
finally:
figure.scene.magnification = int(current_mag)
def sync_camera(reference_figure, target_figure):
""" Synchronise the camera of the target_figure on the camera of the
reference_figure.
"""
reference_figure.scene._renderer.sync_trait('active_camera',
target_figure.scene._renderer)
target_figure.scene._renderer.active_camera.on_trait_change(
lambda: do_later(target_figure.scene.render))
def screenshot(figure=None, mode='rgb', antialiased=False):
""" Return the current figure pixmap as an array.
**Parameters**
:figure: a figure instance or None, optional
If specified, the figure instance to capture the view of.
:mode: {'rgb', 'rgba'}
The color mode of the array captured.
:antialiased: {True, False}
Use anti-aliasing for rendering the screenshot.
Uses the number of aa frames set by
figure.scene.anti_aliasing_frames
**Notes**
On most systems, this works similarly to taking a screenshot of
the rendering window. Thus if it is hidden by another window, you
will capture the other window. This limitation is due to the
heavy use of the hardware graphics system.
**Examples**
This function can be useful for integrating 3D plotting with
Mayavi in a 2D plot created by matplotlib.
>>> from mayavi import mlab
>>> mlab.test_plot3d()
>>> arr = mlab.screenshot()
>>> import pylab as pl
>>> pl.imshow(arr)
>>> pl.axis('off')
>>> pl.show()
"""
if figure is None:
figure = gcf()
x, y = tuple(figure.scene.get_size())
# Try to lift the window
figure.scene._lift()
if mode == 'rgb':
out = tvtk.UnsignedCharArray()
shape = (y, x, 3)
pixel_getter = figure.scene.render_window.get_pixel_data
pg_args = (0, 0, x-1, y-1, 1, out)
elif mode == 'rgba':
out = tvtk.FloatArray()
shape = (y, x, 4)
pixel_getter = figure.scene.render_window.get_rgba_pixel_data
pg_args = (0, 0, x-1, y-1, 1, out)
else:
raise ValueError('mode type not understood')
if antialiased:
# save the current aa value to restore it later
old_aa = figure.scene.render_window.aa_frames
figure.scene.render_window.aa_frames = figure.scene.anti_aliasing_frames
figure.scene.render()
pixel_getter(*pg_args)
figure.scene.render_window.aa_frames = old_aa
figure.scene.render()
else:
pixel_getter(*pg_args)
# Return the array in a way that pylab.imshow plots it right:
out = out.to_array()
out.shape = shape
out = np.flipud(out)
return out
|