Merge pull request #183 from Labelbox/ms/mea-tests

MEA tests

Author: msokoloff1

labelbox/schema/annotation_import.py

@@ -38,10 +38,8 @@ def inputs(self) -> List[Dict[str, Any]]:
         """
         Inputs for each individual annotation uploaded.
         This should match the ndjson annotations that you have uploaded.
-
         Returns:
             Uploaded ndjson.
-
         * This information will expire after 24 hours.
         """
         return self._fetch_remote_ndjson(self.input_file_url)
@@ -50,11 +48,9 @@ def inputs(self) -> List[Dict[str, Any]]:
     def errors(self) -> List[Dict[str, Any]]:
         """
         Errors for each individual annotation uploaded. This is a subset of statuses
-
         Returns:
             List of dicts containing error messages. Empty list means there were no errors
             See `AnnotationImport.statuses` for more details.
-
         * This information will expire after 24 hours.
         """
         self.wait_until_done()
@@ -64,38 +60,32 @@ def errors(self) -> List[Dict[str, Any]]:
     def statuses(self) -> List[Dict[str, Any]]:
         """
         Status for each individual annotation uploaded.
-
         Returns:
             A status for each annotation if the upload is done running.
             See below table for more details
-
         .. list-table::
            :widths: 15 150
            :header-rows: 1
-
            * - Field
              - Description
            * - uuid
              - Specifies the annotation for the status row.
            * - dataRow
              - JSON object containing the Labelbox data row ID for the annotation.
            * - status
-             - Indicates SUCCESS or FAILURE.
+             - Indicates SUCCESS or FAILURE.
            * - errors
-             - An array of error messages included when status is FAILURE. Each error has a name, message and optional (key might not exist) additional_info.
-
+             - An array of error messages included when status is FAILURE. Each error has a name, message and optional (key might not exist) additional_info.
         * This information will expire after 24 hours.
         """
         self.wait_until_done()
         return self._fetch_remote_ndjson(self.status_file_url)
 
     def wait_until_done(self, sleep_time_seconds: int = 10) -> None:
         """Blocks import job until certain conditions are met.
-
         Blocks until the AnnotationImport.state changes either to
         `AnnotationImportState.FINISHED` or `AnnotationImportState.FAILED`,
         periodically refreshing object's state.
-
         Args:
             sleep_time_seconds (str): a time to block between subsequent API calls
         """
@@ -117,7 +107,6 @@ def __backoff_refresh(self) -> None:
     def _fetch_remote_ndjson(self, url: str) -> List[Dict[str, Any]]:
         """
         Fetches the remote ndjson file and caches the results.
-
         Args:
             url (str): Can be any url pointing to an ndjson file.
         Returns:
@@ -132,6 +121,7 @@ def _fetch_remote_ndjson(self, url: str) -> List[Dict[str, Any]]:
 
     @classmethod
     def _build_import_predictions_query(cls, file_args: str, vars: str):
+        cls.validate_cls()
         query_str = """mutation createAnnotationImportPyApi($parent_id : ID!, $name: String!, $predictionType : PredictionType!, %s) {
         createAnnotationImport(data: {
             %s : $parent_id
@@ -148,9 +138,17 @@ def _build_import_predictions_query(cls, file_args: str, vars: str):
         return query_str
 
     @classmethod
-    def _from_name(cls, client, parent_id, name: str, raw=False):
-        query_str = """query
-        getImportPyApi($parent_id : ID!, $name: String!) {
+    def validate_cls(cls):
+        supported_base_classes = {MALPredictionImport, MEAPredictionImport}
+        if cls not in {MALPredictionImport, MEAPredictionImport}:
+            raise TypeError(
+                f"Can't directly use the base AnnotationImport class. Must use one of {supported_base_classes}"
+            )
+
+    @classmethod
+    def from_name(cls, client, parent_id, name: str, raw=False):
+        cls.validate_cls()
+        query_str = """query getImportPyApi($parent_id : ID!, $name: String!) {
             annotationImport(
                 where: {%s: $parent_id, name: $name}){
                     __typename
@@ -194,10 +192,10 @@ def refresh(self) -> None:
         """Synchronizes values of all fields with the database.
         """
         cls = type(self)
-        res = cls._from_name(self.client,
-                             self.get_parent_id(),
-                             self.name,
-                             raw=True)
+        res = cls.from_name(self.client,
+                            self.get_parent_id(),
+                            self.name,
+                            raw=True)
         self._set_field_values(res)
 
     @classmethod

labelbox/schema/model_run.py

@@ -16,6 +16,10 @@ class ModelRun(DbObject):
     created_by_id = Field.String("created_by_id", "createdBy")
     model_id = Field.String("model_id")
 
+    def __init__(self, client, field_values):
+        field_values['createdBy'] = uuid_to_cuid(field_values['createdBy'])
+        super().__init__(client, field_values)
+
     def upsert_labels(self, label_ids):
 
         if len(label_ids) < 1:
@@ -32,38 +36,34 @@ def upsert_labels(self, label_ids):
         return True
 
     def add_predictions(
-            self,
-            name: str,
-            annotations: Union[str, Path, Iterable[Dict]],
-            validate: bool = True) -> 'MEAPredictionImport':  # type: ignore
-        """ Uploads annotations to a new Editor project.
+        self,
+        name: str,
+        predictions: Union[str, Path, Iterable[Dict]],
+    ) -> 'MEAPredictionImport':  # type: ignore
+        """ Uploads predictions to a new Editor project.
         Args:
             name (str): name of the AnnotationImport job
-            annotations (str or Path or Iterable):
+            predictions (str or Path or Iterable):
                 url that is publicly accessible by Labelbox containing an
                 ndjson file
                 OR local path to an ndjson file
                 OR iterable of annotation rows
-            validate (bool):
-                Whether or not to validate the payload before uploading.
         Returns:
             AnnotationImport
         """
-        kwargs = dict(client=self.client,
-                      model_run_id=self.uid,
-                      name=name,
-                      predictions=annotations)
-        if isinstance(annotations, str) or isinstance(annotations, Path):
-            return MEAPredictionImport.create_from_file(**kwargs)
-        elif isinstance(annotations, Iterable):
-            return MEAPredictionImport.create_from_objects(**kwargs)
+        kwargs = dict(client=self.client, model_run_id=self.uid, name=name)
+        if isinstance(predictions, str) or isinstance(predictions, Path):
+            return MEAPredictionImport.create_from_file(path=predictions,
+                                                        **kwargs)
+        elif isinstance(predictions, Iterable):
+            return MEAPredictionImport.create_from_objects(
+                predictions=predictions, **kwargs)
         else:
             raise ValueError(
-                f'Invalid annotations given of type: {type(annotations)}')
+                f'Invalid predictions given of type: {type(predictions)}')
 
     def annotation_groups(self):
-        query_str = """
-            query modelRunPyApi($modelRunId: ID!, $from : String, $first: Int){
+        query_str = """query modelRunPyApi($modelRunId: ID!, $from : String, $first: Int){
                 annotationGroups(where: {modelRunId: {id: $modelRunId}}, after: $from, first: $first)
                 {nodes{%s},pageInfo{endCursor}}
             }
@@ -80,9 +80,8 @@ class AnnotationGroup(DbObject):
     model_run_id = Field.String("model_run_id")
     data_row = Relationship.ToOne("DataRow", False, cache=True)
 
-    def __init__(self, client, model_id, field_values):
-        field_values['labelId'] = uuid_to_cuid(field_values['labelId'])
-        super().__init__(client, field_values)
+    def __init__(self, client, model_id, *args, **kwargs):
+        super().__init__(client, *args, **kwargs)
         self.model_id = model_id
 
     @property

tests/integration/bulk_import/conftest.py

@@ -1,7 +1,9 @@
 import uuid
 import pytest
+import time
 
 from labelbox.schema.labeling_frontend import LabelingFrontend
+from labelbox.schema.annotation_import import MALPredictionImport
 
 IMG_URL = "https://picsum.photos/200/300"
 
@@ -100,7 +102,9 @@ def ontology():
 
 
 @pytest.fixture
-def configured_project(client, project, ontology, dataset):
+def configured_project(client, ontology, rand_gen):
+    project = client.create_project(name=rand_gen(str))
+    dataset = client.create_dataset(name=rand_gen(str), projects=project)
     editor = list(
         client.get_labeling_frontends(
             where=LabelingFrontend.name == "editor"))[0]
@@ -109,13 +113,14 @@ def configured_project(client, project, ontology, dataset):
         dataset.create_data_row(row_data=IMG_URL)
     project.datasets.connect(dataset)
     yield project
+    project.delete()
+    dataset.delete()
 
 
 @pytest.fixture
 def prediction_id_mapping(configured_project):
     #Maps tool types to feature schema ids
     ontology = configured_project.ontology().normalized
-    inferences = []
     datarows = [d for d in list(configured_project.datasets())[0].data_rows()]
     result = {}
 
@@ -266,11 +271,54 @@ def video_checklist_inference(prediction_id_mapping):
 
 
 @pytest.fixture
-def predictions(polygon_inference, rectangle_inference, line_inference,
-                entity_inference, segmentation_inference, checklist_inference,
-                text_inference):
+def model_run_predictions(polygon_inference, rectangle_inference,
+                          line_inference):
+    # Not supporting mask since there isn't a signed url representing a seg mask to upload
+    return [polygon_inference, rectangle_inference, line_inference]
+
+
+@pytest.fixture
+def object_predictions(polygon_inference, rectangle_inference, line_inference,
+                       entity_inference, segmentation_inference):
     return [
         polygon_inference, rectangle_inference, line_inference,
-        entity_inference, segmentation_inference, checklist_inference,
-        text_inference
+        entity_inference, segmentation_inference
     ]
+
+
+@pytest.fixture
+def classification_predictions(checklist_inference, text_inference):
+    return [checklist_inference, text_inference]
+
+
+@pytest.fixture
+def predictions(object_predictions, classification_predictions):
+    return object_predictions + classification_predictions
+
+
+@pytest.fixture
+def model_run(client, rand_gen, configured_project, annotation_submit_fn,
+              model_run_predictions):
+    configured_project.enable_model_assisted_labeling()
+    ontology = configured_project.ontology()
+
+    upload_task = MALPredictionImport.create_from_objects(
+        client, configured_project.uid, f'mal-import-{uuid.uuid4()}',
+        model_run_predictions)
+    upload_task.wait_until_done()
+
+    for data_row_id in {x['dataRow']['id'] for x in model_run_predictions}:
+        annotation_submit_fn(configured_project.uid, data_row_id)
+
+    data = {"name": rand_gen(str), "ontology_id": ontology.uid}
+    model = client.create_model(data["name"], data["ontology_id"])
+    name = rand_gen(str)
+    model_run_s = model.create_model_run(name)
+
+    time.sleep(3)
+    model_run_s.upsert_labels(
+        [label.uid for label in configured_project.labels()])
+    time.sleep(3)
+
+    yield model_run_s
+    # TODO: Delete resources when that is possible ..

tests/integration/bulk_import/test_mea_annotation_import.py

@@ -0,0 +1,98 @@
+import uuid
+import ndjson
+import pytest
+import requests
+
+from labelbox.schema.annotation_import import AnnotationImportState, MEAPredictionImport
+"""
+- 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 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
+
+
+def test_create_from_url(model_run):
+    name = str(uuid.uuid4())
+    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+    annotation_import = model_run.add_predictions(name=name, predictions=url)
+    assert annotation_import.model_run_id == model_run.uid
+    check_running_state(annotation_import, name, url)
+
+
+def test_create_from_objects(model_run, object_predictions):
+    name = str(uuid.uuid4())
+
+    annotation_import = model_run.add_predictions(
+        name=name, predictions=object_predictions)
+
+    assert annotation_import.model_run_id == model_run.uid
+    check_running_state(annotation_import, name)
+    assert_file_content(annotation_import.input_file_url, object_predictions)
+
+
+def test_create_from_local_file(tmp_path, model_run, object_predictions):
+    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)
+
+    annotation_import = model_run.add_predictions(name=name,
+                                                  predictions=str(file_path))
+
+    assert annotation_import.model_run_id == model_run.uid
+    check_running_state(annotation_import, name)
+    assert_file_content(annotation_import.input_file_url, object_predictions)
+
+
+def test_get(client, model_run):
+    name = str(uuid.uuid4())
+    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+    model_run.add_predictions(name=name, predictions=url)
+
+    annotation_import = MEAPredictionImport.from_name(client,
+                                                      parent_id=model_run.uid,
+                                                      name=name)
+
+    assert annotation_import.model_run_id == model_run.uid
+    check_running_state(annotation_import, name, url)
+
+
+@pytest.mark.slow
+def test_wait_till_done(model_run_predictions, model_run):
+    name = str(uuid.uuid4())
+    annotation_import = model_run.add_predictions(
+        name=name, predictions=model_run_predictions)
+
+    assert len(annotation_import.inputs) == len(model_run_predictions)
+    annotation_import.wait_until_done()
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    # Check that the status files are being returned as expected
+    assert len(annotation_import.errors) == 0
+    assert len(annotation_import.inputs) == len(model_run_predictions)
+    input_uuids = [
+        input_annot['uuid'] for input_annot in annotation_import.inputs
+    ]
+    inference_uuids = [pred['uuid'] for pred in model_run_predictions]
+    assert set(input_uuids) == set(inference_uuids)
+    assert len(annotation_import.statuses) == len(model_run_predictions)
+    for status in annotation_import.statuses:
+        assert status['status'] == 'SUCCESS'
+    status_uuids = [
+        input_annot['uuid'] for input_annot in annotation_import.statuses
+    ]
+    assert set(input_uuids) == set(status_uuids)
+
+
+def assert_file_content(url: str, predictions):
+    response = requests.get(url)
+    assert response.text == ndjson.dumps(predictions)