代码收藏家技术教程 2025-01-17

python中matplotlib 获取当前figure为numpy.ndarry或者PIL.Image对象的rgb图像数据bug记录

需求说明

我所搜集到广为流传的做法

1. 使用BytesIO和savefig

2. 使用fig.canvas.tostring_rgb()方法

3. 使用fig.canvas.buffer_rgba()方法 (省流方法)

reference

python中matplotlib 获取当前figure为numpy.ndarry或者PIL.Image对象的rgb图像数据bug记录

需求说明

使用matplotlib.pyplot绘制一个动态图像并记录每一帧RGB图像，便于后续处理这里采用平移sin图像作为演示

import matplotlib.pyplot as plt
import numpy as np
# 初始化图像
fig, ax = plt.subplots()
x = np.linspace(0, 2 * np.pi, 1000)
line, = ax.plot(x, np.sin(x))
# 更新每一帧图像
def update(frame):
	# 更新图像
    line.set_ydata(np.sin(x + frame/100))
    # 这里获取当前帧图像并返回
    # return image_data
    
# frames = []
# 运行100帧并使用列表记录
# for i in range(1000):
#     img = update(i)
#     frames.append(img)

我所搜集到广为流传的做法

1. 使用BytesIO和savefig

采用以下代码获取每一帧图像数据

from io import BytesIO
buffer = BytesIO()
fig.savefig(buffer)
buffer.seek(0)
image_buffer  = Image.open(buffer)
image_data= np.array(image_buffer  ,dtype=np.uint8)
buffer.close()

该代码能正常运行但效率很低

2. 使用fig.canvas.tostring_rgb()方法

# 绘制当前图像
fig.canvas.draw()
# 从画布获取图像数据
image_data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
# 获取图像的宽度和高度
width, height = fig.canvas.get_width_height()
# 重塑图像数据为RGB格式
image_data = image_data.reshape(height, width, 3)

该代码也能正常运行但会报警告

MatplotlibDeprecationWarning: The tostring_rgb function was deprecated in Matplotlib 3.8 and will be removed in 3.10. Use buffer_rgba instead.

到此需求已经实现，如果你不是像我一样的强迫症患者，就可以不用往下看了！

3. 使用fig.canvas.buffer_rgba()方法 (省流方法)

作为强迫症晚期患者，完全不能忍受警告出现！继续修改 (T_T)

按照警告的提示，将**.tostring_rgb() 换为.buffer_rgba()**
然后报错

ValueError: cannot reshape array of size 1228800 into shape (480,640,3)

继续根据乘除法和方法名称发现**.buffer_rgba()生成的rgba格式图像，进而将image_data = image_data.reshape(height, width, 3)改为image_data = image_data.reshape(height, width, 4)**

# 绘制当前图像
fig.canvas.draw()
# 从画布获取图像数据
image_data = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
# 获取图像的宽度和高度
width, height = fig.canvas.get_width_height()
# 重塑图像数据为RGBA格式
image_data = image_data.reshape(height, width, 4)
# 转换为RGB格式
image_data = image_data[:, :, :3]

移植到项目中运行！
没有警告和错误，搞定！！！了吗？

程序跑了30min后，成就满满的显示结果，发现只显示最后一帧，孩子天塌了 (TT)*

于是继续改呀改，改呀改，改呀改… 眼睛一睁一闭，一天过去了
谁懂没有没有报错的bug的感受啊 TT
最后在matplolib的github源码上找到了代码的问题
matplotlib的github源码

 def buffer_rgba(self):
     return memoryview(self._renderer)

 def tostring_argb(self):
     return np.asarray(self._renderer).take([3, 0, 1, 2], axis=2).tobytes()

 @_api.deprecated("3.8", alternative="buffer_rgba")
 def tostring_rgb(self):
     return np.asarray(self._renderer).take([0, 1, 2], axis=2).tobytes()

以下内容纯属瞎猜，不一定正确，请注意辨别！！
简单读一下代码和api的名称。可知**.tostring_rgb()和.buffer_rgba()**区别在于

tostring_rgb()返回的是numpy.array的对象
buffer_rgba()返回的是self._renderer的对象(图像渲染器对象)的内存地址(类似于C/C++语言的指针？)

故之前的代码问题在于，返回的每一帧图像数据只是最后一帧的图像数据
得知了问题原因，则马上把代码改为:

# 绘制当前图像
fig.canvas.draw()
# 从画布获取图像数据
image_data = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
# 获取图像的宽度和高度
width, height = fig.canvas.get_width_height()
# 重塑图像数据为RGBA格式
image_data = image_data.reshape(height, width, 4)
# 转换为RGB格式并将数据复制到本地变量
image_data = image_data[:, :, :3].copy()

为了防止有bug产生影响心情、浪费表情，老老实实采用平移sin图像作为演示
整体代码如下：

import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots()
x = np.linspace(0, 2 * np.pi, 1000)
line, = ax.plot(x, np.sin(x))
fig.canvas.draw()
# 从画布获取图像数据地址
buffer= np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
# 获取图像的宽度和高度
width, height = fig.canvas.get_width_height()
def update(frame):
    line.set_ydata(np.sin(x + frame/100))
    # 绘制当前图像
    fig.canvas.draw()
    # 重塑图像数据为RGBA格式
    image_data = buffer.reshape(height, width, 4).copy()
    # 转换为RGB格式并将数据复制到本地变量
    image_data = image_data[:, :, :3]
    return image_data
    
frames = []
for i in range(1000):
    img = update(i)
    frames.append(img)
print(len(frames))
# 采用OpenCV显示，也可以采用PIL显示
import cv2
for img in frames:
    cv2.imshow('show',img)
    cv2.waitKey(10)
cv2.destroyAllWindows()

至此问题完全解决
如果您觉得有用的，不妨点个赞，谢谢（~ _ ~）