add tests for MEA

Author: Matt Sokoloff

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/conftest.py

@@ -4,6 +4,7 @@
 from random import randint
 from string import ascii_letters
 from types import SimpleNamespace
+import uuid
 import os
 import re
 
@@ -16,6 +17,9 @@
 from labelbox import LabelingFrontend
 from labelbox import Client
 
+from labelbox.schema.ontology import OntologyBuilder, Tool
+from labelbox import LabelingFrontend, MALPredictionImport
+
 IMG_URL = "https://picsum.photos/200/300"
 
 
@@ -110,8 +114,8 @@ def __init__(self, environ: str) -> None:
 
         self.queries = []
 
-    def execute(self, query, params=None, check_naming=True, **kwargs):
-        if check_naming:
+    def execute(self, query=None, params=None, check_naming=True, **kwargs):
+        if check_naming and query is not None:
             assert re.match(r"(?:query|mutation) \w+PyApi", query) is not None
         self.queries.append((query, params))
         return super().execute(query, params, **kwargs)
@@ -241,3 +245,68 @@ def configured_project(project, client, rand_gen):
     yield project
     dataset.delete()
     project.delete()
+
+
+@pytest.fixture
+def annotation_submit_fn(client):
+
+    def submit(project_id, data_row_id):
+        feature_result = client.execute(
+            """query featuresPyApi ($project_id : ID!, $datarow_id: ID!
+            ) {project(where: { id: $project_id }) {
+                    featuresForDataRow(where: {dataRow: { id: $datarow_id }}) {id}}}
+            """, {
+                "project_id": project_id,
+                "datarow_id": data_row_id
+            })
+        features = feature_result['project']['featuresForDataRow']
+        feature_ids = [feature['id'] for feature in features]
+        client.execute(
+            """mutation createLabelPyApi ($project_id : ID!,$datarow_id: ID!,$feature_ids: [ID!]!,$time_seconds : Float!) {
+                createLabelFromFeatures(data: {dataRow: { id: $datarow_id },project: { id: $project_id },
+                    featureIds: $feature_ids,secondsSpent: $time_seconds}) {id}}""",
+            {
+                "project_id": project_id,
+                "datarow_id": data_row_id,
+                "feature_ids": feature_ids,
+                "time_seconds": 10
+            })
+
+    return submit
+
+
+@pytest.fixture
+def configured_project_with_label(client, rand_gen, annotation_submit_fn):
+    project = client.create_project(name=rand_gen(str))
+    dataset = client.create_dataset(name=rand_gen(str), projects=project)
+    data_row = dataset.create_data_row(row_data=IMG_URL)
+    editor = list(
+        project.client.get_labeling_frontends(
+            where=LabelingFrontend.name == "editor"))[0]
+
+    ontology_builder = OntologyBuilder(tools=[
+        Tool(tool=Tool.Type.BBOX, name="test-bbox-class"),
+    ])
+    project.setup(editor, ontology_builder.asdict())
+    project.enable_model_assisted_labeling()
+    ontology = ontology_builder.from_project(project)
+    predictions = [{
+        "uuid": str(uuid.uuid4()),
+        "schemaId": ontology.tools[0].feature_schema_id,
+        "dataRow": {
+            "id": data_row.uid
+        },
+        "bbox": {
+            "top": 20,
+            "left": 20,
+            "height": 50,
+            "width": 50
+        }
+    }]
+    upload_task = MALPredictionImport.create_from_objects(
+        client, project.uid, f'mal-import-{uuid.uuid4()}', predictions)
+    upload_task.wait_until_done()
+    annotation_submit_fn(project.uid, data_row.uid)
+    yield project
+    dataset.delete()
+    project.delete()