Support cloud-storage URIs in Pascal VOC semantic segmentation input

src/labelformat/formats/semantic_segmentation/pascalvoc.py

@@ -9,12 +9,14 @@
 from __future__ import annotations
 
 import json
+import posixpath
 from argparse import ArgumentParser
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 
 import numpy as np
+from fsspec.core import url_to_fs
 from numpy.typing import NDArray
 from PIL import Image as PILImage
 from PIL import ImageDraw
@@ -32,6 +34,7 @@
 )
 from labelformat.model.multipolygon import MultiPolygon
 from labelformat.model.semantic_segmentation import SemanticSegmentationMask
+from labelformat.types import PathLike
 
 """TODO(Malte, 11/2025):
 Support what is already supported in LightlyTrain. https://docs.lightly.ai/train/stable/semantic_segmentation.html#data
@@ -46,8 +49,8 @@
 class PascalVOCSemanticSegmentationInput(InstanceSegmentationInput):
     """Pascal VOC semantic segmentation input format."""
 
-    _images_dir: Path
-    _masks_dir: Path
+    _images_dir: PathLike
+    _masks_dir: PathLike
     _filename_to_image: dict[str, Image]
     _categories: list[Category]
 
@@ -59,24 +62,28 @@ def add_cli_arguments(parser: ArgumentParser) -> None:
     @classmethod
     def from_dirs(
         cls,
-        images_dir: Path,
-        masks_dir: Path,
+        images_dir: PathLike,
+        masks_dir: PathLike,
         class_id_to_name: Mapping[int, str],
     ) -> "PascalVOCSemanticSegmentationInput":
         """Create a PascalVOCSemanticSegmentationInput from directory pairs.
 
         Args:
-            images_dir: Root directory containing images (nested structure allowed).
-            masks_dir: Root directory containing PNG masks mirroring images structure.
+            images_dir: Root directory or URI containing images (nested structure allowed).
+            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.
 
         Raises:
             ValueError: If directories are invalid, a mask is missing or not PNG,
                         or if class_id keys cannot be parsed as integers.
         """
-        if not images_dir.is_dir():
+        images_fs, images_fs_dir = url_to_fs(str(images_dir))
+        if not images_fs.isdir(images_fs_dir):
             raise ValueError(f"Images directory is not a directory: {images_dir}")
-        if not masks_dir.is_dir():
+
+        masks_fs, masks_fs_dir = url_to_fs(str(masks_dir))
+        if not masks_fs.isdir(masks_fs_dir):
             raise ValueError(f"Masks directory is not a directory: {masks_dir}")
 
         # Build categories from mapping
@@ -87,10 +94,12 @@ 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):
-            mask_path = masks_dir / Path(img.filename).with_suffix(".png")
-            if not mask_path.is_file():
+            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):
                 raise ValueError(
-                    f"Missing mask PNG for image '{img.filename}' at path: {mask_path}"
+                    f"Missing mask PNG for image '{img.filename}' at path: "
+                    f"{masks_fs.unstrip_protocol(mask_path)}"
                 )
             images_by_filename[img.filename] = img
 
@@ -138,15 +147,19 @@ def _get_mask(self, image_filepath: str) -> SemanticSegmentationMask:
                 f"Unknown image filepath {image_filepath}. Use one returned by get_images()."
             )
 
-        mask_path = self._masks_dir / Path(image_filepath).with_suffix(".png")
-        if not mask_path.is_file():
+        masks_fs, masks_fs_dir = url_to_fs(str(self._masks_dir))
+        mask_rel_path = str(PurePosixPath(image_filepath).with_suffix(".png"))
+        mask_path = posixpath.join(masks_fs_dir, mask_rel_path)
+        if not masks_fs.isfile(mask_path):
             raise ValueError(
-                f"Mask PNG not found for image '{image_filepath}': {mask_path}"
+                f"Mask PNG not found for image '{image_filepath}': "
+                f"{masks_fs.unstrip_protocol(mask_path)}"
             )
 
         # Load and validate mask by shape and value set.
-        with PILImage.open(mask_path) as mimg:
-            mask_np: NDArray[np.int_] = np.asarray(mimg, dtype=np.int_)
+        with masks_fs.open(mask_path, mode="rb") as mask_file:
+            with PILImage.open(mask_file) as mimg:
+                mask_np: NDArray[np.int_] = np.asarray(mimg, dtype=np.int_)
         _validate_mask(
             image_obj=image_obj,
             mask_np=mask_np,

tests/unit/formats/semantic_segmentation/test_pascalvoc.py

@@ -4,8 +4,10 @@
 from argparse import ArgumentParser
 from pathlib import Path
 from typing import Dict, Iterable
+from uuid import uuid4
 
 import cv2
+import fsspec
 import numpy as np
 import pytest
 from PIL import Image as PILImage
@@ -262,6 +264,33 @@ def test_get_labels(self, tmp_path: Path) -> None:
             [1, 1],
         ]
 
+    def test_from_dirs__supports_memory_uri(self) -> None:
+        mapping = _load_class_mapping_int_keys()
+        root_uri = f"memory://{uuid4().hex}/pascalvoc"
+        images_uri = f"{root_uri}/JPEGImages"
+        masks_uri = f"{root_uri}/SegmentationClass"
+
+        for image_filename in ["2007_000032.jpg", "subdir/2007_000033.jpg"]:
+            with fsspec.open(f"{images_uri}/{image_filename}", "wb") as file:
+                file.write((IMAGES_DIR / image_filename).read_bytes())
+            mask_filename = str(Path(image_filename).with_suffix(".png"))
+            with fsspec.open(f"{masks_uri}/{mask_filename}", "wb") as file:
+                file.write((MASKS_DIR / mask_filename).read_bytes())
+
+        ds = PascalVOCSemanticSegmentationInput.from_dirs(
+            images_dir=images_uri,
+            masks_dir=masks_uri,
+            class_id_to_name=mapping,
+        )
+
+        images = list(ds.get_images())
+        assert len(images) == 2
+        for image in images:
+            mask = ds._get_mask(image.filename)
+            assert sum(run_length for _, run_length in mask.category_id_rle) == (
+                image.width * image.height
+            )
+
 
 class TestPascalVOCSemanticSegmentationOutput:
     def test_save__writes_fixture_masks_and_class_mapping(self, tmp_path: Path) -> None: