Source code for tensorboardX.utils
# Functions for converting
[docs]def figure_to_image(figures, close=True):
"""Render matplotlib figure to numpy format.
Note that this requires the ``matplotlib`` package.
Args:
figure (matplotlib.pyplot.figure) or list of figures: figure or a list of figures
close (bool): Flag to automatically close the figure
Returns:
numpy.array: image in [CHW] order
"""
import numpy as np
try:
import matplotlib.pyplot as plt
import matplotlib.backends.backend_agg as plt_backend_agg
except ModuleNotFoundError:
print('please install matplotlib')
def render_to_rgb(figure):
canvas = plt_backend_agg.FigureCanvasAgg(figure)
canvas.draw()
data = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8)
w, h = figure.canvas.get_width_height()
image_hwc = data.reshape([h, w, 4])[:, :, 0:3]
image_chw = np.moveaxis(image_hwc, source=2, destination=0)
if close:
plt.close(figure)
return image_chw
if not isinstance(figures, list):
image = render_to_rgb(figures)
return image
else:
images = [render_to_rgb(figure) for figure in figures]
return np.stack(images)
def graphviz_to_image():
pass