Merge pull request #462 from Labelbox/modelrunlabels

[AL-1623] - ModelRun Export Annotations

Author: jtsodapop

labelbox/schema/model_run.py

@@ -1,16 +1,21 @@
-from typing import TYPE_CHECKING, Dict, Iterable, Union
+from typing import TYPE_CHECKING, Dict, Iterable, Union, List, Optional, Any
 from pathlib import Path
 import os
 import time
+import logging
+import requests
+import ndjson
 
 from labelbox.pagination import PaginatedCollection
 from labelbox.orm.query import results_query_part
 from labelbox.orm.model import Field, Relationship, Entity
-from labelbox.orm.db_object import DbObject
+from labelbox.orm.db_object import DbObject, experimental
 
 if TYPE_CHECKING:
     from labelbox import MEAPredictionImport
 
+logger = logging.getLogger(__name__)
+
 
 class ModelRun(DbObject):
     name = Field.String("name")
@@ -175,6 +180,54 @@ def delete_model_run_data_rows(self, data_row_ids):
             data_row_ids_param: data_row_ids
         })
 
+    @experimental
+    def export_labels(
+        self,
+        download: bool = False,
+        timeout_seconds: int = 600
+    ) -> Optional[Union[str, List[Dict[Any, Any]]]]:
+        """
+        Experimental. To use, make sure client has enable_experimental=True.
+
+        Fetches Labels from the ModelRun
+
+        Args:
+            download (bool): Returns the url if False
+        Returns:
+            URL of the data file with this ModelRun's labels.
+            If download=True, this instead returns the contents as NDJSON format.
+            If the server didn't generate during the `timeout_seconds` period, 
+            None is returned.
+        """
+        sleep_time = 2
+        query_str = """mutation exportModelRunAnnotationsPyApi($modelRunId: ID!) {
+                exportModelRunAnnotations(data: {modelRunId: $modelRunId}) {
+                    downloadUrl createdAt status
+                }
+            }
+            """
+
+        while True:
+            url = self.client.execute(
+                query_str, {'modelRunId': self.uid},
+                experimental=True)['exportModelRunAnnotations']['downloadUrl']
+
+            if url:
+                if not download:
+                    return url
+                else:
+                    response = requests.get(url)
+                    response.raise_for_status()
+                    return ndjson.loads(response.content)
+
+            timeout_seconds -= sleep_time
+            if timeout_seconds <= 0:
+                return None
+
+            logger.debug("ModelRun '%s' label export, waiting for server...",
+                         self.uid)
+            time.sleep(sleep_time)
+
 
 class ModelRunDataRow(DbObject):
     label_id = Field.String("label_id")

tests/integration/annotation_import/test_model_run.py

@@ -73,3 +73,8 @@ def test_model_run_upsert_data_rows_with_existing_labels(
     ])
     assert n_data_rows == len(
         list(model_run_with_model_run_data_rows.model_run_data_rows()))
+
+
+def test_model_run_export_labels(model_run_with_model_run_data_rows):
+    labels = model_run_with_model_run_data_rows.export_labels(download=True)
+    assert len(labels) == 3