Capture matplotlib output to numpy array

plt_to_numpy.py

def plt_to_numpy_demo(width=15, height=11, dpi=32, offscreen=True):

    if offscreen:
        # Switch to Agg backend to prevent display
        backend = matplotlib.get_backend()
        plt.close('all')
        plt.switch_backend('agg')
    # Create figure with exactly width x height pixels
    fig = plt.figure(figsize=(width/dpi, height/dpi), dpi=dpi, frameon=False)
    ax = fig.add_subplot()
    ax.set_axis_off()
    plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
    plt.margins(0,0)
    ax.set(xlim=(-0.5,width-0.5), ylim=(height-0.5,-0.5))
    try:
        # Drawing goes here...
        # Any pixels not explicitly drawn will be black
        ax.add_artist(plt.Rectangle((-0.5,-0.5),width=width, height=height, ec='none', fc='b'))
        ax.add_artist(plt.Rectangle((1.5,0.5), width=width-3, height=height-3, ec='none', fc='r'))
        # Must use fig.canvas() to avoid margins
        fig.canvas.draw()
        data = np.array(fig.canvas.renderer.buffer_rgba())[:,:,:3]
        assert data.shape == (height, width, 3)
    except Exception as e:
        print(e)
    if offscreen:
        # Switch back to original backend
        plt.close('all')
        plt.switch_backend(backend)
    return data[:,:]

data1 = plt_to_numpy_demo(7, 6, 32, offscreen=False)
plt.savefig('data.png')
data2 = plt.imread('data.png')

print(data2[:,:,0])
data3 = plt_to_numpy_demo(7, 6, 32, offscreen=True)
print(data3[:,:,0])