Merge pull request #505 from Labelbox/gj/batch

Batch SDK Beta

Author: Gareth

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
 from labelbox.schema.review import Review
 from labelbox.schema.user import User
 from labelbox.schema.organization import Organization

labelbox/orm/model.py

@@ -347,6 +347,7 @@ class Entity(metaclass=EntityMeta):
         Invite: Type[labelbox.Invite]
         InviteLimit: Type[labelbox.InviteLimit]
         ProjectRole: Type[labelbox.ProjectRole]
+        Batch: Type[labelbox.Batch]
 
     @classmethod
     def _attributes_of_type(cls, attr_type):

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

@@ -0,0 +1,25 @@
+from labelbox.orm.db_object import DbObject
+from labelbox.orm.model import Field, Relationship
+
+
+class Batch(DbObject):
+    """ A Batch is a group of data rows submitted to a project for labeling
+
+    Attributes:
+        name (str)
+        created_at (datetime)
+        updated_at (datetime)
+        deleted (bool)
+
+        project (Relationship): `ToOne` relationship to Project
+        created_by (Relationship): `ToOne` relationship to User
+
+    """
+    name = Field.String("name")
+    created_at = Field.DateTime("created_at")
+    updated_at = Field.DateTime("updated_at")
+    size = Field.Int("size")
+
+    # Relationships
+    project = Relationship.ToOne("Project")
+    created_by = Relationship.ToOne("User")

labelbox/schema/project.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import time
+import warnings
 from collections import namedtuple
 from datetime import datetime, timezone
 from pathlib import Path
@@ -37,11 +38,6 @@
 logger = logging.getLogger(__name__)
 
 
-class QueueMode(enum.Enum):
-    Batch = "Batch"
-    Dataset = "Dataset"
-
-
 class Project(DbObject, Updateable, Deletable):
     """ A Project is a container that includes a labeling frontend, an ontology,
     datasets and labels.
@@ -89,9 +85,12 @@ class Project(DbObject, Updateable, Deletable):
     benchmarks = Relationship.ToMany("Benchmark", False)
     ontology = Relationship.ToOne("Ontology", True)
 
-    def update(self, **kwargs):
+    class QueueMode(enum.Enum):
+        Batch = "Batch"
+        Dataset = "Dataset"
 
-        mode: Optional[QueueMode] = kwargs.pop("queue_mode", None)
+    def update(self, **kwargs):
+        mode: Optional[Project.QueueMode] = kwargs.pop("queue_mode", None)
         if mode:
             self._update_queue_mode(mode)
 
@@ -569,14 +568,69 @@ 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 _update_queue_mode(self, mode: QueueMode) -> QueueMode:
+    def create_batch(self, name: str, data_rows: List[str], priority: int = 5):
+        """Create a new batch for a project. Batches is in Beta and subject to change
+
+        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() != Project.QueueMode.Batch:
+            raise ValueError("Project must be in batch mode")
+
+        dr_ids = []
+        for dr in data_rows:
+            if isinstance(dr, Entity.DataRow):
+                dr_ids.append(dr.uid)
+            elif isinstance(dr, str):
+                dr_ids.append(dr)
+            else:
+                raise ValueError("You can DataRow ids or DataRow objects")
+
+        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: "Project.QueueMode") -> "Project.QueueMode":
 
         if self.queue_mode() == mode:
             return mode
 
-        if mode == QueueMode.Batch:
+        if mode == Project.QueueMode.Batch:
             status = "ENABLED"
-        elif mode == QueueMode.Dataset:
+        elif mode == Project.QueueMode.Dataset:
             status = "DISABLED"
         else:
             raise ValueError(
@@ -598,7 +652,7 @@ def _update_queue_mode(self, mode: QueueMode) -> QueueMode:
 
         return mode
 
-    def queue_mode(self) -> QueueMode:
+    def queue_mode(self) -> "Project.QueueMode":
         """Provides the status of if queue mode is enabled in the project."""
 
         query_str = """query %s($projectId: ID!) {
@@ -612,9 +666,9 @@ def queue_mode(self) -> QueueMode:
             query_str, {'projectId': self.uid})["project"]["tagSetStatus"]
 
         if status == "ENABLED":
-            return QueueMode.Batch
+            return Project.QueueMode.Batch
         elif status == "DISABLED":
-            return QueueMode.Dataset
+            return Project.QueueMode.Dataset
         else:
             raise ValueError("Status not known")
 

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,28 @@
+import pytest
+
+from labelbox import Dataset, Project
+
+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=Project.QueueMode.Batch)
+
+    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+    batch = configured_project.create_batch("test-batch", data_rows, 3)
+    assert batch.name == 'test-batch'
+    assert batch.size == len(data_rows)

tests/integration/test_project.py

@@ -5,7 +5,6 @@
 
 from labelbox import Project, LabelingFrontend
 from labelbox.exceptions import InvalidQueryError
-from labelbox.schema.project import QueueMode
 
 
 def test_project(client, rand_gen):
@@ -181,6 +180,7 @@ def test_queued_data_row_export(configured_project):
 
 
 def test_queue_mode(configured_project: Project):
-    assert configured_project.queue_mode() == QueueMode.Dataset
-    configured_project.update(queue_mode=QueueMode.Batch)
-    assert configured_project.queue_mode() == QueueMode.Batch
+    assert configured_project.queue_mode(
+    ) == configured_project.QueueMode.Dataset
+    configured_project.update(queue_mode=configured_project.QueueMode.Batch)
+    assert configured_project.queue_mode() == configured_project.QueueMode.Batch