Repository: aimhubio/aim

Issue:
Add matplotlib integration
We need let users to track matplotlib figures.
Need to investigate matplotlib library and understand figure's dump/load procedure and integrate it into our system.

Motivation: Make users more engaged and feel comfortable using aim for their projects


Discussion:

Resolving patch:
diff --git a/aim/sdk/objects/figure.py b/aim/sdk/objects/figure.py
--- a/aim/sdk/objects/figure.py
+++ b/aim/sdk/objects/figure.py
@@ -5,10 +5,12 @@
 
 @CustomObject.alias('aim.figure')
 class Figure(CustomObject):
-    """Figure object used to store plotly figures in Aim repository.
+    """
+    Figure object can be used for storing Plotly or Matplotlib figures into Aim repository.
+    Core functionality is based on Plotly.
 
     Args:
-         obj (:obj:): plotly figure object.
+         obj (:obj:): plotly or matplotlib figure object.
     """
 
     AIM_NAME = 'aim.figure'
@@ -16,7 +18,9 @@ class Figure(CustomObject):
     def __init__(self, obj):
         super().__init__()
 
-        if inst_has_typename(obj, ['Figure', 'BaseFigure']):
+        if inst_has_typename(obj, ['matplotlib', 'Figure']):
+            self._from_matplotlib_figure(obj)
+        elif inst_has_typename(obj, ['plotly', 'Figure', 'BaseFigure']):
             self._prepare(obj)
         else:
             raise TypeError('Object is not a Plotly Figure instance')
@@ -34,6 +38,19 @@ def _prepare(self, obj):
         self.storage['format'] = 'raw_json'
         self.storage['data'] = BLOB(data=obj.to_json())
 
+    def _from_matplotlib_figure(self, obj):
+        try:
+            from plotly.tools import mpl_to_plotly
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError('Plotly is required to track matplotlib figure.')
+
+        try:
+            plotly_obj = mpl_to_plotly(obj)
+        except ValueError as err:
+            raise ValueError(f'Failed to convert matplotlib figure to plotly figure: {err}')
+
+        return self._prepare(plotly_obj)
+
     def json(self):
         blob_data = self.storage['data']
         return blob_data.data
diff --git a/aim/sdk/objects/image.py b/aim/sdk/objects/image.py
--- a/aim/sdk/objects/image.py
+++ b/aim/sdk/objects/image.py
@@ -61,6 +61,8 @@ def __init__(self, image, caption: str = '', format='png', quality=90, optimize=
             self._from_tf_tensor(image, params)
         elif inst_has_typename(image, ['numpy', 'array']):
             self._from_numpy_array(image, params)
+        elif inst_has_typename(image, ['Figure', 'matplotlib', 'figure']):
+            self._from_matplotlib_figure(image, params)
         else:
             raise TypeError(f'Cannot convert to aim.Image. Unsupported type {type(image)}.')
         self.caption = caption
@@ -222,6 +224,18 @@ def _from_tf_tensor(self, tensor, params):
             pil_image = PILImage.fromarray(array)
         self._from_pil_image(pil_image, params)
 
+    def _from_matplotlib_figure(self, figure, params):
+        try:
+            assert hasattr(figure, 'savefig')
+        except AssertionError:
+            raise ValueError('Cannot convert from matplotlib figure.')
+
+        buffer = BytesIO()
+        figure.savefig(buffer)
+        buffer.seek(0)
+
+        return self._from_pil_image(PILImage.open(buffer), params)
+
     def __eq__(self, other):
         if not isinstance(other, Image):
             return False


Test patch: