image_conversion.py 1008 B

1234567891011121314151617181920212223242526272829303132
  1. from PIL import Image
  2. import numpy as np
  3. def fig2data(fig):
  4. """
  5. @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
  6. @param fig a matplotlib figure
  7. @return a numpy 3D array of RGBA values
  8. """
  9. # draw the renderer
  10. fig.canvas.draw()
  11. # Get the RGBA buffer from the figure
  12. w,h = fig.canvas.get_width_height()
  13. buf = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
  14. buf.shape = (w, h, 3)
  15. # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
  16. buf = np.roll(buf, 3, axis=2)
  17. return buf
  18. def fig2img(fig):
  19. """
  20. @brief Convert a Matplotlib figure to a PIL Image in RGBA format and return it
  21. @param fig a matplotlib figure
  22. @return a Python Imaging Library (PIL) image : default size (480,640,3)
  23. """
  24. # put the figure pixmap into a numpy array
  25. buf = fig2data(fig)
  26. w, h, d = buf.shape
  27. return Image.frombytes("RGB", (w, h), buf.tostring())