Custom object embeddings with annotation uploads (#274)

Finally here! + 2 tests

Author: jihan-yin

README.md

@@ -17,12 +17,13 @@ Nucleus is a new way—the right way—to develop ML models, helping us move awa
 
 `$ pip install scale-nucleus`
 
-
 ## CLI installation
+
 We recommend installing the CLI via `pipx` (https://pypa.github.io/pipx/installation/). This makes sure that
 the CLI does not interfere with you system packages and is accessible from your favorite terminal.
 
 For MacOS:
+
 ```bash
 brew install pipx
 pipx ensurepath
@@ -32,6 +33,7 @@ nu install-completions
 ```
 
 Otherwise, install via pip (requires pip 19.0 or later):
+
 ```bash
 python3 -m pip install --user pipx
 python3 -m pipx ensurepath
@@ -45,6 +47,7 @@ nu install-completions
 ### Outdated Client
 
 Nucleus is iterating rapidly and as a result we do not always perfectly preserve backwards compatibility with older versions of the client. If you run into any unexpected error, it's a good idea to upgrade your version of the client by running
+
 ```
 pip install --upgrade scale-nucleus
 ```
@@ -87,33 +90,34 @@ poetry run pytest -m "not integration"
 
 ## Pydantic Models
 
-Prefer using [Pydantic](https://pydantic-docs.helpmanual.io/usage/models/) models rather than creating raw dictionaries 
-or dataclasses to send or receive over the wire as JSONs. Pydantic is created with data validation in mind and provides very clear error 
+Prefer using [Pydantic](https://pydantic-docs.helpmanual.io/usage/models/) models rather than creating raw dictionaries
+or dataclasses to send or receive over the wire as JSONs. Pydantic is created with data validation in mind and provides very clear error
 messages when it encounters a problem with the payload.
 
 The Pydantic model(s) should mirror the payload to send. To represent a JSON payload that looks like this:
+
 ```json
 {
   "example_json_with_info": {
-      "metadata": {
-        "frame": 0
-      },
-      "reference_id": "frame0",
-      "url": "s3://example/scale_nucleus/2021/lidar/0038711321865000.json",
-      "type": "pointcloud"
+    "metadata": {
+      "frame": 0
     },
+    "reference_id": "frame0",
+    "url": "s3://example/scale_nucleus/2021/lidar/0038711321865000.json",
+    "type": "pointcloud"
+  },
   "example_image_with_info": {
-      "metadata": {
-        "author": "Picasso"
-      },
-      "reference_id": "frame0",
-      "url": "s3://bucket/0038711321865000.jpg",
-      "type": "image"
+    "metadata": {
+      "author": "Picasso"
     },
+    "reference_id": "frame0",
+    "url": "s3://bucket/0038711321865000.jpg",
+    "type": "image"
+  }
 }
 ```
 
-Could be represented as the following structure. Note that the field names map to the JSON keys and the usage of field 
+Could be represented as the following structure. Note that the field names map to the JSON keys and the usage of field
 validators (`@validator`).
 
 ```python
@@ -161,29 +165,31 @@ parsed_model = ExampleNestedModel.parse_obj(payload.json())
 requests.post("example/post_to", json=parsed_model.dict())
 ```
 
-
 ### Migrating to Pydantic
+
 - When migrating an interface from a dictionary use `nucleus.pydantic_base.DictCompatibleModel`. That allows you to get
-the benefits of Pydantic but maintaints backwards compatibility with a Python dictionary by delegating `__getitem__` to 
-fields.
-- When migrating a frozen dataclass use `nucleus.pydantic_base.ImmutableModel`. That is a base class set up to be 
-immutable after initialization.
+  the benefits of Pydantic but maintaints backwards compatibility with a Python dictionary by delegating `__getitem__` to
+  fields.
+- When migrating a frozen dataclass use `nucleus.pydantic_base.ImmutableModel`. That is a base class set up to be
+  immutable after initialization.
 
 **Updating documentation:**
 We use [Sphinx](https://www.sphinx-doc.org/en/master/) to autogenerate our API Reference from docstrings.
 
 To test your local docstring changes, run the following commands from the repository's root directory:
+
 ```
 poetry shell
 cd docs
 sphinx-autobuild . ./_build/html --watch ../nucleus
 ```
-`sphinx-autobuild` will spin up a server on localhost (port 8000 by default) that will watch for and automatically rebuild a version of the API reference based on your local docstring changes.
 
+`sphinx-autobuild` will spin up a server on localhost (port 8000 by default) that will watch for and automatically rebuild a version of the API reference based on your local docstring changes.
 
 ## Custom Metrics using Shapely in scale-validate
 
 Certain metrics use `shapely` which is added as an optional dependency.
+
 ```bash
 pip install scale-nucleus[metrics]
 ```
@@ -199,4 +205,4 @@ apt-get install libgeos-dev
 
 To develop it locally use
 
-`poetry install --extra shapely`
+`poetry install --extras shapely`

nucleus/__init__.py

@@ -809,6 +809,7 @@ def make_request(
         Returns:
             Response payload as JSON dict.
         """
+        print(payload, route)
         if payload is None:
             payload = {}
         if requests_command is requests.get:

nucleus/annotation.py

@@ -12,6 +12,7 @@
     CATEGORY_TYPE,
     CUBOID_TYPE,
     DIMENSIONS_KEY,
+    EMBEDDING_VECTOR_KEY,
     GEOMETRY_KEY,
     HEIGHT_KEY,
     INDEX_KEY,
@@ -97,7 +98,8 @@ class BoxAnnotation(Annotation):  # pylint: disable=R0902
             height=10,
             reference_id="image_1",
             annotation_id="image_1_car_box_1",
-            metadata={"vehicle_color": "red"}
+            metadata={"vehicle_color": "red"},
+            embedding_vector=[0.1423, 1.432, ...3.829],
         )
 
     Parameters:
@@ -121,6 +123,9 @@ class BoxAnnotation(Annotation):  # pylint: disable=R0902
             attach to this annotation.  Strings, floats and ints are supported best
             by querying and insights features within Nucleus. For more details see
             our `metadata guide <https://nucleus.scale.com/docs/upload-metadata>`_.
+        embedding_vector: Custom embedding vector for this object annotation.
+            If any custom object embeddings have been uploaded previously to this dataset,
+            this vector must match the dimensions of the previously ingested vectors.
     """
 
     label: str
@@ -131,6 +136,7 @@ class BoxAnnotation(Annotation):  # pylint: disable=R0902
     reference_id: str
     annotation_id: Optional[str] = None
     metadata: Optional[Dict] = None
+    embedding_vector: Optional[list] = None
 
     def __post_init__(self):
         self.metadata = self.metadata if self.metadata else {}
@@ -149,6 +155,7 @@ def from_json(cls, payload: dict):
             reference_id=payload[REFERENCE_ID_KEY],
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
+            embedding_vector=payload.get(EMBEDDING_VECTOR_KEY, None),
         )
 
     def to_payload(self) -> dict:
@@ -164,6 +171,7 @@ def to_payload(self) -> dict:
             REFERENCE_ID_KEY: self.reference_id,
             ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
+            EMBEDDING_VECTOR_KEY: self.embedding_vector,
         }
 
 
@@ -282,7 +290,8 @@ class PolygonAnnotation(Annotation):
             vertices=[Point(100, 100), Point(150, 200), Point(200, 100)],
             reference_id="image_2",
             annotation_id="image_2_bus_polygon_1",
-            metadata={"vehicle_color": "yellow"}
+            metadata={"vehicle_color": "yellow"},
+            embedding_vector=[0.1423, 1.432, ...3.829],
         )
 
     Parameters:
@@ -298,13 +307,17 @@ class PolygonAnnotation(Annotation):
             attach to this annotation.  Strings, floats and ints are supported best
             by querying and insights features within Nucleus. For more details see
             our `metadata guide <https://nucleus.scale.com/docs/upload-metadata>`_.
+        embedding_vector: Custom embedding vector for this object annotation.
+            If any custom object embeddings have been uploaded previously to this dataset,
+            this vector must match the dimensions of the previously ingested vectors.
     """
 
     label: str
     vertices: List[Point]
     reference_id: str
     annotation_id: Optional[str] = None
     metadata: Optional[Dict] = None
+    embedding_vector: Optional[list] = None
 
     def __post_init__(self):
         self.metadata = self.metadata if self.metadata else {}
@@ -333,6 +346,7 @@ def from_json(cls, payload: dict):
             reference_id=payload[REFERENCE_ID_KEY],
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
+            embedding_vector=payload.get(EMBEDDING_VECTOR_KEY, None),
         )
 
     def to_payload(self) -> dict:
@@ -345,6 +359,7 @@ def to_payload(self) -> dict:
             REFERENCE_ID_KEY: self.reference_id,
             ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
+            EMBEDDING_VECTOR_KEY: self.embedding_vector,
         }
         return payload
 

