Added method to upload local ndjson with predictions to Labelbox' GCS

Author: Grzegorz Szpak

labelbox/schema/bulk_import_request.py

@@ -1,12 +1,17 @@
 import json
+from pathlib import Path
 from typing import Iterable
 
+import ndjson
+
 from labelbox import Client
+from labelbox.exceptions import LabelboxError
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field
 from labelbox.orm.model import Relationship
 from labelbox.schema.enums import BulkImportRequestState
+from labelbox.schema.enums import UploadedFileType
 
 
 class BulkImportRequest(DbObject):
@@ -25,7 +30,8 @@ def create(
             predictions: Iterable[dict]) -> 'BulkImportRequest':
         data_str = '\n'.join(json.dumps(prediction) for prediction in predictions)
         data = data_str.encode('utf-8')
-        input_file_url = client.upload_data(data)
+        input_file_url = client.upload_data(
+            data, uploaded_file_type=UploadedFileType.PREDICTIONS)
         query_str = """
         mutation CreateBulkImportRequestPyApi {
             createBulkImportRequest(data: {
@@ -44,3 +50,27 @@ def create(
         )
         bulk_import_request_kwargs = client.execute(query_str)["createBulkImportRequest"]
         return BulkImportRequest(client, bulk_import_request_kwargs)
+
+    @staticmethod
+    def upload_local_predictions_file(
+            client: Client, local_predictions_file_path: Path) -> str:
+        """
+        Uploads local NDJSON file containing predictions to Labelbox' object store
+        and returns a URL of created file.
+
+        Args:
+            client (Client): The Labelbox client
+            local_predictions_file_path (str): local NDJSON file containing predictions
+        Returns:
+            A URL of uploaded NDJSON file
+        Raises:
+            LabelboxError: if local file is not a valid NDJSON file
+        """
+        with local_predictions_file_path.open("rb") as f:
+            try:
+                data = ndjson.load(f)
+                return client.upload_data(
+                    data, uploaded_file_type=UploadedFileType.PREDICTIONS)
+            except ValueError:
+                raise LabelboxError(
+                    f"File {local_predictions_file_path} is not a valid ndjson file")

requirements.txt

@@ -1 +1,2 @@
 requests==2.22.0
+ndjson==0.3.1

setup.py

@@ -15,7 +15,10 @@
     long_description_content_type="text/markdown",
     url="https://labelbox.com",
     packages=setuptools.find_packages(),
-    install_requires=["requests>=2.22.0"],
+    install_requires=[
+        "ndjson>=0.3.1",
+        "requests>=2.22.0"
+    ],
     classifiers=[
         'Development Status :: 3 - Alpha',
         'License :: OSI Approved :: Apache Software License',