Fix conversion for rect photo

MLSTRUCT · Nov 21, 2023 · fef383b · fef383b
1 parent 0715281
commit fef383b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git 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__ = '[email protected]'
-__version__ = '0.3.6'
+__version__ = '0.3.7'
 
 # URL
 __url__ = 'https://github.com/MLSTRUCT/MLSTRUCT-FP'

diff --git a/MLStructFP/db/image/_rect_photo.py b/MLStructFP/db/image/_rect_photo.py
@@ -315,14 +315,16 @@ def _get_floor_image(self, floor: 'Floor') -> Tuple['np.ndarray', 'GeomPoint2D']
         # Make default empty color
         pixels: 'np.ndarray'
         if self._empty_color >= 0:
-            image = cv2.imread(ip, cv2.IMREAD_UNCHANGED)
-
-            # make mask of where the transparent bits are
-            trans_mask = image[:, :, 3] == 0
+            image: 'np.ndarray' = cv2.imread(ip, cv2.IMREAD_UNCHANGED)
+            if image.shape[2] == 4:
+                # make mask of where the transparent bits are
+                trans_mask = image[:, :, 3] == 0
 
-            # replace areas of transparency with white and not transparent
-            image[trans_mask] = [self._empty_color, self._empty_color, self._empty_color, 255]
-            pixels = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
+                # replace areas of transparency with white and not transparent
+                image[trans_mask] = [self._empty_color, self._empty_color, self._empty_color, 255]
+                pixels = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
+            else:
+                pixels = 255 - image
         else:
             pixels = cv2.imread(ip)
 
@@ -331,7 +333,9 @@ def _get_floor_image(self, floor: 'Floor') -> Tuple['np.ndarray', 'GeomPoint2D']
                 image = cv2.imread(ip, cv2.IMREAD_UNCHANGED)
                 trans_mask = image[:, :, 3] == 0
                 image[trans_mask] = [255, 255, 255, 255]  # Turn all black to white to invert
-                pixels = 255 - cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)  # Invert colors
+            else:
+                image = pixels
+            pixels = 255 - cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)  # Invert colors
 
         # Flip image
         if floor.mutator_scale_x < 0: