diff --git a/src/labelformat/formats/semantic_segmentation/pascalvoc.py b/src/labelformat/formats/semantic_segmentation/pascalvoc.py index 025ccf6..0e20c00 100644 --- a/src/labelformat/formats/semantic_segmentation/pascalvoc.py +++ b/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, diff --git a/tests/unit/formats/semantic_segmentation/test_pascalvoc.py b/tests/unit/formats/semantic_segmentation/test_pascalvoc.py index aa2a04f..9a2d1b8 100644 --- a/tests/unit/formats/semantic_segmentation/test_pascalvoc.py +++ b/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: