Allow embedding info upload for Dataset Items (#405)

Author: ntamas92

CHANGELOG.md

@@ -5,6 +5,11 @@ 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.16.7](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.7) - 2023-11-03
+
+### Added
+- Allow direct embedding vector upload together with dataset items. `DatasetItem` now has an additional parameter called `embedding_info` which can be used to directly upload embeddings when a dataset is uploaded.
+
 
 ## [0.16.6](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.6) - 2023-11-01
 

nucleus/constants.py

@@ -48,6 +48,7 @@
 DEFAULT_ANNOTATION_UPDATE_MODE = False
 DEFAULT_NETWORK_TIMEOUT_SEC = 120
 DIMENSIONS_KEY = "dimensions"
+EMBEDDING_INFO_KEY = "embedding_info"
 EMBEDDING_VECTOR_KEY = "embedding_vector"
 EMBEDDINGS_URL_KEY = "embeddings_urls"
 EMBEDDING_DIMENSION_KEY = "embedding_dimension"
@@ -70,6 +71,7 @@
 IMAGE_LOCATION_KEY = "image_location"
 IMAGE_URL_KEY = "image_url"
 INDEX_KEY = "index"
+INDEX_ID_KEY = "index_id"
 INDEX_CONTINUOUS_ENABLE_KEY = "enable"
 IOU_KEY = "iou"
 ITEMS_KEY = "items"

nucleus/dataset_item.py

@@ -10,7 +10,10 @@
 from .constants import (
     BACKEND_REFERENCE_ID_KEY,
     CAMERA_PARAMS_KEY,
+    EMBEDDING_INFO_KEY,
+    EMBEDDING_VECTOR_KEY,
     IMAGE_URL_KEY,
+    INDEX_ID_KEY,
     METADATA_KEY,
     ORIGINAL_IMAGE_URL_KEY,
     POINTCLOUD_URL_KEY,
@@ -26,6 +29,18 @@ class DatasetItemType(Enum):
     POINTCLOUD = "pointcloud"
 
 
+@dataclass
+class DatasetItemEmbeddingInfo:
+    index_id: str
+    embedding_vector: list
+
+    def to_payload(self) -> dict:
+        return {
+            INDEX_ID_KEY: self.index_id,
+            EMBEDDING_VECTOR_KEY: self.embedding_vector,
+        }
+
+
 @dataclass  # pylint: disable=R0902
 class DatasetItem:  # pylint: disable=R0902
     """A dataset item is an image or pointcloud that has associated metadata.
@@ -113,16 +128,23 @@ class DatasetItem:  # pylint: disable=R0902
     metadata: Optional[dict] = None
     pointcloud_location: Optional[str] = None
     upload_to_scale: Optional[bool] = True
+    embedding_info: Optional[DatasetItemEmbeddingInfo] = None
 
     def __post_init__(self):
         assert self.reference_id != "DUMMY_VALUE", "reference_id is required."
         assert bool(self.image_location) != bool(
             self.pointcloud_location
         ), "Must specify exactly one of the image_location or pointcloud_location parameters"
+        if self.pointcloud_location and self.embedding_info:
+            raise AssertionError(
+                "Cannot upload embedding vector if pointcloud_location is set"
+            )
+
         if (self.pointcloud_location) and not self.upload_to_scale:
             raise NotImplementedError(
                 "Skipping upload to Scale is not currently implemented for pointclouds."
             )
+
         self.local = (
             is_local_path(self.image_location) if self.image_location else None
         )
@@ -179,6 +201,9 @@ def to_payload(self, is_scene=False) -> dict:
 
         payload[REFERENCE_ID_KEY] = self.reference_id
 
+        if self.embedding_info:
+            payload[EMBEDDING_INFO_KEY] = self.embedding_info.to_payload()
+
         if is_scene:
             if self.image_location:
                 payload[URL_KEY] = self.image_location

pyproject.toml

@@ -25,7 +25,7 @@ ignore = ["E501", "E741", "E731", "F401"]  # Easy ignore for getting it running
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.16.6"
+version = "0.16.7"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]