ADD: batch creation on project

Author: gdj0nes

labelbox/__init__.py

@@ -1,14 +1,15 @@
 name = "labelbox"
 __version__ = "3.16.0"
 
-from labelbox.schema.project import Project
 from labelbox.client import Client
+from labelbox.schema.project import Project
 from labelbox.schema.model import Model
 from labelbox.schema.bulk_import_request import BulkImportRequest
 from labelbox.schema.annotation_import import MALPredictionImport, MEAPredictionImport, LabelImport
 from labelbox.schema.dataset import Dataset
 from labelbox.schema.data_row import DataRow
 from labelbox.schema.label import Label
+from labelbox.schema.batch import Batch, BatchPriority
 from labelbox.schema.review import Review
 from labelbox.schema.user import User
 from labelbox.schema.organization import Organization

labelbox/schema/__init__.py

@@ -18,4 +18,5 @@
 import labelbox.schema.user
 import labelbox.schema.webhook
 import labelbox.schema.data_row_metadata
+import labelbox.schema.batch
 import labelbox.schema.iam_integration

labelbox/schema/batch.py

@@ -1,6 +1,10 @@
+from typing import NewType
+
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field, Relationship
 
+BatchPriority = NewType('BatchPriority', int)
+
 
 class Batch(DbObject):
     """ A Batch is a group of data rows submitted to a project for labeling

labelbox/schema/project.py

@@ -569,6 +569,60 @@ def setup(self, labeling_frontend, labeling_frontend_options) -> None:
         timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
         self.update(setup_complete=timestamp)
 
+    def create_batch(self, name: str, data_rows: List[Union[str, "DataRow"]], priority: "BatchPriority" = 5):
+        """Create a new batch for a project
+
+        Args:
+            name: a name for the batch, must be unique within a project
+            data_rows: Either a list of `DataRows` or Data Row ids
+            priority: An optional priority for the Data Rows in the Batch. 1 highest -> 5 lowest
+
+        """
+
+        # @TODO: make this automatic?
+        if self.queue_mode() != QueueMode.Batch:
+            raise ValueError("Project must be in batch mode")
+
+        dr_ids = []
+        for dr in data_rows:
+            from labelbox import DataRow
+            if isinstance(dr, DataRow):
+                dr_ids.append(dr.uid)
+            elif isinstance(dr, str):
+                dr_ids.append(dr)
+            else:
+                raise ValueError("Must pass")
+
+        if len(dr_ids) > 25_000:
+            raise ValueError(
+                f"Batch exceeds max size, break into smaller batches"
+            )
+        if not len(dr_ids):
+            raise ValueError("You need at least one data row in a batch")
+
+        method = 'createBatch'
+        query_str = """mutation %sPyApi($projectId: ID!, $batchInput: CreateBatchInput!) {
+              project(where: {id: $projectId}) {
+                %s(input: $batchInput) {
+                  %s
+                }
+              }
+            }
+        """ % (method, method, query.results_query_part(Entity.Batch))
+
+        params = {
+            "projectId": self.uid,
+            "batchInput": {
+                "name": name,
+                "dataRowIds": dr_ids,
+                "priority": priority
+            }
+        }
+
+        res = self.client.execute(query_str, params, experimental=True)["project"][method]
+        res['size'] = len(dr_ids)
+        return Entity.Batch(self.client, res)
+
     def _update_queue_mode(self, mode: QueueMode) -> QueueMode:
 
         if self.queue_mode() == mode:
@@ -885,7 +939,7 @@ class LabelingParameterOverride(DbObject):
 
 LabelerPerformance = namedtuple(
     "LabelerPerformance", "user count seconds_per_label, total_time_labeling "
-    "consensus average_benchmark_agreement last_activity_time")
+                          "consensus average_benchmark_agreement last_activity_time")
 LabelerPerformance.__doc__ = (
     "Named tuple containing info about a labeler's performance.")
 

tests/integration/conftest.py

@@ -1,24 +1,21 @@
 import os
 import re
-import uuid
 import time
-from datetime import datetime
+import uuid
 from enum import Enum
-from random import randint
-from string import ascii_letters
 from types import SimpleNamespace
 
 import pytest
 import requests
 
 from labelbox import Client
 from labelbox import LabelingFrontend
+from labelbox import OntologyBuilder, Tool, Option, Classification
 from labelbox.orm import query
 from labelbox.pagination import PaginatedCollection
+from labelbox.schema.annotation_import import LabelImport
 from labelbox.schema.invite import Invite
 from labelbox.schema.user import User
-from labelbox import OntologyBuilder, Tool, Option, Classification
-from labelbox.schema.annotation_import import LabelImport
 
 IMG_URL = "https://picsum.photos/200/300"
 
@@ -256,7 +253,7 @@ def configured_project_with_label(client, rand_gen, image_url, project, dataset,
         Tool(tool=Tool.Type.BBOX, name="test-bbox-class"),
     ])
     project.setup(editor, ontology_builder.asdict())
-    #TODO: ontology may not be synchronous after setup. remove sleep when api is more consistent
+    # TODO: ontology may not be synchronous after setup. remove sleep when api is more consistent
     time.sleep(2)
 
     ontology = ontology_builder.from_project(project)

tests/integration/test_batch.py

@@ -0,0 +1,27 @@
+import pytest
+
+from labelbox import Dataset, Project
+from labelbox.schema.project import QueueMode
+
+IMAGE_URL = "https://storage.googleapis.com/diagnostics-demo-data/coco/COCO_train2014_000000000034.jpg"
+
+
+@pytest.fixture
+def big_dataset(dataset: Dataset):
+    task = dataset.create_data_rows([
+                                        {
+                                            "row_data": IMAGE_URL,
+                                            "external_id": "my-image"
+                                        },
+                                    ] * 250)
+    task.wait_till_done()
+
+    yield dataset
+    dataset.delete()
+
+
+def test_create_batch(configured_project: Project, big_dataset: Dataset):
+    configured_project.update(queue_mode=QueueMode.Batch)
+
+    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+    queue_res = configured_project.create_batch("test-batch", data_rows, 3)