Search code examples
pythonmayavi

Screenshot in mayavi with transparency


I am trying to save the screenshot with transparency of mayavi scene. Based this discussion I wrote a script to save the figure, but the result is a mess in output. Below the minimal working example: in mayavi I create a scene with two spheres and save it first in RGB, then in RGBA format. Saving to RGB works while RGBA file is a mess. I believe, the problem is on mayavi side because if I take RGB image from mayavi, add transparency channel and save the file with PIL, the result is what I expect.

Is this a bug or the mayavi rgba format should be somehow converted to the format accepted by PIL?

from mayavi import mlab
from PIL import Image

fig=mlab.figure(1, bgcolor=(1, 1, 1), size=(700, 700))
# Set camera position and properties
fig.scene.parallel_projection = True
fig.scene.show_axes = True

# Draw atoms
x, y, z, t = [0.0,1.0] , [0.0,1.0], [0.0,0.0], [1,2]
dat = mlab.pipeline.scalar_scatter(x, y, z, t, figure=fig)
fig = mlab.pipeline.glyph(dat,scale_mode='none', scale_factor=0.5, figure=fig)

imgmap_RGB = mlab.screenshot(figure=fig, mode='rgb', antialiased=True)
img_RGB = Image.fromarray(imgmap_RGB, 'RGB')
img_RGB.save('foo_RGB.png')

imgmap_RGBA = mlab.screenshot(figure=fig, mode='rgba', antialiased=True)
img_RGBA = Image.fromarray(imgmap_RGBA, 'RGBA')
img_RGBA.save('foo_RGBA.png')

mlab.show()

Solution

  • For some reason that I don't know, mayavi will return floats between 0 and 1 for the RGBA data and unsigned ints for the RGB data, see https://github.com/enthought/mayavi/blob/master/mayavi/tools/figure.py#L304 (I couldn't find the info in the docs).

    To convert, replace the img_RGBA = ... line with

    img_RGBA = Image.fromarray(np.array(imgmap_RGBA*255, dtype=np.uint8))
    

    I could successfully view the png file afterwards.