Add PredictionModel and Prediction

Author: florijanstamenkovic

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## Not yet released
+
+### Added
+* `Prediction` and `PredictionModel` data types.
+
 ## Version 2.3 (2019-11-12)
 
 ### Changed

labelbox/__init__.py

@@ -12,3 +12,4 @@
 from labelbox.schema.labeling_frontend import LabelingFrontend
 from labelbox.schema.asset_metadata import AssetMetadata
 from labelbox.schema.webhook import Webhook
+from labelbox.schema.prediction import Prediction, PredictionModel

labelbox/schema/__init__.py

@@ -10,3 +10,4 @@
 import labelbox.schema.task
 import labelbox.schema.user
 import labelbox.schema.webhook
+import labelbox.schema.prediction

labelbox/schema/data_row.py

@@ -19,6 +19,7 @@ class DataRow(DbObject, Updateable, BulkDeletable):
     organization = Relationship.ToOne("Organization", False)
     labels = Relationship.ToMany("Label", True)
     metadata = Relationship.ToMany("AssetMetadata", False, "metadata")
+    predictions = Relationship.ToMany("Prediction", False)
 
     @staticmethod
     def bulk_delete(data_rows):

labelbox/schema/prediction.py

@@ -0,0 +1,31 @@
+from labelbox.orm.db_object import DbObject
+from labelbox.orm.model import Field, Relationship
+
+
+class PredictionModel(DbObject):
+    """ A prediction model represents a specific version of a model. """
+    updated_at = Field.DateTime("updated_at")
+    created_at = Field.DateTime("created_at")
+    created_by = Relationship.ToOne("User", False, "created_by")
+    organization = Relationship.ToOne("Organization", False)
+
+    name = Field.String("name")
+    slug = Field.String("slug")
+    version = Field.Int("version")
+
+    created_predictions = Relationship.ToMany("Prediction", False,
+                                              "created_predictions")
+
+
+class Prediction(DbObject):
+    """ A prediction created by a PredictionModel. """
+    updated_at = Field.DateTime("updated_at")
+    created_at = Field.DateTime("created_at")
+    organization = Relationship.ToOne("Organization", False)
+
+    label = Field.String("label")
+    agreement = Field.Float("agreement")
+
+    prediction_model = Relationship.ToOne("PredictionModel", False)
+    data_row = Relationship.ToOne("DataRow", False)
+    project = Relationship.ToOne("Project", False)

labelbox/schema/project.py

@@ -40,6 +40,9 @@ class Project(DbObject, Updateable, Deletable):
         "LabelingParameterOverride", False, "labeling_parameter_overrides")
     webhooks = Relationship.ToMany("Webhook", False)
     benchmarks = Relationship.ToMany("Benchmark", False)
+    active_prediction_model = Relationship.ToOne("PredictionModel", False,
+                                                 "active_prediction_model")
+    predictions = Relationship.ToMany("Prediction", False)
 
     def create_label(self, **kwargs):
         """ Creates a label on this Project.
@@ -283,6 +286,59 @@ def extend_reservations(self, queue_type):
         res = self.client.execute(query_str, {id_param: self.uid})
         return res["extendReservations"]
 
+    def create_prediction_model(self, name, version):
+        """ Creates a PredictionModel connected to this Project.
+        Args:
+            name (str): The new PredictionModel's name.
+            version (int): The new PredictionModel's version.
+        Return:
+            A newly created PredictionModel.
+        """
+        PM = Entity.PredictionModel
+        model =  self.client._create(
+            PM, {PM.name.name: name, PM.version.name: version})
+        self.active_prediction_model.connect(model)
+        return model
+
+    def create_prediction(self, label, data_row, prediction_model=None):
+        """ Creates a Prediction within this Project.
+        Args:
+            label (str): The `label` field of the new Prediction.
+            data_row (DataRow): The DataRow for which the Prediction is created.
+            prediction_model (PredictionModel or None): The PredictionModel
+                within which the new Prediction is created. If None then this
+                Project's active_prediction_model is used.
+        Return:
+            A newly created Prediction.
+        Raises:
+            labelbox.excepions.InvalidQueryError: if given `prediction_model`
+                is None and this Project's active_prediction_model is also
+                None.
+        """
+        if prediction_model is None:
+            prediction_model = self.active_prediction_model()
+            if prediction_model is None:
+                raise InvalidQueryError(
+                    "Project '%s' has no active prediction model" % self.name)
+
+        label_param = "label"
+        model_param = "prediction_model_id"
+        project_param = "project_id"
+        data_row_param = "data_row_id"
+
+        Prediction = Entity.Prediction
+        query_str = """mutation CreatePredictionPyApi(
+            $%s: String!, $%s: ID!, $%s: ID!, $%s: ID!) {createPrediction(
+            data: {label: $%s, predictionModelId: $%s, projectId: $%s,
+                   dataRowId: $%s})
+            {%s}}""" % (label_param, model_param, project_param, data_row_param,
+                        label_param, model_param, project_param, data_row_param,
+                        query.results_query_part(Prediction))
+        params = {label_param: label, model_param: prediction_model.uid,
+                  data_row_param: data_row.uid, project_param: self.uid}
+        res = self.client.execute(query_str, params)
+        return Prediction(self.client, res["createPrediction"])
+
 
 class LabelingParameterOverride(DbObject):
     priority = Field.Int("priority")

tests/integration/test_predictions.py

@@ -0,0 +1,36 @@
+def test_prediction_model(project, rand_gen):
+    model_name = rand_gen(str)
+    version = 42
+    model_1 = project.create_prediction_model(model_name, version)
+    assert model_1.name == model_name
+    assert model_1.version == version
+    assert project.active_prediction_model() == model_1
+
+    model_2 = project.create_prediction_model(rand_gen(str), 22)
+    assert project.active_prediction_model() == model_2
+
+
+def test_predictions(label_pack, rand_gen):
+    project, dataset, data_row, label = label_pack
+    model_1 = project.create_prediction_model(rand_gen(str), 12)
+
+    assert set(project.predictions()) == set()
+    pred_1 = project.create_prediction("l1", data_row)
+    assert pred_1.label == "l1"
+    assert set(model_1.created_predictions()) == {pred_1}
+    assert set(project.predictions()) == {pred_1}
+    assert set(data_row.predictions()) == {pred_1}
+    assert pred_1.prediction_model() == model_1
+    assert pred_1.data_row() == data_row
+    assert pred_1.project() == project
+    label_2 = project.create_label(data_row=data_row, label="test",
+                                   seconds_to_label=0.0)
+
+    model_2 = project.create_prediction_model(rand_gen(str), 12)
+    assert set(project.predictions()) == {pred_1}
+    pred_2 = project.create_prediction("l2", data_row)
+    assert pred_2.label == "l2"
+    assert set(model_1.created_predictions()) == {pred_1}
+    assert set(model_2.created_predictions()) == {pred_2}
+    assert set(project.predictions()) == {pred_1, pred_2}
+    assert set(data_row.predictions()) == {pred_1, pred_2}

tools/db_object_doc_gen.py

@@ -44,6 +44,7 @@
     labelbox.Project, labelbox.Dataset, labelbox.DataRow, labelbox.Label,
     labelbox.AssetMetadata, labelbox.LabelingFrontend, labelbox.Task,
     labelbox.Webhook, labelbox.User, labelbox.Organization, labelbox.Review,
+    labelbox.Prediction, labelbox.PredictionModel,
     LabelerPerformance]
 
 ERROR_CLASSES = [LabelboxError] + LabelboxError.__subclasses__()