[AL-7738] Support custom metrics for predictions (#1399)

Author: kozikkamil

labelbox/data/annotation_types/annotation.py

@@ -3,13 +3,13 @@
 from labelbox.data.annotation_types.base_annotation import BaseAnnotation
 from labelbox.data.annotation_types.geometry.geometry import Geometry
 
-from labelbox.data.mixins import ConfidenceMixin
+from labelbox.data.mixins import ConfidenceMixin, CustomMetricsMixin
 
 from labelbox.data.annotation_types.classification.classification import ClassificationAnnotation
 from .ner import DocumentEntity, TextEntity, ConversationEntity
 
 
-class ObjectAnnotation(BaseAnnotation, ConfidenceMixin):
+class ObjectAnnotation(BaseAnnotation, ConfidenceMixin, CustomMetricsMixin):
     """Generic localized annotation (non classifications)
 
     >>> ObjectAnnotation(
@@ -27,5 +27,6 @@ class ObjectAnnotation(BaseAnnotation, ConfidenceMixin):
         classifications (Optional[List[ClassificationAnnotation]]): Optional sub classification of the annotation
         extra (Dict[str, Any])
     """
+
     value: Union[TextEntity, ConversationEntity, DocumentEntity, Geometry]
     classifications: List[ClassificationAnnotation] = []

labelbox/data/annotation_types/classification/classification.py

@@ -2,7 +2,7 @@
 import warnings
 from labelbox.data.annotation_types.base_annotation import BaseAnnotation
 
-from labelbox.data.mixins import ConfidenceMixin
+from labelbox.data.mixins import ConfidenceMixin, CustomMetricsMixin
 
 try:
     from typing import Literal
@@ -23,7 +23,7 @@ def dict(self, *args, **kwargs):
         return res
 
 
-class ClassificationAnswer(FeatureSchema, ConfidenceMixin):
+class ClassificationAnswer(FeatureSchema, ConfidenceMixin, CustomMetricsMixin):
     """
     - Represents a classification option.
     - Because it inherits from FeatureSchema
@@ -47,7 +47,7 @@ def dict(self, *args, **kwargs) -> Dict[str, str]:
         return res
 
 
-class Radio(ConfidenceMixin, BaseModel):
+class Radio(ConfidenceMixin, CustomMetricsMixin, BaseModel):
     """ A classification with only one selected option allowed
 
     >>> Radio(answer = ClassificationAnswer(name = "dog"))
@@ -66,7 +66,7 @@ class Checklist(_TempName):
     answer: List[ClassificationAnswer]
 
 
-class Text(ConfidenceMixin, BaseModel):
+class Text(ConfidenceMixin, CustomMetricsMixin, BaseModel):
     """ Free form text
 
     >>> Text(answer = "some text answer")
@@ -93,7 +93,8 @@ def __init__(self, **data: Any):
                       "removed in a future release")
 
 
-class ClassificationAnnotation(BaseAnnotation, ConfidenceMixin):
+class ClassificationAnnotation(BaseAnnotation, ConfidenceMixin,
+                               CustomMetricsMixin):
     """Classification annotations (non localized)
 
     >>> ClassificationAnnotation(

labelbox/data/annotation_types/video.py

@@ -6,7 +6,7 @@
 
 from labelbox.data.annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation
 from labelbox.data.annotation_types.feature import FeatureSchema
-from labelbox.data.mixins import ConfidenceNotSupportedMixin
+from labelbox.data.mixins import ConfidenceNotSupportedMixin, CustomMetricsNotSupportedMixin
 from labelbox.utils import _CamelCaseMixin, is_valid_uri
 
 
@@ -24,7 +24,8 @@ class VideoClassificationAnnotation(ClassificationAnnotation):
     segment_index: Optional[int] = None
 
 
-class VideoObjectAnnotation(ObjectAnnotation, ConfidenceNotSupportedMixin):
+class VideoObjectAnnotation(ObjectAnnotation, ConfidenceNotSupportedMixin,
+                            CustomMetricsNotSupportedMixin):
     """Video object annotation
     >>> VideoObjectAnnotation(
     >>>     keyframe=True,

labelbox/data/mixins.py

@@ -1,32 +1,73 @@
-from typing import Optional
+from typing import Optional, List
 
 from pydantic import BaseModel, validator
 
-from labelbox.exceptions import ConfidenceNotSupportedException
+from labelbox.exceptions import ConfidenceNotSupportedException, CustomMetricsNotSupportedException
 
 
 class ConfidenceMixin(BaseModel):
     confidence: Optional[float] = None
 
-    @validator('confidence')
+    @validator("confidence")
     def confidence_valid_float(cls, value):
         if value is None:
             return value
         if not isinstance(value, (int, float)) or not 0 <= value <= 1:
-            raise ValueError('must be a number within [0,1] range')
+            raise ValueError("must be a number within [0,1] range")
         return value
 
     def dict(self, *args, **kwargs):
         res = super().dict(*args, **kwargs)
-        if 'confidence' in res and res['confidence'] is None:
-            res.pop('confidence')
+        if "confidence" in res and res["confidence"] is None:
+            res.pop("confidence")
         return res
 
 
 class ConfidenceNotSupportedMixin:
 
     def __new__(cls, *args, **kwargs):
-        if 'confidence' in kwargs:
+        if "confidence" in kwargs:
             raise ConfidenceNotSupportedException(
-                'Confidence is not supported for this annotaiton type yet')
+                "Confidence is not supported for this annotation type yet")
+        return super().__new__(cls)
+
+
+class CustomMetric(BaseModel):
+    name: str
+    value: float
+
+    @validator("name")
+    def confidence_valid_float(cls, value):
+        if not isinstance(value, str):
+            raise ValueError("Name must be a string")
+        return value
+
+    @validator("value")
+    def value_valid_float(cls, value):
+        if not isinstance(value, (int, float)):
+            raise ValueError("Value must be a number")
+        return value
+
+
+class CustomMetricsMixin(BaseModel):
+    custom_metrics: Optional[List[CustomMetric]] = None
+
+    def dict(self, *args, **kwargs):
+        res = super().dict(*args, **kwargs)
+
+        if "customMetrics" in res and res["customMetrics"] is None:
+            res.pop("customMetrics")
+
+        if "custom_metrics" in res and res["custom_metrics"] is None:
+            res.pop("custom_metrics")
+
+        return res
+
+
+class CustomMetricsNotSupportedMixin:
+
+    def __new__(cls, *args, **kwargs):
+        if "custom_metrics" in kwargs:
+            raise CustomMetricsNotSupportedException(
+                "Custom metrics is not supported for this annotation type yet")
         return super().__new__(cls)

labelbox/data/serialization/ndjson/classification.py

@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Union, Optional
 
 from pydantic import BaseModel, Field, root_validator
-from labelbox.data.mixins import ConfidenceMixin
+from labelbox.data.mixins import ConfidenceMixin, CustomMetric, CustomMetricsMixin
 from labelbox.data.serialization.ndjson.base import DataRow, NDAnnotation
 
 from labelbox.utils import camel_case
@@ -12,7 +12,7 @@
 from ...annotation_types.data import TextData, VideoData, ImageData
 
 
-class NDAnswer(ConfidenceMixin):
+class NDAnswer(ConfidenceMixin, CustomMetricsMixin):
     name: Optional[str] = None
     schema_id: Optional[Cuid] = None
     classifications: Optional[List['NDSubclassificationType']] = []
@@ -64,15 +64,20 @@ class NDTextSubclass(NDAnswer):
     answer: str
 
     def to_common(self) -> Text:
-        return Text(answer=self.answer, confidence=self.confidence)
+        return Text(answer=self.answer,
+                    confidence=self.confidence,
+                    custom_metrics=self.custom_metrics)
 
     @classmethod
     def from_common(cls, text: Text, name: str,
                     feature_schema_id: Cuid) -> "NDTextSubclass":
-        return cls(answer=text.answer,
-                   name=name,
-                   schema_id=feature_schema_id,
-                   confidence=text.confidence)
+        return cls(
+            answer=text.answer,
+            name=name,
+            schema_id=feature_schema_id,
+            confidence=text.confidence,
+            custom_metrics=text.custom_metrics,
+        )
 
 
 class NDChecklistSubclass(NDAnswer):
@@ -87,7 +92,8 @@ def to_common(self) -> Checklist:
                                  classifications=[
                                      NDSubclassification.to_common(annot)
                                      for annot in answer.classifications
-                                 ])
+                                 ],
+                                 custom_metrics=answer.custom_metrics)
             for answer in self.answer
         ])
 
@@ -101,7 +107,8 @@ def from_common(cls, checklist: Checklist, name: str,
                      classifications=[
                          NDSubclassification.from_common(annot)
                          for annot in answer.classifications
-                     ])
+                     ],
+                     custom_metrics=answer.custom_metrics)
             for answer in checklist.answer
         ],
                    name=name,
@@ -126,7 +133,7 @@ def to_common(self) -> Radio:
                 NDSubclassification.to_common(annot)
                 for annot in self.answer.classifications
             ],
-        ))
+            custom_metrics=self.answer.custom_metrics))
 
     @classmethod
     def from_common(cls, radio: Radio, name: str,
@@ -137,7 +144,8 @@ def from_common(cls, radio: Radio, name: str,
                                    classifications=[
                                        NDSubclassification.from_common(annot)
                                        for annot in radio.answer.classifications
-                                   ]),
+                                   ],
+                                   custom_metrics=radio.answer.custom_metrics),
                    name=name,
                    schema_id=feature_schema_id)
 
@@ -165,21 +173,25 @@ def from_common(cls,
             uuid=uuid,
             message_id=message_id,
             confidence=text.confidence,
+            custom_metrics=text.custom_metrics,
         )
 
 
 class NDChecklist(NDAnnotation, NDChecklistSubclass, VideoSupported):
 
     @classmethod
-    def from_common(cls,
-                    uuid: str,
-                    checklist: Checklist,
-                    name: str,
-                    feature_schema_id: Cuid,
-                    extra: Dict[str, Any],
-                    data: Union[VideoData, TextData, ImageData],
-                    message_id: str,
-                    confidence: Optional[float] = None) -> "NDChecklist":
+    def from_common(
+            cls,
+            uuid: str,
+            checklist: Checklist,
+            name: str,
+            feature_schema_id: Cuid,
+            extra: Dict[str, Any],
+            data: Union[VideoData, TextData, ImageData],
+            message_id: str,
+            confidence: Optional[float] = None,
+            custom_metrics: Optional[List[CustomMetric]] = None
+    ) -> "NDChecklist":
 
         return cls(answer=[
             NDAnswer(name=answer.name,
@@ -188,7 +200,8 @@ def from_common(cls,
                      classifications=[
                          NDSubclassification.from_common(annot)
                          for annot in answer.classifications
-                     ])
+                     ],
+                     custom_metrics=answer.custom_metrics)
             for answer in checklist.answer
         ],
                    data_row=DataRow(id=data.uid, global_key=data.global_key),
@@ -220,7 +233,8 @@ def from_common(
                                    classifications=[
                                        NDSubclassification.from_common(annot)
                                        for annot in radio.answer.classifications
-                                   ]),
+                                   ],
+                                   custom_metrics=radio.answer.custom_metrics),
                    data_row=DataRow(id=data.uid, global_key=data.global_key),
                    name=name,
                    schema_id=feature_schema_id,