From 681224c61098915d0cfad476e65bb987cbaf9d06 Mon Sep 17 00:00:00 2001 From: Piotr Jander Date: Tue, 29 Sep 2020 11:18:28 +0200 Subject: [PATCH] Fix --- neptune/internal/utils/image.py | 9 +++++++-- tests/neptune/internal/utils/test_image.py | 4 ++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/neptune/internal/utils/image.py b/neptune/internal/utils/image.py index 970f9298f..67a048e00 100644 --- a/neptune/internal/utils/image.py +++ b/neptune/internal/utils/image.py @@ -50,14 +50,19 @@ def get_image_content(image): try: from matplotlib import figure if isinstance(image, figure.Figure): - return _get_pil_image_data(image) + return _get_figure_as_image(image) except ImportError: pass raise InvalidChannelValue(expected_type='image', actual_type=type(image).__name__) -def _get_pil_image_data(figure): +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() diff --git a/tests/neptune/internal/utils/test_image.py b/tests/neptune/internal/utils/test_image.py index e095d0846..f2aa40b49 100644 --- a/tests/neptune/internal/utils/test_image.py +++ b/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)) \ No newline at end of file + self.assertEqual(get_image_content(figure), _get_figure_as_image(figure))