matplotlib figure转为numpy array或者PIL图像进行显示

matplotlib figure转为numpy array或者PIL图像进行显示

实现将matplotlib绘制的图像转换为numpy数组,并使用PIL或者OpenCV进行显示

参考资料:http://www.icare.univ-lille1.fr/tutorials/convert_a_matplotlib_figure

# -*- coding: utf-8 -*-
"""
# --------------------------------------------------------
# @Author : panjq
# @E-mail : [email protected]
# @Date   : 2020-02-05 11:01:49
# --------------------------------------------------------
"""

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


def fig2data(fig):
    """
    fig = plt.figure()
    image = fig2data(fig)
    @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
    @param fig a matplotlib figure
    @return a numpy 3D array of RGBA values
    """
    import PIL.Image as Image
    # draw the renderer
    fig.canvas.draw()

    # Get the RGBA buffer from the figure
    w, h = fig.canvas.get_width_height()
    buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)
    buf.shape = (w, h, 4)

    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
    buf = np.roll(buf, 3, axis=2)
    image = Image.frombytes("RGBA", (w, h), buf.tostring())
    image = np.asarray(image)
    return image


if __name__ == "__main__":
    # Generate a figure with matplotlib
    figure = plt.figure()
    plot = figure.add_subplot(111)

    # draw a cardinal sine plot
    x = np.arange(1, 100, 0.1)
    y = np.sin(x) / x
    plot.plot(x, y)
    plt.show()
    ##
    image = fig2data(figure)
    cv2.imshow("image", image)
    cv2.waitKey(0)

如果在Pycharm中运行出现"AttributeError: 'FigureCanvasInterAgg' object has no attribute 'renderer'" 的错误,请参考修改:

https://stackoverflow.com/questions/51214140/attributeerror-figurecanvasinteragg-object-has-no-attribute-renderer解决