File size: 1,587 Bytes
8026e91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import matplotlib.pyplot as plt
import unittest
from tensorboardX import SummaryWriter
class FigureTest(unittest.TestCase):
def test_figure(self):
writer = SummaryWriter()
figure, axes = plt.figure(), plt.gca()
circle1 = plt.Circle((0.2, 0.5), 0.2, color='r')
circle2 = plt.Circle((0.8, 0.5), 0.2, color='g')
axes.add_patch(circle1)
axes.add_patch(circle2)
plt.axis('scaled')
plt.tight_layout()
writer.add_figure("add_figure/figure", figure, 0, close=False)
assert plt.fignum_exists(figure.number) is True
writer.add_figure("add_figure/figure", figure, 1)
assert plt.fignum_exists(figure.number) is False
writer.close()
def test_figure_list(self):
writer = SummaryWriter()
figures = []
for i in range(5):
figure = plt.figure()
plt.plot([i * 1, i * 2, i * 3], label="Plot " + str(i))
plt.xlabel("X")
plt.xlabel("Y")
plt.legend()
plt.tight_layout()
figures.append(figure)
writer.add_figure("add_figure/figure_list", figures, 0, close=False)
assert all([plt.fignum_exists(figure.number) is True for figure in figures])
writer.add_figure("add_figure/figure_list", figures, 1)
assert all([plt.fignum_exists(figure.number) is False for figure in figures])
writer.close()
|