nucleus/constants.py

@@ -41,6 +41,7 @@
 DEFAULT_ANNOTATION_UPDATE_MODE = False
 DEFAULT_NETWORK_TIMEOUT_SEC = 120
 DIMENSIONS_KEY = "dimensions"
+EMBEDDING_VECTOR_KEY = "embedding_vector"
 EMBEDDINGS_URL_KEY = "embeddings_urls"
 EMBEDDING_DIMENSION_KEY = "embedding_dimension"
 ERRORS_KEY = "errors"

nucleus/prediction.py

@@ -26,6 +26,7 @@
     CONFIDENCE_KEY,
     CUBOID_TYPE,
     DIMENSIONS_KEY,
+    EMBEDDING_VECTOR_KEY,
     GEOMETRY_KEY,
     HEIGHT_KEY,
     LABEL_KEY,
@@ -153,6 +154,9 @@ class BoxPrediction(BoxAnnotation):
             annotation. Each value should be between 0 and 1 (inclusive), and sum up to
             1 as a complete distribution. This can be useful for computing entropy to
             surface places where the model is most uncertain.
+        embedding_vectorOptional[List]): Custom embedding vector for this object annotation.
+            If any custom object embeddings have been uploaded previously to this dataset,
+            this vector must match the dimensions of the previously ingested vectors.
     """
 
     def __init__(
@@ -167,6 +171,7 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
         class_pdf: Optional[Dict] = None,
+        embedding_vector: Optional[list] = None,
     ):
         super().__init__(
             label=label,
@@ -177,6 +182,7 @@ def __init__(
             reference_id=reference_id,
             annotation_id=annotation_id,
             metadata=metadata,
+            embedding_vector=embedding_vector,
         )
         self.confidence = confidence
         self.class_pdf = class_pdf
@@ -204,6 +210,7 @@ def from_json(cls, payload: dict):
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
             class_pdf=payload.get(CLASS_PDF_KEY, None),
+            embedding_vector=payload.get(EMBEDDING_VECTOR_KEY, None),
         )
 
 
@@ -296,6 +303,9 @@ class PolygonPrediction(PolygonAnnotation):
             annotation. Each value should be between 0 and 1 (inclusive), and sum up to
             1 as a complete distribution. This can be useful for computing entropy to
             surface places where the model is most uncertain.
+        embedding_vector: Custom embedding vector for this object annotation.
+            If any custom object embeddings have been uploaded previously to this dataset,
+            this vector must match the dimensions of the previously ingested vectors.
     """
 
     def __init__(
@@ -307,13 +317,15 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
         class_pdf: Optional[Dict] = None,
+        embedding_vector: Optional[list] = None,
     ):
         super().__init__(
             label=label,
             vertices=vertices,
             reference_id=reference_id,
             annotation_id=annotation_id,
             metadata=metadata,
+            embedding_vector=embedding_vector,
         )
         self.confidence = confidence
         self.class_pdf = class_pdf
