Implement support of Matplotlib figures

Hello @lanpa,

First of all, thanks a lot for investing the time and developing this project. I find it super useful whenever I try (and usually fail 😉 hacking PyTorch.

Frequently, when I work there’s a need to create Matplotlib plot and add it to TensorBoard. There are known and well documented ways to do it by converting the figure to image. One of your examples in this repo actually demonstrates it and there’s a few other ways to achieve the same result.

I was wondering if it would make sense to implement SummaryWriter.add_figure(...) method that would accept matplotlib.Figure, convert it into a RGB image and then add it to TensorBoard. I think the implementation could look roughly like this:

import matplotlib.pyplot as plt
import matplotlib.backends.backend_agg as plt_backend_agg

class SummaryWriter:
    def add_figure(self, tag, figure, global_step=None, close_figure=True):
        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 = data.reshape([h, w, 4])[:, :, 0:3]
        if close_figure:
            plt.close(figure)
        self.add_image(tag, image, global_step)

I’ll submit a PR if you think adding the method makes sense and it fits well into the API.

Have a good one!

[lanpa]

thanks!

[vojtamolda]

Agreed. Thanks. PR coming in a little bit.