import Image
def fig2img ( fig ): """ @brief Convert a Matplotlib figure to a PIL Image in RGBA format and return it @param fig a matplotlib figure @return a Python Imaging Library ( PIL ) image """ # put the figure pixmap into a numpy array buf = fig2data ( fig ) w, h, d = buf.shape return Image.frombytes( "RGBA", ( w ,h ), buf.tostring( ) )
def 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 """ # 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 )
return buf
def rgba2rgb(img, c=(0, 0, 0), path='foo.jpg', is_already_saved=False, if_load=True): if not is_already_saved: background = Image.new("RGB", img.size, c) background.paste(img, mask=img.split()[3]) # 3 is the alpha channel
background.save(path, 'JPEG', quality=100)
is_already_saved = True
if if_load:
if is_already_saved:
im = Image.open(path)
return np.array(im)
else:
raise ValueError('No image to load.')