Enhance restore plot backend for binary image

MLStructFP/__init__.py

@@ -9,7 +9,7 @@
 __description__ = 'Machine learning structural floor plan dataset'
 __keywords__ = ['ml', 'ai', 'floor plan', 'architectural', 'dataset', 'cnn']
 __email__ = 'pablo@ppizarror.com'
-__version__ = '0.6.7'
+__version__ = '0.6.8'
 
 # URL
 __url__ = 'https://github.com/MLSTRUCT/MLSTRUCT-FP'

MLStructFP/db/image/__init__.py

@@ -4,5 +4,5 @@
 Dataset image operations.
 """
 
-from MLStructFP.db.image._rect_binary import RectBinaryImage
+from MLStructFP.db.image._rect_binary import RectBinaryImage, restore_plot_backend
 from MLStructFP.db.image._rect_photo import RectFloorPhoto

MLStructFP/db/image/_rect_binary.py

@@ -4,7 +4,7 @@
 Image of the surroundings of a rect.
 """
 
-__all__ = ['RectBinaryImage']
+__all__ = ['RectBinaryImage', 'restore_plot_backend']
 
 from MLStructFP.db.image._base import BaseImage, TYPE_IMAGE
 from MLStructFP.utils import make_dirs
@@ -31,6 +31,14 @@
 PLOT_BACKEND: str = 'agg'
 
 
+def restore_plot_backend() -> None:
+    """
+    Restores plot backend.
+    """
+    if plt.get_backend() == PLOT_BACKEND:
+        plt.switch_backend(INITIAL_BACKEND)
+
+
 class RectBinaryImage(BaseImage):
     """
     Rect binary image.
@@ -260,7 +268,7 @@ def close(self, restore_plot: bool = True) -> None:
         self._names.clear()
 
         # Restore plot
-        if restore_plot and plt.get_backend() == PLOT_BACKEND:
-            plt.switch_backend(INITIAL_BACKEND)
+        if restore_plot:
+            restore_plot_backend()
 
         self._initialized = False

MLStructFP/db/image/_rect_photo.py

@@ -579,7 +579,7 @@ def _get_crop_image(
 
         return image_kernel
 
-    def close(self) -> None:
+    def close(self, **kwargs) -> None:
         """
         Close and delete all generated figures.
         """