From 0834128bdabd942ccb236673da67cc1d452072d3 Mon Sep 17 00:00:00 2001
From: Viet-Anh Nguyen <vietanh@nrl.ai>
Date: Sun, 30 Aug 2026 09:31:25 +0700
Subject: [PATCH] fix: skip invalid images during SAM preload

---
 anylabeling/services/auto_labeling/model.py |  6 +++-
 tests/test_model_image_loading.py           | 31 +++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_model_image_loading.py

diff --git a/anylabeling/services/auto_labeling/model.py b/anylabeling/services/auto_labeling/model.py
index 4f89ff1..0e11c28 100644
--- a/anylabeling/services/auto_labeling/model.py
+++ b/anylabeling/services/auto_labeling/model.py
@@ -115,7 +115,7 @@ def unload(self):
 
     @staticmethod
     def load_image_from_filename(filename):
-        """Load image from labeling file and return image data and image path."""
+        """Load an image for background inference, or return ``None``."""
         label_file = os.path.splitext(filename)[0] + ".json"
         if QFile.exists(label_file) and LabelFile.is_label_file(label_file):
             try:
@@ -126,9 +126,13 @@ def load_image_from_filename(filename):
             image_data = label_file.image_data
         else:
             image_data = LabelFile.load_image_file(filename)
+        if not image_data:
+            logging.error(f"Error reading {filename}")
+            return None
         image = QImage.fromData(image_data)
         if image.isNull():
             logging.error(f"Error reading {filename}")
+            return None
         return image
 
     def on_next_files_changed(self, next_files):
diff --git a/tests/test_model_image_loading.py b/tests/test_model_image_loading.py
new file mode 100644
index 0000000..e0472f0
--- /dev/null
+++ b/tests/test_model_image_loading.py
@@ -0,0 +1,31 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+from PyQt6.QtGui import QImage
+
+from anylabeling.services.auto_labeling.model import Model
+
+
+class TestModelImageLoading(unittest.TestCase):
+    def test_corrupt_image_returns_none(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path = Path(temp_dir) / "corrupt.png"
+            path.write_bytes(b"not an image")
+
+            self.assertIsNone(Model.load_image_from_filename(str(path)))
+
+    def test_valid_image_returns_non_null_qimage(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path = Path(temp_dir) / "valid.png"
+            image = QImage(2, 2, QImage.Format.Format_RGB32)
+            self.assertTrue(image.save(str(path)))
+
+            loaded = Model.load_image_from_filename(str(path))
+
+            self.assertIsInstance(loaded, QImage)
+            self.assertFalse(loaded.isNull())
+
+
+if __name__ == "__main__":
+    unittest.main()