Add summaries for Label Extension

Author: duckontheweb

pystac/extensions/label.py

@@ -4,13 +4,13 @@
 """
 
 from enum import Enum
-from pystac.extensions.base import ExtensionManagementMixin
+from pystac.extensions.base import ExtensionManagementMixin, SummariesExtension
 from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
 
 import pystac
 from pystac.serialization.identify import STACJSONDescription, STACVersionID
 from pystac.extensions.hooks import ExtensionHooks
-from pystac.utils import get_required
+from pystac.utils import get_required, map_opt
 
 SCHEMA_URI = "https://stac-extensions.github.io/label/v1.0.0/schema.json"
 
@@ -703,6 +703,83 @@ def ext(cls, obj: pystac.Item, add_if_missing: bool = False) -> "LabelExtension"
                 f"Label extension does not apply to type {type(obj)}"
             )
 
+    @staticmethod
+    def summaries(obj: pystac.Collection) -> "SummariesLabelExtension":
+        """Returns the extended summaries object for the given collection."""
+        return SummariesLabelExtension(obj)
+
+
+class SummariesLabelExtension(SummariesExtension):
+    """A concrete implementation of :class:`~SummariesExtension` that extends
+    the ``summaries`` field of a :class:`~pystac.Collection` to include properties
+    defined in the :stac-ext:`Label Extension <label>`.
+    """
+
+    @property
+    def label_properties(self) -> Optional[List[str]]:
+        """Get or sets the summary of :attr:`LabelExtension.label_properties` values
+        for this Collection.
+        """
+
+        return self.summaries.get_list(PROPERTIES_PROP)
+
+    @label_properties.setter
+    def label_properties(self, v: Optional[List[LabelClasses]]) -> None:
+        self._set_summary(PROPERTIES_PROP, v)
+
+    @property
+    def label_classes(self) -> Optional[List[LabelClasses]]:
+        """Get or sets the summary of :attr:`LabelExtension.label_classes` values
+        for this Collection.
+        """
+
+        return map_opt(
+            lambda classes: [LabelClasses(c) for c in classes],
+            self.summaries.get_list(CLASSES_PROP),
+        )
+
+    @label_classes.setter
+    def label_classes(self, v: Optional[List[LabelClasses]]) -> None:
+        self._set_summary(
+            CLASSES_PROP, map_opt(lambda classes: [c.to_dict() for c in classes], v)
+        )
+
+    @property
+    def label_type(self) -> Optional[List[LabelType]]:
+        """Get or sets the summary of :attr:`LabelExtension.label_type` values
+        for this Collection.
+        """
+
+        return self.summaries.get_list(TYPE_PROP)
+
+    @label_type.setter
+    def label_type(self, v: Optional[List[LabelType]]) -> None:
+        self._set_summary(TYPE_PROP, v)
+
+    @property
+    def label_tasks(self) -> Optional[List[Union[LabelTask, str]]]:
+        """Get or sets the summary of :attr:`LabelExtension.label_tasks` values
+        for this Collection.
+        """
+
+        return self.summaries.get_list(TASKS_PROP)
+
+    @label_tasks.setter
+    def label_tasks(self, v: Optional[List[Union[LabelTask, str]]]) -> None:
+        self._set_summary(TASKS_PROP, v)
+
+    @property
+    def label_methods(self) -> Optional[List[str]]:
+        """Get or sets the summary of :attr:`LabelExtension.label_methods` values
+        for this Collection.
+        """
+
+        return self.summaries.get_list(METHODS_PROP)
+
+    @label_methods.setter
+    def label_methods(self, v: Optional[List[str]]) -> None:
+        self._set_summary(METHODS_PROP, v)
+
 
 class LabelExtensionHooks(ExtensionHooks):
     schema_uri: str = SCHEMA_URI

tests/data-files/label/spacenet-roads/roads_collection.json

@@ -0,0 +1,55 @@
+{
+    "stac_version": "1.0.0-rc.1",
+    "type": "Collection",
+    "id": "spacenet-roads-sample",
+    "description": "A sample of the SpaceNet Roads dataset built during STAC Sprint 4. The dataset contains hand-labeled roads.",
+    "keywords": [
+        "spacenet",
+        "roads",
+        "labels"
+    ],
+    "license": "CC-BY-SA-4.0",
+    "providers": [
+        {
+            "name": "SpaceNet",
+            "roles": [
+                "licensor",
+                "host",
+                "producer",
+                "processor"
+            ],
+            "url": "https://spacenet.ai"
+        }
+    ],
+    "extent": {
+        "spatial": {
+            "bbox": [
+                [
+                    2.23379639995,
+                    49.0178709,
+                    2.23730639995,
+                    49.0213809
+                ]
+            ]
+        },
+        "temporal": {
+            "interval": [
+                [
+                    "2016-08-26T22:41:55.000000Z",
+                    null
+                ]
+            ]
+        }
+    },
+    "links": [
+        {
+            "href": "roads_collection.json",
+            "rel": "root",
+            "title": "sample SpaceNet roads label collection"
+        },
+        {
+            "rel": "item",
+            "href": "roads_item.json"
+        }
+    ]
+}

tests/extensions/test_label.py

@@ -5,13 +5,14 @@
 from typing import List, Union
 
 import pystac
-from pystac import Catalog, Item, CatalogType
+from pystac import Catalog, Collection, Item, CatalogType
 from pystac.extensions.label import (
     LabelExtension,
     LabelClasses,
     LabelCount,
     LabelOverview,
     LabelStatistics,
+    LabelTask,
     LabelType,
     LabelRelType,
 )
@@ -439,3 +440,85 @@ def test_ext_add_to(self) -> None:
         _ = LabelExtension.ext(item, add_if_missing=True)
 
         self.assertIn(LabelExtension.get_schema_uri(), item.stac_extensions)
+
+
+class LabelSummariesTest(unittest.TestCase):
+    EXAMPLE_COLLECTION = TestCases.get_path(
+        "data-files/label/spacenet-roads/roads_collection.json"
+    )
+
+    def test_label_properties_summary(self) -> None:
+        label_properties = ["road_type", "lane_number", "paved"]
+        collection = Collection.from_file(self.EXAMPLE_COLLECTION)
+        label_ext_summaries = LabelExtension.summaries(collection)
+
+        label_ext_summaries.label_properties = label_properties
+
+        summaries = collection.summaries
+        assert summaries is not None
+        label_properties_summary = summaries.get_list("label:properties")
+        assert label_properties_summary is not None
+        self.assertListEqual(label_properties, label_properties_summary)
+
+        label_properties_summary_ext = label_ext_summaries.label_properties
+        assert label_properties_summary_ext is not None
+        self.assertListEqual(label_properties, label_properties_summary_ext)
+
+    def test_label_classes_summary(self) -> None:
+        label_classes = [
+            LabelClasses(
+                {"name": "road_type", "classes": ["1", "2", "3", "4", "5", "6"]}
+            ),
+            LabelClasses({"name": "lane_number", "classes": ["1", "2", "3", "4", "5"]}),
+            LabelClasses({"name": "paved", "classes": ["0", "1"]}),
+        ]
+        collection = Collection.from_file(self.EXAMPLE_COLLECTION)
+        label_ext_summaries = LabelExtension.summaries(collection)
+
+        label_ext_summaries.label_classes = label_classes
+
+        summaries = collection.summaries
+        assert summaries is not None
+        label_classes_summary = summaries.get_list("label:classes")
+        assert label_classes_summary is not None
+        self.assertListEqual(
+            [lc.to_dict() for lc in label_classes], label_classes_summary
+        )
+
+        label_classes_summary_ext = label_ext_summaries.label_classes
+        assert label_classes_summary_ext is not None
+        self.assertListEqual(label_classes, label_classes_summary_ext)
+
+    def test_label_type_summary(self) -> None:
+        label_types = [LabelType.VECTOR]
+        collection = Collection.from_file(self.EXAMPLE_COLLECTION)
+        label_ext_summaries = LabelExtension.summaries(collection)
+
+        label_ext_summaries.label_type = label_types
+
+        summaries = collection.summaries
+        assert summaries is not None
+        label_type_summary = summaries.get_list("label:type")
+        assert label_type_summary is not None
+        self.assertListEqual(label_types, label_type_summary)
+
+        label_type_summary_ext = label_ext_summaries.label_type
+        assert label_type_summary_ext is not None
+        self.assertListEqual(label_types, label_type_summary_ext)
+
+    def test_label_task_summary(self) -> None:
+        label_tasks: List[Union[LabelTask, str]] = [LabelTask.REGRESSION]
+        collection = Collection.from_file(self.EXAMPLE_COLLECTION)
+        label_ext_summaries = LabelExtension.summaries(collection)
+
+        label_ext_summaries.label_tasks = label_tasks
+
+        summaries = collection.summaries
+        assert summaries is not None
+        label_tasks_summary = summaries.get_list("label:tasks")
+        assert label_tasks_summary is not None
+        self.assertListEqual(label_tasks, label_tasks_summary)
+
+        label_tasks_summary_ext = label_ext_summaries.label_tasks
+        assert label_tasks_summary_ext is not None
+        self.assertListEqual(label_tasks, label_tasks_summary_ext)