Merge pull request #277 from neptune-ai/pj/fix-saving-matplotlib-figures

Fixes plot cutoffs

neptune/internal/utils/image.py

@@ -50,19 +50,19 @@ def get_image_content(image):
         try:
             from matplotlib import figure
             if isinstance(image, figure.Figure):
-                return _get_pil_image_data(_figure_to_pil_image(image))
+                return _get_figure_as_image(image)
         except ImportError:
             pass
 
     raise InvalidChannelValue(expected_type='image', actual_type=type(image).__name__)
 
 
+def _get_figure_as_image(figure):
+    with io.BytesIO() as image_buffer:
+        figure.savefig(image_buffer, format='png', bbox_inches="tight")
+        return image_buffer.getvalue()
+
 def _get_pil_image_data(image):
     with io.BytesIO() as image_buffer:
         image.save(image_buffer, format='PNG')
         return image_buffer.getvalue()
-
-
-def _figure_to_pil_image(figure):
-    figure.canvas.draw()
-    return Image.frombytes('RGB', figure.canvas.get_width_height(), figure.canvas.tostring_rgb())

tests/neptune/internal/utils/test_image.py

@@ -24,7 +24,7 @@
 from uuid import uuid4
 import numpy
 
-from neptune.internal.utils.image import get_image_content, _get_pil_image_data
+from neptune.internal.utils.image import get_image_content, _get_pil_image_data, _get_figure_as_image
 
 
 class TestImage(unittest.TestCase):
@@ -102,4 +102,4 @@ def test_get_image_content_from_figure(self):
         expected_image = Image.frombytes('RGB', figure.canvas.get_width_height(), figure.canvas.tostring_rgb())
 
         # expect
-        self.assertEqual(get_image_content(figure), _get_pil_image_data(expected_image))
+        self.assertEqual(get_image_content(figure), _get_figure_as_image(figure))