From 4d5eaa8121b0351c8e2a9a9693d3984cfaa23fec Mon Sep 17 00:00:00 2001
From: ppizarror <pablo@ppizarror.com>
Date: Tue, 12 Mar 2024 13:10:48 -0300
Subject: [PATCH] Fix rect photo generation for non-3D images

---
 MLStructFP/__init__.py             | 2 +-
 MLStructFP/db/image/_base.py       | 4 ++++
 MLStructFP/db/image/_rect_photo.py | 9 ++++++---
 test/test_db.py                    | 1 +
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/MLStructFP/__init__.py b/MLStructFP/__init__.py
index 407c4d9..468f007 100644
--- a/MLStructFP/__init__.py
+++ b/MLStructFP/__init__.py
@@ -9,7 +9,7 @@
 __description__ = 'Machine learning structural floor plan dataset'
 __keywords__ = ['ml', 'ai', 'dataset', 'calc', 'matrix analysis', 'cnn', 'structural analysis', 'structural design']
 __email__ = 'pablo@ppizarror.com'
-__version__ = '0.5.2'
+__version__ = '0.5.3'
 
 # URL
 __url__ = 'https://github.com/MLSTRUCT/MLSTRUCT-FP'
diff --git a/MLStructFP/db/image/_base.py b/MLStructFP/db/image/_base.py
index 475c3a8..2f1e413 100644
--- a/MLStructFP/db/image/_base.py
+++ b/MLStructFP/db/image/_base.py
@@ -57,6 +57,10 @@ def __init__(self, path: str, save_images: bool, image_size_px: int) -> None:
 
         self.save = True
 
+    @property
+    def image_shape(self) -> Tuple[int, int]:
+        return self._image_size, self._image_size
+
     def make_rect(self, rect: 'Rect', crop_length: NumberType) -> Tuple[int, 'np.ndarray']:
         """
         Generate image for the perimeter of a given rectangle.
diff --git a/MLStructFP/db/image/_rect_photo.py b/MLStructFP/db/image/_rect_photo.py
index 27eda96..f5f12ea 100644
--- a/MLStructFP/db/image/_rect_photo.py
+++ b/MLStructFP/db/image/_rect_photo.py
@@ -516,8 +516,12 @@ def _get_crop_image(
         y1, y2 = min(y), max(y)
         ww = int(x2 - x1)
         hh = int(y2 - y1)
-        out_img = np.zeros((ww, hh, 3))
-        w, h, _ = image.shape  # Real image size
+        if len(image.shape) == 2:
+            w, h = image.shape
+            out_img = np.zeros((ww, hh))
+        else:
+            w, h, c = image.shape
+            out_img = np.zeros((ww, hh, c))
         dx = 0
         if x1 < 0:
             dx = - x1
@@ -560,7 +564,6 @@ def _get_crop_image(
             adjusted: 'np.ndarray' = cv2.convertScaleAbs(im, alpha=_alpha, beta=0)
         else:
             adjusted = im
-        # _swap_colors(adjusted, 0, self._empty_color)
 
         # Apply kernel
         image_kernel = cv2.filter2D(adjusted, -1, self._kernel)
diff --git a/test/test_db.py b/test/test_db.py
index f2cabbe..ba788df 100644
--- a/test/test_db.py
+++ b/test/test_db.py
@@ -108,6 +108,7 @@ def test_image(self) -> None:
 
         image_binary = RectBinaryImage(image_size_px=256).init()
         image_photo = RectFloorPhoto(image_size_px=256)
+        self.assertEqual(image_binary.image_shape, (256, 256))
 
         r = f[0].rect[3]  # Selected rectangle
         image_binary.make_rect(r)