pythonanimationvisualizationmayavimayavi.mlab

Saving a 3D animated mayavi figure into a in-memory file (buffer)


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


Solution

  • 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
    

    Digital Elevation Model