Fix rect photo generation for non-3D images

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__ = '[email protected]'
-__version__ = '0.5.2'
+__version__ = '0.5.3'
 
 # URL
 __url__ = 'https://github.com/MLSTRUCT/MLSTRUCT-FP'

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.

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)

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)