[PLT-630] Add support for export embeddings from SDK (#1573)

Author: attila-papai

libs/labelbox/src/labelbox/schema/catalog.py

@@ -26,7 +26,7 @@ def export_v2(
     ) -> Union[Task, ExportTask]:
         """
         Creates a catalog export task with the given params, filters and returns the task.
-        
+
         >>>     import labelbox as lb
         >>>     client = lb.Client(<API_KEY>)
         >>>     catalog = client.get_catalog()
@@ -98,6 +98,7 @@ def _export(self,
 
         _params = params or CatalogExportParams({
             "attachments": False,
+            "embeddings": False,
             "metadata_fields": False,
             "data_row_details": False,
             "project_details": False,
@@ -142,6 +143,8 @@ def _export(self,
                         if media_type_override is not None else None,
                     "includeAttachments":
                         _params.get('attachments', False),
+                    "includeEmbeddings":
+                        _params.get('embeddings', False),
                     "includeMetadata":
                         _params.get('metadata_fields', False),
                     "includeDataRowDetails":

libs/labelbox/src/labelbox/schema/data_row.py

@@ -266,6 +266,7 @@ def _export(
     ) -> Tuple[Task, bool]:
         _params = params or CatalogExportParams({
             "attachments": False,
+            "embeddings": False,
             "metadata_fields": False,
             "data_row_details": False,
             "project_details": False,
@@ -325,6 +326,8 @@ def _export(
                         if media_type_override is not None else None,
                     "includeAttachments":
                         _params.get('attachments', False),
+                    "includeEmbeddings":
+                        _params.get('embeddings', False),
                     "includeMetadata":
                         _params.get('metadata_fields', False),
                     "includeDataRowDetails":

libs/labelbox/src/labelbox/schema/dataset.py

@@ -717,6 +717,7 @@ def _export(
     ) -> Tuple[Task, bool]:
         _params = params or CatalogExportParams({
             "attachments": False,
+            "embeddings": False,
             "metadata_fields": False,
             "data_row_details": False,
             "project_details": False,
@@ -763,6 +764,8 @@ def _export(
                         if media_type_override is not None else None,
                     "includeAttachments":
                         _params.get('attachments', False),
+                    "includeEmbeddings":
+                        _params.get('embeddings', False),
                     "includeMetadata":
                         _params.get('metadata_fields', False),
                     "includeDataRowDetails":

libs/labelbox/src/labelbox/schema/export_params.py

@@ -15,6 +15,7 @@ class DataRowParams(TypedDict):
     data_row_details: Optional[bool]
     metadata_fields: Optional[bool]
     attachments: Optional[bool]
+    embeddings: Optional[bool]
     media_type_override: Optional[MediaType]
 
 

libs/labelbox/src/labelbox/schema/model_run.py

@@ -57,7 +57,7 @@ def upsert_labels(self,
                       label_ids: Optional[List[str]] = None,
                       project_id: Optional[str] = None,
                       timeout_seconds=3600):
-        """ 
+        """
         Adds data rows and labels to a Model Run
 
         Args:
@@ -273,16 +273,16 @@ def add_predictions(
         name: str,
         predictions: Union[str, Path, Iterable[Dict], Iterable["Label"]],
     ) -> 'MEAPredictionImport':  # type: ignore
-        """ 
+        """
         Uploads predictions to a new Editor project.
-        
+
         Args:
             name (str): name of the AnnotationImport job
             predictions (str or Path or Iterable): url that is publicly accessible by Labelbox containing an
                 ndjson file
                 OR local path to an ndjson file
                 OR iterable of annotation rows
-        
+
         Returns:
             AnnotationImport
         """
@@ -566,6 +566,8 @@ def _export(
                         _params.get('media_type_override', None),
                     "includeAttachments":
                         _params.get('attachments', False),
+                    "includeEmbeddings":
+                        _params.get('embeddings', False),
                     "includeMetadata":
                         _params.get('metadata_fields', False),
                     "includeDataRowDetails":

libs/labelbox/src/labelbox/schema/project.py

@@ -518,6 +518,7 @@ def _export(
     ) -> Tuple[Task, bool]:
         _params = params or ProjectExportParams({
             "attachments": False,
+            "embeddings": False,
             "metadata_fields": False,
             "data_row_details": False,
             "project_details": False,
@@ -560,6 +561,8 @@ def _export(
                         if media_type_override is not None else None,
                     "includeAttachments":
                         _params.get('attachments', False),
+                    "includeEmbeddings":
+                        _params.get('embeddings', False),
                     "includeMetadata":
                         _params.get('metadata_fields', False),
                     "includeDataRowDetails":

libs/labelbox/src/labelbox/schema/slice.py

@@ -167,6 +167,7 @@ def _export(
     ) -> Tuple[Task, bool]:
         _params = params or CatalogExportParams({
             "attachments": False,
+            "embeddings": False,
             "metadata_fields": False,
             "data_row_details": False,
             "project_details": False,
@@ -201,6 +202,8 @@ def _export(
                         if media_type_override is not None else None,
                     "includeAttachments":
                         _params.get('attachments', False),
+                    "includeEmbeddings":
+                        _params.get('embeddings', False),
                     "includeMetadata":
                         _params.get('metadata_fields', False),
                     "includeDataRowDetails":

libs/labelbox/tests/conftest.py

@@ -78,7 +78,7 @@ class Environ(Enum):
 
 @pytest.fixture
 def image_url() -> str:
-    return IMAGE_URL
+    return MASKABLE_IMG_URL
 
 
 @pytest.fixture
@@ -376,14 +376,6 @@ def client(environ: str):
     return IntegrationClient(environ)
 
 
-@pytest.fixture(scope="session")
-def image_url(client):
-    return client.upload_data(requests.get(MASKABLE_IMG_URL).content,
-                              content_type="image/jpeg",
-                              filename="image.jpeg",
-                              sign=True)
-
-
 @pytest.fixture(scope="session")
 def pdf_url(client):
     pdf_url = client.upload_file('tests/assets/loremipsum.pdf')
@@ -1042,4 +1034,12 @@ def configured_project_with_complex_ontology(client, initial_dataset, rand_gen,
     project.setup(editor, ontology.asdict())
 
     yield [project, data_row]
-    project.delete()
+    project.delete()
+
+
+@pytest.fixture
+def embedding(client: Client):
+    uuid_str = uuid.uuid4().hex
+    embedding = client.create_embedding(f"sdk-int-{uuid_str}", 8)
+    yield embedding
+    embedding.delete()

libs/labelbox/tests/data/export/streamable/test_export_embeddings_streamable.py

@@ -0,0 +1,73 @@
+import json
+import random
+
+from labelbox import StreamType, JsonConverter
+
+
+class TestExportEmbeddings:
+
+    def test_export_embeddings_precomputed(self, client, dataset, environ,
+                                           image_url):
+        data_row_specs = [{
+            "row_data": image_url,
+            "external_id": "image",
+        }]
+        task = dataset.create_data_rows(data_row_specs)
+        task.wait_till_done()
+        export_task = dataset.export(params={"embeddings": True})
+        export_task.wait_till_done()
+        assert export_task.status == "COMPLETE"
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+
+        results = []
+        export_task.get_stream(converter=JsonConverter(),
+                               stream_type=StreamType.RESULT).start(
+                                   stream_handler=lambda output: results.append(
+                                       json.loads(output.json_str)))
+
+        assert len(results) == len(data_row_specs)
+
+        result = results[0]
+        assert "embeddings" in result
+        assert len(result["embeddings"]) > 0
+        assert result["embeddings"][0][
+            "name"] == "Image Embedding V2 (CLIP ViT-B/32)"
+        assert len(result["embeddings"][0]["values"]) == 1
+
+    def test_export_embeddings_custom(self, client, dataset, image_url,
+                                      embedding):
+        vector = [random.uniform(1.0, 2.0) for _ in range(embedding.dims)]
+        import_task = dataset.create_data_rows([{
+            "row_data": image_url,
+            "embeddings": [{
+                "embedding_id": embedding.id,
+                "vector": vector,
+            }],
+        }])
+        import_task.wait_till_done()
+        assert import_task.status == "COMPLETE"
+
+        export_task = dataset.export(params={"embeddings": True})
+        export_task.wait_till_done()
+        assert export_task.status == "COMPLETE"
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+
+        results = []
+        export_task.get_stream(converter=JsonConverter(),
+                               stream_type=StreamType.RESULT).start(
+                                   stream_handler=lambda output: results.append(
+                                       json.loads(output.json_str)))
+
+        assert len(results) == 1
+        assert "embeddings" in results[0]
+        assert (len(results[0]["embeddings"])
+                >= 1)  # should at least contain the custom embedding
+        for emb in results[0]["embeddings"]:
+            if emb["id"] == embedding.id:
+                assert emb["name"] == embedding.name
+                assert emb["dimensions"] == embedding.dims
+                assert emb["is_custom"] == True
+                assert len(emb["values"]) == 1
+                assert emb["values"][0]["value"] == vector

libs/labelbox/tests/integration/test_embedding.py

@@ -1,7 +1,6 @@
 import json
 import random
 import threading
-import uuid
 from tempfile import NamedTemporaryFile
 from typing import List, Dict, Any
 
@@ -12,14 +11,6 @@
 from labelbox.schema.embedding import Embedding
 
 
-@pytest.fixture
-def embedding(client: Client):
-    uuid_str = uuid.uuid4().hex
-    embedding = client.create_embedding(f"sdk-int-{uuid_str}", 8)
-    yield embedding
-    embedding.delete()
-
-
 def test_get_embedding_by_id(client: Client, embedding: Embedding):
     e = client.get_embedding_by_id(embedding.id)
     assert e.id == embedding.id
@@ -43,7 +34,7 @@ def test_get_embeddings(client: Client, embedding: Embedding):
 @pytest.mark.parametrize('data_rows', [10], indirect=True)
 def test_import_vectors_from_file(data_rows: List[DataRow],
                                   embedding: Embedding):
-    vector = [random.uniform(1.0, 2.0) for _ in range(8)]
+    vector = [random.uniform(1.0, 2.0) for _ in range(embedding.dims)]
     event = threading.Event()
 
     def callback(_: Dict[str, Any]):
@@ -66,7 +57,7 @@ def callback(_: Dict[str, Any]):
 def test_get_imported_vector_count(dataset: Dataset, embedding: Embedding):
     assert embedding.get_imported_vector_count() == 0
 
-    vector = [random.uniform(1.0, 2.0) for _ in range(8)]
+    vector = [random.uniform(1.0, 2.0) for _ in range(embedding.dims)]
     dataset.create_data_row(row_data="foo",
                             embeddings=[{
                                 "embedding_id": embedding.id,