I want to save a gif or a mp4 of an animated 3D figure I build with mayavi.
When I was doing similar things with 2D figures and imagio, I could save the data to a buffer and append them to a imagio writer. That was very fast and saved the trouble of generating thousands of intermediary png files. However, I can not find a way to save them mayavi figure to a buffer rather than a file.
@mlab.animate(delay=10, ui=False)
def update_animation(Z, surface, writer):
# Sets nan pixels to white
surface.module_manager.scalar_lut_manager.lut.nan_color = 0, 0, 0, 0
surface.update_pipeline()
f = mlab.gcf()
# Sets background to white
f.scene.background = (1,1,1)
t = 2.0
while t <= Z.shape[0]:
f.scene.camera.azimuth(1)
f.scene.render()
surface.mlab_source.scalars = get_band(Z, int(t)-1)
t += 0.1
#mlab.savefig('tmp/' + str(t) + '.png') # this used to work, but generates thousands of png
with io.BytesIO() as buff: # so I want instead to try this strategy
mlab.savefig(buff, format='raw')
buff.seek(0)
data = numpy.frombuffer(buff.getvalue(), dtype=numpy.uint8)
w, h = fig.canvas.get_width_height()
im = data.reshape((int(h), int(w), -1))
writer.append_data(im)
yield
In the main function:
# Z is a 3D np array built from reading a multiband raster with rasterio
with imageio.get_writer(output, mode='I') as writer:
surface = mlab.imshow(get_band(Z, 0), colormap='viridis')
a = update_animation(Z, surface, writer)
mlab.show()
But the error message states:
TypeError: expected str, bytes or os.PathLike object, not _io.BytesIO
From mayavi documentation:
Starting from Mayavi version 3.4.0, the mlab
screenshot()
can be used to take a screenshot of the current figure, to integrate in a matplotlib plot.
So one can actually totally bypass the code relative to BytesIO
, and directly invoke the imagio writer with writer.append_data(mlab.screenshot())
:
# Z is a 3D np array built from reading a multiband raster with rasterio
with imageio.get_writer(output, mode='I') as writer:
surface = mlab.imshow(get_band(Z, 0), colormap='viridis')
lut_manager = mlab.scalarbar(orientation='vertical')
mlab.colorbar()
a = update_animation(Z, surface, writer)
mlab.show()
And the animation function itself:
@mlab.animate(delay=10, ui=False)
def update_animation(Z, surface, writer):
# Sets nan pixels to white
surface.module_manager.scalar_lut_manager.lut.nan_color = 0, 0, 0, 0
surface.update_pipeline()
f = mlab.gcf()
t = 2.0
while t <= Z.shape[0]:
f.scene.camera.azimuth(1)
f.scene.render()
surface.mlab_source.scalars = get_band(Z, int(t)-1)
t += 0.1
writer.append_data(mlab.screenshot()) # here !
yield