@@ -340,6 +352,7 @@ def from_json(cls, payload: dict):
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
             class_pdf=payload.get(CLASS_PDF_KEY, None),
+            embedding_vector=payload.get(EMBEDDING_VECTOR_KEY, None),
         )
 
 

tests/helpers.py

@@ -177,6 +177,13 @@ def reference_id_from_url(url):
     for i in range(len(TEST_IMG_URLS))
 ]
 
+
+TEST_BOX_ANNOTATIONS_EMBEDDINGS = [
+    {**ann, "embedding_vector": [i, i, i]}
+    for i, ann in enumerate(TEST_BOX_ANNOTATIONS)
+]
+
+
 TEST_LINE_ANNOTATIONS = [
     {
         "label": f"[Pytest] Line Annotation ${i}",
@@ -355,6 +362,20 @@ def reference_id_from_url(url):
     for i in range(len(TEST_BOX_ANNOTATIONS))
 ]
 
+TEST_BOX_PREDICTIONS_EMBEDDINGS = [
+    {
+        **TEST_BOX_ANNOTATIONS_EMBEDDINGS[i],
+        "confidence": 0.10 * i,
+        "class_pdf": TEST_BOX_MODEL_PDF,
+    }
+    if i != 0
+    else {
+        **TEST_BOX_ANNOTATIONS_EMBEDDINGS[i],
+        "confidence": 0.10 * i,
+    }
+    for i in range(len(TEST_BOX_ANNOTATIONS_EMBEDDINGS))
+]
+
 TEST_LINE_PREDICTIONS = [
     {
         **TEST_LINE_ANNOTATIONS[i],