sdk for label imports

Author: GrantEating

labelbox/__init__.py

@@ -5,7 +5,7 @@
 from labelbox.client import Client
 from labelbox.schema.model import Model
 from labelbox.schema.bulk_import_request import BulkImportRequest
-from labelbox.schema.annotation_import import MALPredictionImport, MEAPredictionImport
+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

labelbox/schema/annotation_import.py

@@ -330,7 +330,7 @@ def create_from_file(cls, client: "labelbox.Client", project_id: str,
         """
         if os.path.exists(path):
             with open(path, 'rb') as f:
-                return cls._create_mea_import_from_bytes(
+                return cls._create_mal_import_from_bytes(
                     client, project_id, name, f,
                     os.stat(path).st_size)
         else:
@@ -355,7 +355,7 @@ def create_from_objects(
         if not data_str:
             raise ValueError('annotations cannot be empty')
         data = data_str.encode('utf-8')
-        return cls._create_mea_import_from_bytes(client, project_id, name, data,
+        return cls._create_mal_import_from_bytes(client, project_id, name, data,
                                                  len(data))
 
     @classmethod
@@ -440,7 +440,7 @@ def _get_file_mutation(cls) -> str:
         }""" % query.results_query_part(cls)
 
     @classmethod
-    def _create_mea_import_from_bytes(
+    def _create_mal_import_from_bytes(
             cls, client: "labelbox.Client", project_id: str, name: str,
             bytes_data: BinaryIO, content_len: int) -> "MALPredictionImport":
         file_name = f"{project_id}__{name}.ndjson"
@@ -454,3 +454,155 @@ def _create_mea_import_from_bytes(
         res = cls._create_from_bytes(client, variables, query_str, file_name,
                                      bytes_data)
         return cls(client, res["createModelAssistedLabelingPredictionImport"])
+
+
+class LabelImport(AnnotationImport):
+    project = Relationship.ToOne("Project", cache=True)
+
+    @property
+    def parent_id(self) -> str:
+        """
+        Identifier for this import. Used to refresh the status
+        """
+        return self.project().uid
+
+    @classmethod
+    def create_from_file(cls, client: "labelbox.Client", project_id: str,
+                         name: str, path: str) -> "LabelImport":
+        """
+        Create a label import job from a file of annotations
+
+        Args:
+            client: Labelbox Client for executing queries
+            project_id: Project to import labels into
+            name: Name of the import job. Can be used to reference the task later
+            path: Path to ndjson file containing annotations
+        Returns:
+            LabelImport
+        """
+        if os.path.exists(path):
+            with open(path, 'rb') as f:
+                return cls._create_label_import_from_bytes(
+                    client, project_id, name, f,
+                    os.stat(path).st_size)
+        else:
+            raise ValueError(f"File {path} is not accessible")
+
+    @classmethod
+    def create_from_objects(
+            cls, client: "labelbox.Client", project_id: str, name: str,
+            labels: List[Dict[str, Any]]) -> "LabelImport":
+        """
+        Create an label import job from an in memory dictionary
+
+        Args:
+            client: Labelbox Client for executing queries
+            project_id: Project to import labels into
+            name: Name of the import job. Can be used to reference the task later
+            labels: List of labels 
+        Returns:
+            LabelImport
+        """
+        data_str = ndjson.dumps(labels)
+        if not data_str:
+            raise ValueError('labels cannot be empty')
+        data = data_str.encode('utf-8')
+        return cls._create_label_import_from_bytes(client, project_id, name, data,
+                                                 len(data))
+
+    @classmethod
+    def create_from_url(cls, client: "labelbox.Client", project_id: str,
+                        name: str, url: str) -> "LabelImport":
+        """
+        Create an label annotation import job from a url
+        The url must point to a file containing label annotations.
+
+        Args:
+            client: Labelbox Client for executing queries
+            project_id: Project to import labels into
+            name: Name of the import job. Can be used to reference the task later
+            url: Url pointing to file to upload
+        Returns:
+            LabelImport
+        """
+        if requests.head(url):
+            query_str = cls._get_url_mutation()
+            return cls(
+                client,
+                client.execute(
+                    query_str,
+                    params={
+                        "fileUrl": url,
+                        "projectId": project_id,
+                        'name': name
+                    })["createLabelImport"])
+        else:
+            raise ValueError(f"Url {url} is not reachable")
+
+    @classmethod
+    def from_name(cls,
+                  client: "labelbox.Client",
+                  project_id: str,
+                  name: str,
+                  as_json: bool = False) -> "LabelImport":
+        """
+        Retrieves an label import job.
+
+        Args:
+            client: Labelbox Client for executing queries
+            project_id:  ID used for querying import jobs
+            name: Name of the import job.
+        Returns:
+            LabelImport
+        """
+        query_str = """query getLabelImportPyApi($projectId : ID!, $name: String!) {
+            labelImport(
+                where: {projectId: $projectId, name: $name}){
+                    %s
+                }}""" % query.results_query_part(cls)
+        params = {
+            "projectId": project_id,
+            "name": name,
+        }
+        response = client.execute(query_str, params)
+        if response is None:
+            raise labelbox.exceptions.ResourceNotFoundError(
+                LabelImport, params)
+        response = response["labelImport"]
+        if as_json:
+            return response
+        return cls(client, response)
+
+    @classmethod
+    def _get_url_mutation(cls) -> str:
+        return """mutation createLabelImportPyApi($projectId : ID!, $name: String!, $fileUrl: String!) {
+            createLabelImport(data: {
+                projectId: $projectId
+                name: $name
+                fileUrl: $fileUrl
+            }) {%s}
+        }""" % query.results_query_part(cls)
+
+    @classmethod
+    def _get_file_mutation(cls) -> str:
+        return """mutation createLabelImportPyApi($projectId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {
+            createLabelImport(data: {
+                projectId: $projectId name: $name filePayload: { file: $file, contentLength: $contentLength}
+        }) {%s}
+        }""" % query.results_query_part(cls)
+
+    @classmethod
+    def _create_label_import_from_bytes(
+            cls, client: "labelbox.Client", project_id: str, name: str,
+            bytes_data: BinaryIO, content_len: int) -> "LabelImport":
+        file_name = f"{project_id}__{name}.ndjson"
+        variables = {
+            "file": None,
+            "contentLength": content_len,
+            "projectId": project_id,
+            "name": name
+        }
+        query_str = cls._get_file_mutation()
+        res = cls._create_from_bytes(client, variables, query_str, file_name,
+                                     bytes_data)
+        return cls(client, res["createLabelImport"])

tests/integration/mal_and_mea/conftest.py renamed to tests/integration/annotation_import/conftest.py

@@ -2,9 +2,12 @@
 
 import pytest
 import time
+import requests
+import ndjson
 
+from typing import Type
 from labelbox.schema.labeling_frontend import LabelingFrontend
-from labelbox.schema.annotation_import import MALPredictionImport
+from labelbox.schema.annotation_import import MALPredictionImport, AnnotationImportState 
 
 
 @pytest.fixture
@@ -277,6 +280,7 @@ def model_run_predictions(polygon_inference, rectangle_inference,
     return [polygon_inference, rectangle_inference, line_inference]
 
 
+# also used for label imports
 @pytest.fixture
 def object_predictions(polygon_inference, rectangle_inference, line_inference,
                        entity_inference, segmentation_inference):
@@ -339,3 +343,25 @@ def model_run_annotation_groups(client, configured_project,
     time.sleep(3)
     yield model_run
     # TODO: Delete resources when that is possible ..
+
+
+class AnnotationImportTestHelpers:
+    @staticmethod
+    def assert_file_content(url: str, predictions):
+        response = requests.get(url)
+        assert response.text == ndjson.dumps(predictions)
+
+    @staticmethod
+    def check_running_state(req, name, url=None): 
+        assert req.name == name
+        if url is not None:
+            assert req.input_file_url == url
+        assert req.error_file_url is None
+        assert req.status_file_url is None
+        assert req.state == AnnotationImportState.RUNNING
+
+
+@pytest.fixture
+def annotation_import_test_helpers() -> Type[AnnotationImportTestHelpers]:
+    return AnnotationImportTestHelpers()
+

tests/integration/mal_and_mea/test_bulk_import_request.py renamed to tests/integration/annotation_import/test_bulk_import_request.py

@@ -39,7 +39,7 @@ def test_validate_file(client, configured_project):
         #Schema ids shouldn't match
 
 
-def test_create_from_objects(configured_project, predictions):
+def test_create_from_objects(configured_project, predictions, annotation_import_test_helpers):
     name = str(uuid.uuid4())
 
     bulk_import_request = configured_project.upload_annotations(
@@ -50,7 +50,7 @@ def test_create_from_objects(configured_project, predictions):
     assert bulk_import_request.error_file_url is None
     assert bulk_import_request.status_file_url is None
     assert bulk_import_request.state == BulkImportRequestState.RUNNING
-    assert_file_content(bulk_import_request.input_file_url, predictions)
+    annotation_import_test_helpers.assert_file_content(bulk_import_request.input_file_url, predictions)
 
 
 def test_create_from_local_file(tmp_path, predictions, configured_project):
@@ -68,7 +68,7 @@ def test_create_from_local_file(tmp_path, predictions, configured_project):
     assert bulk_import_request.error_file_url is None
     assert bulk_import_request.status_file_url is None
     assert bulk_import_request.state == BulkImportRequestState.RUNNING
-    assert_file_content(bulk_import_request.input_file_url, predictions)
+    annotation_import_test_helpers.assert_file_content(bulk_import_request.input_file_url, predictions)
 
 
 def test_get(client, configured_project):
@@ -144,11 +144,6 @@ def test_wait_till_done(rectangle_inference, configured_project):
         'uuid']
 
 
-def assert_file_content(url: str, predictions):
-    response = requests.get(url)
-    assert response.text == ndjson.dumps(predictions)
-
-
 def test_project_bulk_import_requests(client, configured_project, predictions):
     result = configured_project.bulk_import_requests()
     assert len(list(result)) == 0

tests/integration/annotation_import/test_label_import.py

@@ -0,0 +1,83 @@
+import uuid
+import ndjson
+import pytest
+import requests
+
+from labelbox.schema.annotation_import import AnnotationImportState, LabelImport
+"""
+- Here we only want to check that the uploads are calling the validation
+- Then with unit tests we can check the types of errors raised
+
+"""
+
+
+def test_create_from_url(client, project, annotation_import_test_helpers):
+    name = str(uuid.uuid4())
+    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+    label_import = LabelImport.create_from_url(client=client, project_id=project.uid, name=name, url=url)
+    assert label_import.parent_id == project.uid 
+    annotation_import_test_helpers.check_running_state(label_import, name, url)
+
+
+def test_create_from_objects(client, project, object_predictions, annotation_import_test_helpers):
+    name = str(uuid.uuid4())
+
+    label_import = LabelImport.create_from_objects(client=client, project_id=project.uid, name=name, labels=object_predictions)
+    assert label_import.parent_id == project.uid 
+    annotation_import_test_helpers.check_running_state(label_import, name)
+    annotation_import_test_helpers.assert_file_content(label_import.input_file_url, object_predictions)
+
+
+#   TODO: add me when we add this ability
+# def test_create_from_local_file(client, tmp_path, project,
+#                                 object_predictions, annotation_import_test_helpers):
+#     name = str(uuid.uuid4())
+#     file_name = f"{name}.ndjson"
+#     file_path = tmp_path / file_name
+#     with file_path.open("w") as f:
+#         ndjson.dump(object_predictions, f)
+
+#     label_import = LabelImport.create_from_url(client=client, project_id=project.uid, name=name, url=str(file_path))
+
+#     assert label_import.parent_id == project.uid
+#     annotation_import_test_helpers.check_running_state(label_import, name)
+#     annotation_import_test_helpers.assert_file_content(label_import.input_file_url, object_predictions)
+
+
+def test_get(client, project, annotation_import_test_helpers):
+    name = str(uuid.uuid4())
+    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+
+    label_import = LabelImport.create_from_url(client=client, project_id=project.uid, name=name, url=url)
+
+    assert label_import.parent_id == project.uid
+    annotation_import_test_helpers.check_running_state(label_import, name, url)
+
+
+@pytest.mark.slow
+def test_wait_till_done(client, project, model_run_predictions):
+    name = str(uuid.uuid4())
+    label_import = LabelImport.create_from_objects(client=client, project_id=project.uid, name=name, labels=model_run_predictions)
+
+    assert len(label_import.inputs) == len(model_run_predictions)
+    label_import.wait_until_done()
+    # TODO(grant): some of this is commented out
+    # TODO(grant): since the pipeline is not complete, you will get a failed status
+
+    # assert label_import.state == AnnotationImportState.FINISHED
+    # # Check that the status files are being returned as expected
+    # assert len(label_import.errors) == 0
+    assert len(label_import.inputs) == len(model_run_predictions)
+    input_uuids = [
+        input_annot['uuid'] for input_annot in label_import.inputs
+    ]
+    inference_uuids = [pred['uuid'] for pred in model_run_predictions]
+    assert set(input_uuids) == set(inference_uuids)
+    assert len(label_import.statuses) == len(model_run_predictions)
+    # for status in label_import.statuses:
+    #     assert status['status'] == 'SUCCESS'
+    status_uuids = [
+        input_annot['uuid'] for input_annot in label_import.statuses
+    ]
+    assert set(input_uuids) == set(status_uuids)
+