Added check for image load error prevention

MLSTRUCT · Feb 27, 2024 · ce55479 · ce55479
1 parent 8a9b440
commit ce55479
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/MLStructFP_benchmarks/utils/_fp_db.py b/MLStructFP_benchmarks/utils/_fp_db.py
@@ -11,6 +11,9 @@
 from MLStructFP.utils import make_dirs
 from MLStructFP_benchmarks.utils._fp_patch_generator import FloorPatchGenerator
 
+# noinspection PyProtectedMember
+from MLStructFP.db.image._rect_photo import RectFloorPhotoFileLoadException
+
 import datetime
 import functools
 import gc
@@ -175,7 +178,10 @@ def _process_floor(
             if verbose:
                 print(f'\tGenerating patches with angle={angle} ... ', end='')
             floor.mutate(angle)
-            self._gen.process(floor)
+            try:
+                self._gen.process(floor)
+            except RectFloorPhotoFileLoadException:
+                print(f'\tSkipping floor ID {floor.id} as its image "{floor.image_path}" could not be loaded')
             added += self._gen._test_last_added
             ignored += len(self._gen._test_ignored_patches)
             if verbose:

diff --git a/MLStructFP_benchmarks/utils/_fp_patch_generator.py b/MLStructFP_benchmarks/utils/_fp_patch_generator.py
@@ -14,7 +14,7 @@
 
 from MLStructFP.db.image import RectBinaryImage, RectFloorPhoto
 # noinspection PyProtectedMember
-from MLStructFP.db.image._rect_photo import RectFloorShapeException
+from MLStructFP.db.image._rect_photo import RectFloorPhotoShapeException
 from MLStructFP.utils import make_dirs, DEFAULT_PLOT_DPI, configure_figure
 
 import gc
@@ -204,7 +204,7 @@ def process(self, floor: 'Floor') -> 'FloorPatchGenerator':
             patch_b = self._gen_binary.make_region(xmin, xmax, ymin, ymax, floor)[1]
             try:
                 patch_p = self._process_photo(xmin, xmax, ymin, ymax, floor)
-            except RectFloorShapeException:  # Try to crop a larger plan
+            except RectFloorPhotoShapeException:  # Try to crop a larger plan
                 ignore = True
                 patch_p = patch_b
             sb, sp = np.sum(patch_b), np.sum(patch_p)

diff --git a/setup.py b/setup.py
@@ -18,7 +18,7 @@
     'Keras <= 2.3.1',
     'keras_tqdm <= 2.0.1',
     'matplotlib <= 3.5.3',
-    'MLStructFP >= 0.5.1',
+    'MLStructFP >= 0.5.2',
     'numpy <= 1.18.5',
     'Pillow >= 10.2.0',
     'scikit-image <= 0.18.1',