Add create_slice_by_ids (#383)

Author: Anthony Krivonos

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.15.4](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.4) - 2023-03-21
+
+### Changed
+- Added `create_slice_by_ids` to create slices from dataset item, scene, and object IDs
+
+
 ## [0.15.3](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.3) - 2023-03-02
 
 ### Changed

nucleus/constants.py

@@ -38,6 +38,7 @@
 DATASET_ID_KEY = "dataset_id"
 DATASET_IS_SCENE_KEY = "is_scene"
 DATASET_ITEM_ID_KEY = "dataset_item_id"
+DATASET_ITEM_IDS_KEY = "dataset_item_ids"
 DATASET_ITEMS_KEY = "dataset_items"
 DATASET_LENGTH_KEY = "length"
 DATASET_MODEL_RUNS_KEY = "model_run_ids"
@@ -104,6 +105,7 @@
 PAGE_TOKEN_KEY = "pageToken"
 NEXT_TOKEN_KEY = "nextPageToken"
 OVERWRITE_KEY = "overwrite"
+OBJECT_IDS_KEY = "object_ids"
 P1_KEY = "p1"
 P2_KEY = "p2"
 POINTCLOUD_KEY = "pointcloud"
@@ -120,6 +122,7 @@
 SCALE_TASK_INFO_KEY = "scale_task_info"
 SCENE_KEY = "scene"
 SCENES_KEY = "scenes"
+SCENE_IDS_KEY = "scene_ids"
 SERIALIZED_REQUEST_KEY = "serialized_request"
 SEGMENTATIONS_KEY = "segmentations"
 SLICE_ID_KEY = "slice_id"

nucleus/dataset.py

@@ -35,6 +35,7 @@
     BACKFILL_JOB_KEY,
     DATASET_ID_KEY,
     DATASET_IS_SCENE_KEY,
+    DATASET_ITEM_IDS_KEY,
     DATASET_ITEMS_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EMBEDDING_DIMENSION_KEY,
@@ -47,8 +48,10 @@
     KEEP_HISTORY_KEY,
     MESSAGE_KEY,
     NAME_KEY,
+    OBJECT_IDS_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    SCENE_IDS_KEY,
     SLICE_ID_KEY,
     TRACK_REFERENCE_IDS_KEY,
     TRACKS_KEY,
@@ -895,6 +898,51 @@ def create_slice(
         )
         return Slice(response[SLICE_ID_KEY], self._client)
 
+    def create_slice_by_ids(
+        self,
+        name: str,
+        dataset_item_ids: Optional[List[str]] = None,
+        scene_ids: Optional[List[str]] = None,
+        annotation_ids: Optional[List[str]] = None,
+        prediction_ids: Optional[List[str]] = None,
+    ) -> Slice:
+        """Creates a :class:`Slice` of dataset items, scenes, annotations, or predictions within a dataset by their IDs.
+
+        .. note::
+            Dataset item, scene, and object (annotation or prediction) IDs may not be mixed.
+            However, when creating an object slice, both annotation and prediction IDs may be supplied.
+
+        Parameters:
+            name: A human-readable name for the slice.
+            dataset_item_ids: List of internal IDs of dataset items to add to the slice::
+            scene_ids: List of internal IDs of scenes to add to the slice::
+            annotation_ids: List of internal IDs of Annotations to add to the slice::
+            prediction_ids: List of internal IDs of Predictions to add to the slice::
+
+        Returns:
+            :class:`Slice`: The newly constructed slice item.
+        """
+        if (
+            not dataset_item_ids
+            and not scene_ids
+            and not annotation_ids
+            and not prediction_ids
+        ):
+            raise Exception("Must provide at least one list of internal IDs")
+        payload = {
+            NAME_KEY: name,
+            DATASET_ITEM_IDS_KEY: dataset_item_ids
+            if dataset_item_ids
+            else None,
+            SCENE_IDS_KEY: scene_ids if scene_ids else None,
+            OBJECT_IDS_KEY: [*(annotation_ids or []), *(prediction_ids or [])]
+            or None,
+        }
+        response = self._client.make_request(
+            payload, f"dataset/{self.id}/create_slice"
+        )
+        return Slice(response[SLICE_ID_KEY], self._client)
+
     def build_slice(
         self,
         name: str,

pyproject.toml

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