Add opt-in on_error hook for unreadable images during folder scan

src/labelformat/formats/kitti.py

@@ -1,7 +1,7 @@
 import logging
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import Iterable
+from typing import Iterable, Optional
 
 from labelformat import utils
 from labelformat.cli.registry import Task, cli_register
@@ -49,6 +49,9 @@ def __init__(
     ) -> None:
         self._input_folder = input_folder
         self._images_rel_path = images_rel_path
+        # Optional hook to tolerate unreadable images during folder scan.
+        # Default None re-raises; see utils.get_images_from_folder.
+        self.on_error: Optional[utils.OnImageErrorHook] = None
         self._categories = [
             Category(id=idx, name=name)
             for idx, name in enumerate(category_names.split(","))
@@ -59,7 +62,8 @@ def get_categories(self) -> Iterable[Category]:
 
     def get_images(self) -> Iterable[Image]:
         yield from utils.get_images_from_folder(
-            folder=self._input_folder / self._images_rel_path
+            folder=self._input_folder / self._images_rel_path,
+            on_error=self.on_error,
         )
 
     def get_labels(self) -> Iterable[ImageObjectDetection]:

src/labelformat/formats/lightly.py

@@ -2,7 +2,7 @@
 import json
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence
 
 from labelformat import utils
 from labelformat.cli.registry import Task, cli_register
@@ -49,14 +49,18 @@ def __init__(
         self._input_folder = input_folder
         self._images_rel_path = images_rel_path
         self._skip_labels_without_image = skip_labels_without_image
+        # Optional hook to tolerate unreadable images during folder scan.
+        # Default None re-raises; see utils.get_images_from_folder.
+        self.on_error: Optional[utils.OnImageErrorHook] = None
         self._categories = self._get_categories()
 
     def get_categories(self) -> Iterable[Category]:
         return self._categories
 
     def get_images(self) -> Iterable[Image]:
         yield from utils.get_images_from_folder(
-            folder=self._input_folder / self._images_rel_path
+            folder=self._input_folder / self._images_rel_path,
+            on_error=self.on_error,
         )
 
     def get_labels(self) -> Iterable[ImageObjectDetection]:

src/labelformat/formats/maskpair.py

@@ -4,7 +4,7 @@
 
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import Iterable, List, Literal, Union
+from typing import Iterable, List, Literal, Optional, Union
 
 from labelformat.cli.registry import Task, cli_register
 from labelformat.mask_utils import (
@@ -23,7 +23,11 @@
     SingleInstanceSegmentation,
 )
 from labelformat.model.multipolygon import MultiPolygon
-from labelformat.utils import get_image_dimensions
+from labelformat.utils import (
+    ImageDimensionError,
+    OnImageErrorHook,
+    get_image_dimensions,
+)
 
 
 @cli_register(format="maskpair", task=Task.INSTANCE_SEGMENTATION)
@@ -128,6 +132,9 @@ def __init__(
         self._morph_close = morph_close
         self._segmentation_type = segmentation_type
         self._approx_epsilon = approx_epsilon
+        # Optional hook to tolerate unreadable images during folder scan.
+        # Default None re-raises; see utils.get_images_from_folder.
+        self.on_error: Optional[OnImageErrorHook] = None
 
         # Parse category names
         self._categories = []
@@ -147,9 +154,20 @@ def get_categories(self) -> Iterable[Category]:
         yield from self._categories
 
     def get_images(self) -> Iterable[Image]:
-        """Get images from the image/mask pairs."""
+        """Get images from the image/mask pairs.
+
+        Image ids are the positional index into the image/mask pairs so that
+        they stay aligned with ``get_labels()``. Images skipped via ``on_error``
+        leave a gap in the id sequence.
+        """
         for image_id, (image_path, _) in enumerate(self._image_mask_pairs):
-            width, height = get_image_dimensions(image_path)
+            try:
+                width, height = get_image_dimensions(image_path)
+            except ImageDimensionError as error:
+                if self.on_error is None:
+                    raise
+                self.on_error(Path(image_path), error)
+                continue
             yield Image(
                 id=image_id,
                 filename=image_path.name,
@@ -162,6 +180,11 @@ def get_labels(self) -> Iterable[ImageInstanceSegmentation]:
         images = {img.id: img for img in self.get_images()}
 
         for image_id, (image_path, mask_path) in enumerate(self._image_mask_pairs):
+            # Skip pairs whose image was dropped by the on_error hook in
+            # get_images() so labels stay aligned with readable images.
+            if image_id not in images:
+                continue
+
             # Load and binarize the mask
             binary_mask = binarize_mask(
                 mask_path=mask_path,

src/labelformat/formats/semantic_segmentation/pascalvoc.py

@@ -14,6 +14,7 @@
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
 from pathlib import Path, PurePosixPath
+from typing import Optional
 
 import numpy as np
 from fsspec.core import url_to_fs
@@ -65,6 +66,7 @@ def from_dirs(
         images_dir: PathLike,
         masks_dir: PathLike,
         class_id_to_name: Mapping[int, str],
+        on_error: Optional[utils.OnImageErrorHook] = None,
     ) -> "PascalVOCSemanticSegmentationInput":
         """Create a PascalVOCSemanticSegmentationInput from directory pairs.
 
@@ -73,6 +75,10 @@ def from_dirs(
             masks_dir: Root directory or URI containing PNG masks mirroring images
                 structure.
             class_id_to_name: Mapping of class_id -> class name, with integer keys.
+            on_error: Optional hook invoked when an image cannot be read during the
+                folder scan. When ``None`` (the default), a failure raises
+                ``ImageDimensionError``. When provided, the hook is called with the
+                offending path and the error, and the image is skipped.
 
         Raises:
             ValueError: If directories are invalid, a mask is missing or not PNG,
@@ -93,7 +99,7 @@ def from_dirs(
 
         # Collect images using helper and ensure a PNG mask exists for each.
         images_by_filename: dict[str, Image] = {}
-        for img in utils.get_images_from_folder(images_dir):
+        for img in utils.get_images_from_folder(images_dir, on_error=on_error):
             mask_rel_path = str(PurePosixPath(img.filename).with_suffix(".png"))
             mask_path = posixpath.join(masks_fs_dir, mask_rel_path)
             if not masks_fs.isfile(mask_path):

src/labelformat/formats/yolov8.py

@@ -1,7 +1,7 @@
 import logging
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import Iterable, List
+from typing import Iterable, List, Optional
 
 import yaml
 
@@ -41,6 +41,9 @@ def add_cli_arguments(parser: ArgumentParser) -> None:
     def __init__(self, input_file: Path, input_split: str) -> None:
         self._config_file = input_file
         self._split = input_split
+        # Optional hook to tolerate unreadable images during folder scan.
+        # Default None re-raises; see utils.get_images_from_folder.
+        self.on_error: Optional[utils.OnImageErrorHook] = None
         with self._config_file.open() as file:
             self._config_data = yaml.safe_load(file)
 
@@ -56,7 +59,9 @@ def get_categories(self) -> Iterable[Category]:
             yield Category(id=category_id, name=names[category_id])
 
     def get_images(self) -> Iterable[Image]:
-        yield from utils.get_images_from_folder(folder=self._images_dir())
+        yield from utils.get_images_from_folder(
+            folder=self._images_dir(), on_error=self.on_error
+        )
 
     def _root_dir(self) -> Path:
         """Return the root directory of the dataset.

src/labelformat/utils.py

@@ -1,7 +1,7 @@
 import logging
 import posixpath
 from pathlib import Path, PurePosixPath
-from typing import Iterable, Tuple
+from typing import Callable, Iterable, Optional, Tuple
 from urllib.parse import urlsplit
 
 import fsspec
@@ -46,9 +46,21 @@
 
 
 class ImageDimensionError(Exception):
-    """Raised when unable to extract image dimensions using fast methods."""
+    """Raised when unable to extract image dimensions.
+
+    Carries the offending image path (when available) so callers and error
+    hooks can identify which file failed.
+    """
+
+    def __init__(self, message: str, path: Optional[PathLike] = None) -> None:
+        super().__init__(message)
+        self.path = path
 
-    pass
+
+# Hook invoked when an image cannot be read during a folder scan. Receives the
+# offending path and the raised ImageDimensionError. Returning normally skips
+# the file; raising propagates.
+OnImageErrorHook = Callable[[Path, ImageDimensionError], None]
 
 
 def _path_suffix(path: PathLike) -> str:
@@ -177,18 +189,32 @@ def get_image_dimensions(image_path: PathLike) -> Tuple[int, int]:
         except ImageDimensionError:
             pass
 
-    with fsspec.open(str(image_path), "rb") as img_file:
-        with PIL.Image.open(img_file) as img:
-            return img.size
-
-
-def get_images_from_folder(folder: PathLike) -> Iterable[Image]:
+    try:
+        with fsspec.open(str(image_path), "rb") as img_file:
+            with PIL.Image.open(img_file) as img:
+                return img.size
+    except (PIL.UnidentifiedImageError, OSError) as error:
+        raise ImageDimensionError(
+            f"Failed to read image dimensions for '{image_path}': {error}",
+            path=image_path,
+        ) from error
+
+
+def get_images_from_folder(
+    folder: PathLike,
+    on_error: Optional[OnImageErrorHook] = None,
+) -> Iterable[Image]:
     """Yields an Image structure for all images in the given folder.
 
     The order of the images is arbitrary. Images in nested folders are included.
 
     Args:
         folder: Path or URI to the folder containing images.
+        on_error: Optional hook invoked when an image cannot be read. When
+            ``None`` (the default), a failure raises ``ImageDimensionError`` and
+            aborts iteration. When provided, the hook is called with the
+            offending path and the error, and the file is skipped so iteration
+            continues.
     """
     image_id = 0
     logger.debug(f"Listing images in '{folder}'...")
@@ -205,7 +231,13 @@ def get_images_from_folder(folder: PathLike) -> Iterable[Image]:
             continue
         image_filename = _relative_path(path=image_path, root=fs_folder)
         image_uri = fs.unstrip_protocol(image_path)
-        image_width, image_height = get_image_dimensions(image_uri)
+        try:
+            image_width, image_height = get_image_dimensions(image_uri)
+        except ImageDimensionError as error:
+            if on_error is None:
+                raise
+            on_error(Path(image_uri), error)
+            continue
         yield Image(
             id=image_id,
             filename=image_filename,

tests/integration/test_maskpair_integration.py

@@ -2,7 +2,7 @@
 
 import json
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, List, Tuple
 
 import cv2
 import numpy as np
@@ -11,6 +11,7 @@
 
 from labelformat.formats.coco import COCOInstanceSegmentationOutput
 from labelformat.formats.maskpair import MaskPairInstanceSegmentationInput
+from labelformat.utils import ImageDimensionError
 
 
 def create_test_data(base_path: Path) -> None:
@@ -265,3 +266,59 @@ def test_empty_category_fallback(self, tmp_path: Path) -> None:
         # Should create a default category
         assert len(categories) > 0
         # Default category should be created in the get_labels method
+
+    def test_get_images__reraises_by_default(self, tmp_path: Path) -> None:
+        maskpair_input = _make_on_error_input(tmp_path)
+        with pytest.raises(ImageDimensionError):
+            list(maskpair_input.get_images())
+
+    def test_get_images__on_error_hook_skips_broken_image(self, tmp_path: Path) -> None:
+        maskpair_input = _make_on_error_input(tmp_path)
+        errors: List[Tuple[Path, ImageDimensionError]] = []
+        maskpair_input.on_error = lambda path, error: errors.append((path, error))
+
+        images = list(maskpair_input.get_images())
+        assert [img.filename for img in images] == ["good.jpg"]
+        assert len(errors) == 1
+        assert Path(errors[0][0]).name == "broken.jpg"
+        assert isinstance(errors[0][1], ImageDimensionError)
+
+    def test_get_labels__on_error_hook_skips_broken_image(self, tmp_path: Path) -> None:
+        maskpair_input = _make_on_error_input(tmp_path)
+        maskpair_input.on_error = lambda path, error: None
+
+        labels = list(maskpair_input.get_labels())
+        assert [label.image.filename for label in labels] == ["good.jpg"]
+
+
+def _make_data_with_broken_image(base_path: Path) -> None:
+    """Create one valid image/mask pair and one broken image with a valid mask."""
+    images_dir = base_path / "images"
+    masks_dir = base_path / "masks"
+    images_dir.mkdir(parents=True)
+    masks_dir.mkdir(parents=True)
+
+    # One valid image/mask pair.
+    image = np.random.randint(100, 200, (100, 100, 3), dtype=np.uint8)
+    cv2.imwrite(str(images_dir / "good.jpg"), image)
+    good_mask: NDArray[np.uint8] = np.zeros((100, 100), dtype=np.uint8)
+    good_mask[20:40, 20:40] = 255
+    cv2.imwrite(str(masks_dir / "good.png"), good_mask)
+
+    # One broken image with a valid mask (so the pair is matched).
+    (images_dir / "broken.jpg").write_bytes(b"not a valid jpeg")
+    broken_mask: NDArray[np.uint8] = np.zeros((100, 100), dtype=np.uint8)
+    broken_mask[10:30, 10:30] = 255
+    cv2.imwrite(str(masks_dir / "broken.png"), broken_mask)
+
+
+def _make_on_error_input(tmp_path: Path) -> MaskPairInstanceSegmentationInput:
+    _make_data_with_broken_image(tmp_path)
+    return MaskPairInstanceSegmentationInput(
+        image_glob="images/*.jpg",
+        mask_glob="masks/*.png",
+        base_path=tmp_path,
+        pairing_mode="stem",
+        category_names="object",
+        min_area=10.0,
+    )