From 04de40e11ec6e7bacc656320caa8979a7f1eb68b Mon Sep 17 00:00:00 2001
From: "Pablo Pizarro R." <pablo@ppizarror.com>
Date: Tue, 13 Aug 2024 00:49:29 -0400
Subject: [PATCH] Fix restore plot

---
 MLStructFP/__init__.py              |  2 +-
 MLStructFP/db/image/_base.py        |  5 +++++
 MLStructFP/db/image/_rect_binary.py | 11 ++++++++---
 MLStructFP/db/image/_rect_photo.py  |  5 ++++-
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/MLStructFP/__init__.py b/MLStructFP/__init__.py
index c611220..fcf4ec8 100644
--- a/MLStructFP/__init__.py
+++ b/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.1'
+__version__ = '0.6.2'
 
 # URL
 __url__ = 'https://github.com/MLSTRUCT/MLSTRUCT-FP'
diff --git a/MLStructFP/db/image/_base.py b/MLStructFP/db/image/_base.py
index 012c426..b38b18e 100644
--- a/MLStructFP/db/image/_base.py
+++ b/MLStructFP/db/image/_base.py
@@ -28,6 +28,7 @@ class BaseImage(ABC):
     """
     _image_size: int
     _images: List['np.ndarray']  # List of stored images during make_region
+    _last_make_region_time: float  # Total time for last make region
     _names: List[str]
     _path: str
     _save_images: bool
@@ -98,6 +99,10 @@ def make_region(self, xmin: NumberType, xmax: NumberType, ymin: NumberType, ymax
         """
         raise NotImplementedError()
 
+    @property
+    def make_region_last_time(self) -> float:
+        return self._last_make_region_time
+
     def export(self, filename: str, close: bool = True, compressed: bool = True) -> None:
         """
         Export saved images to numpy format and then removes all data.
diff --git a/MLStructFP/db/image/_rect_binary.py b/MLStructFP/db/image/_rect_binary.py
index 28bda82..0f154a8 100644
--- a/MLStructFP/db/image/_rect_binary.py
+++ b/MLStructFP/db/image/_rect_binary.py
@@ -17,6 +17,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import os
+import time
 
 if TYPE_CHECKING:
     from MLStructFP.db._c_rect import Rect
@@ -74,7 +75,7 @@ def init(self) -> 'RectBinaryImage':
         """
         plt.switch_backend('agg')
         self._initialized = True
-        self.close()
+        self.close(restore_plot=False)
         self._initialized = True
         return self
 
@@ -168,6 +169,7 @@ def make_region(self, xmin: NumberType, xmax: NumberType, ymin: NumberType, ymax
         """
         if not self._initialized:
             raise RuntimeError('Exporter not initialized, use .init()')
+        t0 = time.time()
         store_matplotlib_figure = not HIGHLIGHT_RECT
         fig, ax = self._get_floor_plot(floor, rect, store=store_matplotlib_figure)
 
@@ -226,12 +228,15 @@ def make_region(self, xmin: NumberType, xmax: NumberType, ymin: NumberType, ymax
         del im, im2, im3, im4, fig, ax
 
         # Returns the image index on the library array
+        self._last_make_region_time = time.time() - t0
         return len(self._images) - 1, array
 
-    def close(self) -> None:
+    def close(self, restore_plot: bool = True) -> None:
         """
         Close and delete all generated figures.
         This function also restores plot engine.
+
+        :param restore_plot: Restores plotting engine
         """
         if not self._initialized:
             raise RuntimeError('Exporter not initialized, it cannot be closed')
@@ -247,7 +252,7 @@ def close(self) -> None:
         self._names.clear()
 
         # Restore plot
-        if plt.get_backend() == 'agg':
+        if restore_plot and plt.get_backend() == 'agg':
             plt.switch_backend(INITIAL_BACKEND)
 
         self._initialized = False
diff --git a/MLStructFP/db/image/_rect_photo.py b/MLStructFP/db/image/_rect_photo.py
index 9f4dc21..2aa3cf9 100644
--- a/MLStructFP/db/image/_rect_photo.py
+++ b/MLStructFP/db/image/_rect_photo.py
@@ -381,9 +381,12 @@ def make_region(self, xmin: NumberType, xmax: NumberType, ymin: NumberType, ymax
         :return: Returns the image index on the library array
         """
         assert xmax > xmin and ymax > ymin
+        t0 = time.time()
         dx = (xmax - xmin) / 2
         dy = (ymax - ymin) / 2
-        return self._make(floor, GeomPoint2D(xmin + dx, ymin + dy), dx, dy, rect)
+        mk = self._make(floor, GeomPoint2D(xmin + dx, ymin + dy), dx, dy, rect)
+        self._last_make_region_time = time.time() - t0
+        return mk
 
     def _make(self, floor: 'Floor', cr: 'GeomPoint2D', dx: float, dy: float, rect: Optional['Rect']
               ) -> Tuple[int, 'np.ndarray']: