Support Cloud Storage for COCO (#66)

Author: MalteEbner

README.md

@@ -68,6 +68,28 @@ This package is compatible with the following platforms:
 pip install labelformat
 ```
 
+### ☁️ Using Cloud Storage
+
+To work with annotations stored in cloud storage (like AWS S3,
+GCS, or Azure), install the cloud storage dependencies:
+
+```shell
+pip install "labelformat[cloud-storage]"
+```
+
+This installs the required libraries: `s3fs` (for S3), `gcsfs` (for GCS), and
+`adlfs` (for Azure).
+
+Labelformat uses `fsspec`, which also supports other file systems. If you need a
+different provider (for example FTP or SSH), check the
+[fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#other-known-implementations)
+and install the matching implementation manually (for example `pip install sftpfs`).
+
+**Current Support Limitations:**
+
+- **Input format:** Cloud URIs are currently supported only for COCO input (`--input-format coco`).
+
+
 ## Usage
 
 ### CLI

pyproject.toml

@@ -13,11 +13,19 @@ dependencies = [
   "pydantic==2.10.6",
   "pydantic-xml==2.17.3",
   "numpy",
+  "fsspec>=2023.1.0",
 ]
 
 [project.scripts]
 labelformat = "labelformat.cli.cli:main"
 
+[project.optional-dependencies]
+cloud-storage = [
+  "s3fs>=2023.1.0",
+  "gcsfs>=2023.1.0",
+  "adlfs>=2023.1.0",
+]
+
 [dependency-groups]
 dev = [
   "mypy==1.14.1",
@@ -60,6 +68,10 @@ strict = true
 disallow_untyped_defs = true
 ignore_missing_imports = false
 
+[[tool.mypy.overrides]]
+module = "fsspec.*"
+ignore_missing_imports = true
+
 # Print format
 show_error_codes = true
 show_error_context = true

src/labelformat/formats/coco.py

@@ -5,6 +5,8 @@
 from pathlib import Path
 from typing import Dict, Iterable, List
 
+import fsspec
+
 import labelformat.formats.coco_segmentation_helpers as segmentation_helpers
 from labelformat.cli.registry import Task, cli_register
 from labelformat.formats.coco_segmentation_helpers import (
@@ -31,21 +33,21 @@
     ObjectDetectionOutput,
     SingleObjectDetection,
 )
-from labelformat.types import JsonDict, ParseError
+from labelformat.types import JsonDict, ParseError, PathLike
 
 
 class _COCOBaseInput:
     @staticmethod
     def add_cli_arguments(parser: ArgumentParser) -> None:
         parser.add_argument(
             "--input-file",
-            type=Path,
+            type=str,
             required=True,
-            help="Path to input COCO JSON file",
+            help="Path or URI to input COCO JSON file",
         )
 
-    def __init__(self, input_file: Path) -> None:
-        with input_file.open() as file:
+    def __init__(self, input_file: PathLike) -> None:
+        with fsspec.open(str(input_file), mode="r") as file:
             self._data = json.load(file)
 
     def get_categories(self) -> Iterable[Category]:

src/labelformat/types.py

@@ -1,7 +1,11 @@
-from typing import Any, Dict
+from pathlib import Path
+from typing import Any, Dict, Union
 
 JsonDict = Dict[str, Any]
 
 
 class ParseError(Exception):
     pass
+
+
+PathLike = Union[str, Path]

tests/integration/test_coco_cloud_storage__unmocked.py

@@ -0,0 +1,31 @@
+import pytest
+from fsspec.core import url_to_fs
+
+from labelformat.formats.coco import COCOObjectDetectionInput
+
+COCO_S3_ANNOTATION_FILE = (
+    "s3://studio-test-datasets-eu/coco_subset_128_images/instances_train2017.json"
+)
+COCO_S3_IMAGES_DIR = "s3://studio-test-datasets-eu/coco_subset_128_images/images/"
+
+
+@pytest.mark.skip(reason="Requires access to S3 dataset")
+def test_coco_od_inputs_read_from_s3__unmocked() -> None:
+    object_detection_input = COCOObjectDetectionInput(
+        input_file=COCO_S3_ANNOTATION_FILE
+    )
+
+    object_detection_categories = list(object_detection_input.get_categories())
+    object_detection_images = list(object_detection_input.get_images())
+    object_detection_labels = list(object_detection_input.get_labels())
+
+    assert len(object_detection_categories) == 80
+    assert len(object_detection_images) == 128
+    assert len(object_detection_labels) == 128
+    assert sum(len(label.objects) for label in object_detection_labels) == 900
+
+    # Assert that the images found in the annotations can be accessed in S3
+    for image in object_detection_images:
+        image_path = COCO_S3_IMAGES_DIR + image.filename
+        fs, fs_path = url_to_fs(image_path)
+        assert fs.exists(fs_path)

tests/unit/formats/test_coco_cloud_storage.py

@@ -0,0 +1,28 @@
+import io
+import json
+
+from pytest_mock import MockerFixture
+
+from labelformat.formats.coco import COCOObjectDetectionInput
+
+COCO_S3_ANNOTATION_FILE = "s3://some_bucket/some_file.json"
+MOCK_COCO_PAYLOAD = {
+    "categories": [{"id": 1, "name": "cat"}],
+    "images": [],
+    "annotations": [],
+}
+
+
+def test_coco_od_inputs_read_from_s3__mocked(mocker: MockerFixture) -> None:
+    mock_open = mocker.patch(
+        "labelformat.formats.coco.fsspec.open",
+        return_value=io.StringIO(json.dumps(MOCK_COCO_PAYLOAD)),
+    )
+    object_detection_input = COCOObjectDetectionInput(
+        input_file=COCO_S3_ANNOTATION_FILE
+    )
+    object_detection_categories = list(object_detection_input.get_categories())
+
+    assert len(object_detection_categories) == 1
+    assert object_detection_categories[0].name == "cat"
+    mock_open.assert_called_once_with(COCO_S3_ANNOTATION_FILE, mode="r")