CODEWITHSUNDEEP
pythonimagenumpymatplotlibpython-imageio
Physics
Physics

Reputation: 113

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)

Upvotes: 1

Views: 1431

Answers (1)

gboffi
gboffi

Reputation: 25093

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.

Upvotes: 2

Related Questions