Search code examples
pythonimagenumpymatplotlibpython-imageio

Save 3D plot with imageio


I want to save a 3D plot as a gif. However, for reasons I cannot explain, the following code does not work. I get the error message: Image must be 2D (grayscale, RGB, or RGBA), when I use imageio.mimsave(...), but I saved my image as RGB:

import numpy as np
import matplotlib.pyplot as plt
import imageio

x = [0,1,0,0]
y = [0,0,1,0]
z = [0,0,0,1]

fig = plt.figure(figsize=(15,9))
ax = fig.add_subplot(projection='3d')
ax.scatter(x,y,z,s=1500)

images = []
for n in range(0, 100):
    if n >= 20: 
        ax.azim = ax.azim+1.1 
    fig.canvas.draw()
    image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')
    images.append(image)
   
imageio.mimsave('test.gif', images)


Solution

  • np.frombuffer returns a flat array, its shape(X*Y*3,), on the contrary imageio needs an Y × X × 3 (or maybe × 4 for a RGBA buffer). You must reshape your image

    In [43]: import numpy as np
        ...: import matplotlib.pyplot as plt
        ...: import imageio
        ...: 
        ...: x = [0,1,0,0]
        ...: y = [0,0,1,0]
        ...: z = [0,0,0,1]
        ...: 
        ...: fig = plt.figure(figsize=(5, 3)) # @100 dpi it's 500×300 pixels
        ...: ax = fig.add_subplot(projection='3d')
        ...: ax.scatter(x,y,z,s=500)
        ...: 
        ...: images = []
        ...: for n in range(0, 25):
        ...:     if n >= 15:
        ...:         ax.azim = ax.azim+1.1
        ...:     fig.canvas.draw()
        ...:     image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')
        ...:     images.append(image.reshape(300, 500, 3)) ## Y×X
        ...: 
        ...: imageio.mimsave('test.gif', images)
        ...: 
    
    In [44]: ls -l test.gif
    -rw-r--r-- 1 boffi boffi 825988 May 21 19:46 test.gif
    
    In [45]: 
    

    enter image description here

    Note that I have modified the dimensions and the number of frames to have a smaller GIF.