Add basic metadata update for scenes/items (#216)

Author: jean-lucas

nucleus/constants.py

@@ -59,7 +59,6 @@
 INDEX_CONTINUOUS_ENABLE_KEY = "enable"
 ITEMS_KEY = "items"
 ITEM_KEY = "item"
-ITEMS_KEY = "items"
 ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 JOB_ID_KEY = "job_id"
 KEEP_HISTORY_KEY = "keep_history"

nucleus/dataset.py

@@ -59,6 +59,7 @@
 from .dataset_item_uploader import DatasetItemUploader
 from .deprecation_warning import deprecated
 from .errors import DatasetItemRetrievalError
+from .metadata_manager import ExportMetadataType, MetadataManager
 from .payload_constructor import (
     construct_append_scenes_payload,
     construct_model_run_creation_payload,
@@ -1392,3 +1393,45 @@ def _upload_items(
 
         populator = DatasetItemUploader(self.id, self._client)
         return populator.upload(dataset_items, batch_size, update)
+
+    def update_scene_metadata(self, mapping: Dict[str, dict]):
+        """
+        Update (merge) scene metadata for each reference_id given in the mapping.
+        The backed will join the specified mapping metadata to the exisiting metadata.
+        If there is a key-collision, the value given in the mapping will take precedence.
+
+        Args:
+            mapping: key-value pair of <reference_id>: <metadata>
+
+        Examples:
+            >>> mapping = {"scene_ref_1": {"new_key": "foo"}, "scene_ref_2": {"some_value": 123}}
+            >>> dataset.update_scene_metadata(mapping)
+
+        Returns:
+            A dictionary outlining success or failures.
+        """
+        mm = MetadataManager(
+            self.id, self._client, mapping, ExportMetadataType.SCENES
+        )
+        return mm.update()
+
+    def update_item_metadata(self, mapping: Dict[str, dict]):
+        """
+        Update (merge) dataset item metadata for each reference_id given in the mapping.
+        The backed will join the specified mapping metadata to the exisiting metadata.
+        If there is a key-collision, the value given in the mapping will take precedence.
+
+        Args:
+            mapping: key-value pair of <reference_id>: <metadata>
+
+        Examples:
+            >>> mapping = {"item_ref_1": {"new_key": "foo"}, "item_ref_2": {"some_value": 123}}
+            >>> dataset.update_item_metadata(mapping)
+
+        Returns:
+            A dictionary outlining success or failures.
+        """
+        mm = MetadataManager(
+            self.id, self._client, mapping, ExportMetadataType.DATASET_ITEMS
+        )
+        return mm.update()

nucleus/metadata_manager.py

@@ -0,0 +1,45 @@
+from enum import Enum
+from typing import TYPE_CHECKING, Dict
+
+if TYPE_CHECKING:
+    from . import NucleusClient
+
+
+# Wording set to match with backend enum
+class ExportMetadataType(Enum):
+    SCENES = "scene"
+    DATASET_ITEMS = "item"
+
+
+class MetadataManager:
+    """
+    Helper class for managing metadata updates on a scene or dataset item.
+    Do not call directly, use the dataset class methods: `update_scene_metadata` or `update_item_metadata`
+    """
+
+    def __init__(
+        self,
+        dataset_id: str,
+        client: "NucleusClient",
+        raw_mappings: Dict[str, dict],
+        level: ExportMetadataType,
+    ):
+        self.dataset_id = dataset_id
+        self._client = client
+        self.raw_mappings = raw_mappings
+        self.level = level
+
+        self._payload = self._format_mappings()
+
+    def _format_mappings(self):
+        payload = []
+        for ref_id, meta in self.raw_mappings.items():
+            payload.append({"reference_id": ref_id, "metadata": meta})
+        return payload
+
+    def update(self):
+        payload = {"metadata": self._payload, "level": self.level.value}
+        resp = self._client.make_request(
+            payload=payload, route=f"dataset/{self.dataset_id}/metadata"
+        )
+        return resp

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.6.1"
+version = "0.6.2"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/test_dataset.py

@@ -1,16 +1,9 @@
 import copy
 import math
-import os
 
 import pytest
 
-from nucleus import (
-    Dataset,
-    DatasetItem,
-    NucleusAPIError,
-    NucleusClient,
-    UploadResponse,
-)
+from nucleus import Dataset, DatasetItem, NucleusClient, UploadResponse
 from nucleus.annotation import (
     BoxAnnotation,
     CategoryAnnotation,
@@ -34,8 +27,6 @@
     UPDATED_ITEMS,
 )
 from nucleus.job import AsyncJob, JobError
-from nucleus.model import Model
-from nucleus.prediction import BoxPrediction
 
 from .helpers import (
     LOCAL_FILENAME,
@@ -556,3 +547,22 @@ def sort_labelmap(segmentation_annotation):
         exported[0][ANNOTATIONS_KEY][MULTICATEGORY_TYPE][0]
         == multicategory_annotation
     )
+
+
+def test_dataset_item_metadata_update(dataset):
+    items = make_dataset_items()
+    dataset.append(items)
+
+    expected_metadata = {}
+    new_metadata = {}
+    for item in dataset.items:
+        data = {"a_new_key": 123}
+        new_metadata[item.reference_id] = data
+        expected_metadata[item.reference_id] = {**item.metadata, **data}
+
+    dataset.update_item_metadata(new_metadata)
+    actual_metadata = {
+        item.reference_id: item.metadata for item in dataset.items
+    }
+
+    assert actual_metadata == expected_metadata

tests/test_scene.py

@@ -477,3 +477,25 @@ def test_scene_upload_async_item_dataset(dataset_item):
 
     with pytest.raises(Exception):
         dataset_item.append(scenes, update=update, asynchronous=True)
+
+
+@pytest.mark.integration
+def test_scene_metadata_update(dataset_scene):
+    payload = TEST_LIDAR_SCENES
+    scenes = [
+        LidarScene.from_json(scene_json) for scene_json in payload[SCENES_KEY]
+    ]
+    update = payload[UPDATE_KEY]
+
+    job = dataset_scene.append(scenes, update=update, asynchronous=True)
+    job.sleep_until_complete()
+
+    scene_ref_id = scenes[0].reference_id
+    additional_metadata = {"some_new_key": 123}
+    dataset_scene.update_scene_metadata({scene_ref_id: additional_metadata})
+
+    expected_new_metadata = {**scenes[0].metadata, **additional_metadata}
+
+    updated_scene = dataset_scene.get_scene(scene_ref_id)
+    actual_metadata = updated_scene.metadata
+    assert expected_new_metadata == actual_metadata