Created
August 12, 2024 00:59
-
-
Save dkirkby/f071c97c7b522e854ead5e74a2177d84 to your computer and use it in GitHub Desktop.
Capture matplotlib output to numpy array
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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]) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment