Merge pull request #291 from Labelbox/grant/sdk-label-imports

SDK for label imports [DIAG-604]

Author: msokoloff1

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

@@ -266,7 +266,7 @@ def from_name(cls,
 
     @classmethod
     def _get_url_mutation(cls) -> str:
-        return """mutation createMEAPredictionImportPyApi($modelRunId : ID!, $name: String!, $fileUrl: String!) {
+        return """mutation createMEAPredictionImportByUrlPyApi($modelRunId : ID!, $name: String!, $fileUrl: String!) {
             createModelErrorAnalysisPredictionImport(data: {
                 modelRunId: $modelRunId
                 name: $name
@@ -276,7 +276,7 @@ def _get_url_mutation(cls) -> str:
 
     @classmethod
     def _get_file_mutation(cls) -> str:
-        return """mutation createMEAPredictionImportPyApi($modelRunId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {
+        return """mutation createMEAPredictionImportByFilePyApi($modelRunId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {
             createModelErrorAnalysisPredictionImport(data: {
                 modelRunId: $modelRunId name: $name filePayload: { file: $file, contentLength: $contentLength}
         }) {%s}
@@ -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
@@ -423,7 +423,7 @@ def from_name(cls,
 
     @classmethod
     def _get_url_mutation(cls) -> str:
-        return """mutation createMALPredictionImportPyApi($projectId : ID!, $name: String!, $fileUrl: String!) {
+        return """mutation createMALPredictionImportByUrlPyApi($projectId : ID!, $name: String!, $fileUrl: String!) {
             createModelAssistedLabelingPredictionImport(data: {
                 projectId: $projectId
                 name: $name
@@ -433,14 +433,14 @@ def _get_url_mutation(cls) -> str:
 
     @classmethod
     def _get_file_mutation(cls) -> str:
-        return """mutation createMALPredictionImportPyApi($projectId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {
+        return """mutation createMALPredictionImportByFilePyApi($projectId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {
             createModelAssistedLabelingPredictionImport(data: {
                 projectId: $projectId name: $name filePayload: { file: $file, contentLength: $contentLength}
         }) {%s}
         }""" % 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,154 @@ 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 a 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 a 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 createLabelImportByUrlPyApi($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 createLabelImportByFilePyApi($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,8 @@ 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,10 +51,12 @@ 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):
+def test_create_from_local_file(tmp_path, predictions, configured_project,
+                                annotation_import_test_helpers):
     name = str(uuid.uuid4())
     file_name = f"{name}.ndjson"
     file_path = tmp_path / file_name
@@ -68,7 +71,8 @@ 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 +148,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