From ce554792951b753cd4fe25301de54ec4e1a7fe23 Mon Sep 17 00:00:00 2001
From: ppizarror <pablo@ppizarror.com>
Date: Tue, 27 Feb 2024 08:12:27 -0800
Subject: [PATCH] Added check for image load error prevention

---
 MLStructFP_benchmarks/utils/_fp_db.py              | 8 +++++++-
 MLStructFP_benchmarks/utils/_fp_patch_generator.py | 4 ++--
 setup.py                                           | 2 +-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/MLStructFP_benchmarks/utils/_fp_db.py b/MLStructFP_benchmarks/utils/_fp_db.py
index a897c9c..375f033 100644
--- 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
index 41652af..f83a13b 100644
--- 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
index 0fe33c7..2b2f56b 100644
--- 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',