Merge pull request #761 from Labelbox/sdubinin/al-3578

[AL-3578] Confidence support in SDK

Author: dubininsergey

labelbox/data/annotation_types/annotation.py

@@ -1,6 +1,8 @@
 import abc
 from typing import Any, Dict, List, Optional, Union
 
+from labelbox.data.mixins import ConfidenceNotSupportedMixin, ConfidenceMixin
+
 from .classification import Checklist, Dropdown, Radio, Text
 from .feature import FeatureSchema
 from .geometry import Geometry, Rectangle, Point
@@ -31,7 +33,7 @@ class ClassificationAnnotation(BaseAnnotation):
     value: Union[Text, Checklist, Radio, Dropdown]
 
 
-class ObjectAnnotation(BaseAnnotation):
+class ObjectAnnotation(BaseAnnotation, ConfidenceMixin):
     """Generic localized annotation (non classifications)
 
     >>> ObjectAnnotation(
@@ -53,7 +55,7 @@ class ObjectAnnotation(BaseAnnotation):
     classifications: List[ClassificationAnnotation] = []
 
 
-class VideoObjectAnnotation(ObjectAnnotation):
+class VideoObjectAnnotation(ObjectAnnotation, ConfidenceNotSupportedMixin):
     """Video object annotation
 
     >>> VideoObjectAnnotation(
@@ -76,6 +78,7 @@ class VideoObjectAnnotation(ObjectAnnotation):
         classifications (List[ClassificationAnnotation]) = []
         extra (Dict[str, Any])
     """
+
     frame: int
     keyframe: bool
     segment_index: Optional[int] = None

labelbox/data/annotation_types/classification/classification.py

@@ -1,6 +1,8 @@
 from typing import Any, Dict, List, Union, Optional
 import warnings
 
+from labelbox.data.mixins import ConfidenceMixin
+
 try:
     from typing import Literal
 except:
@@ -20,7 +22,7 @@ def dict(self, *args, **kwargs):
         return res
 
 
-class ClassificationAnswer(FeatureSchema):
+class ClassificationAnswer(FeatureSchema, ConfidenceMixin):
     """
     - Represents a classification option.
     - Because it inherits from FeatureSchema

labelbox/data/mixins.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from pydantic import BaseModel, validator
+
+from labelbox.exceptions import ConfidenceNotSupportedException
+
+
+class ConfidenceMixin(BaseModel):
+    confidence: Optional[float] = None
+
+    @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')
+        return value
+
+    def dict(self, *args, **kwargs):
+        res = super().dict(*args, **kwargs)
+        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:
+            raise ConfidenceNotSupportedException(
+                'Confidence is not supported for this annotaiton type yet')
+        return super().__new__(cls)

labelbox/data/serialization/ndjson/classification.py

@@ -1,6 +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.utils import camel_case
 from ...annotation_types.annotation import ClassificationAnnotation, VideoClassificationAnnotation
@@ -10,7 +11,7 @@
 from .base import NDAnnotation
 
 
-class NDFeature(BaseModel):
+class NDFeature(ConfidenceMixin):
     name: Optional[str] = None
     schema_id: Optional[Cuid] = None
 
@@ -41,7 +42,7 @@ class FrameLocation(BaseModel):
 
 
 class VideoSupported(BaseModel):
-    #Note that frames are only allowed as top level inferences for video
+    # Note that frames are only allowed as top level inferences for video
     frames: Optional[List[FrameLocation]] = None
 
     def dict(self, *args, **kwargs):
@@ -70,15 +71,18 @@ class NDChecklistSubclass(NDFeature):
     def to_common(self) -> Checklist:
         return Checklist(answer=[
             ClassificationAnswer(name=answer.name,
-                                 feature_schema_id=answer.schema_id)
+                                 feature_schema_id=answer.schema_id,
+                                 confidence=answer.confidence)
             for answer in self.answer
         ])
 
     @classmethod
     def from_common(cls, checklist: Checklist, name: str,
                     feature_schema_id: Cuid) -> "NDChecklistSubclass":
         return cls(answer=[
-            NDFeature(name=answer.name, schema_id=answer.feature_schema_id)
+            NDFeature(name=answer.name,
+                      schema_id=answer.feature_schema_id,
+                      confidence=answer.confidence)
             for answer in checklist.answer
         ],
                    name=name,
@@ -95,19 +99,22 @@ class NDRadioSubclass(NDFeature):
     answer: NDFeature
 
     def to_common(self) -> Radio:
-        return Radio(answer=ClassificationAnswer(
-            name=self.answer.name, feature_schema_id=self.answer.schema_id))
+        return Radio(
+            answer=ClassificationAnswer(name=self.answer.name,
+                                        feature_schema_id=self.answer.schema_id,
+                                        confidence=self.answer.confidence))
 
     @classmethod
     def from_common(cls, radio: Radio, name: str,
                     feature_schema_id: Cuid) -> "NDRadioSubclass":
         return cls(answer=NDFeature(name=radio.answer.name,
-                                    schema_id=radio.answer.feature_schema_id),
+                                    schema_id=radio.answer.feature_schema_id,
+                                    confidence=radio.answer.confidence),
                    name=name,
                    schema_id=feature_schema_id)
 
 
-### ====== End of subclasses
+# ====== End of subclasses
 
 
 class NDText(NDAnnotation, NDTextSubclass):
@@ -133,7 +140,9 @@ def from_common(
             extra: Dict[str, Any], data: Union[VideoData, TextData,
                                                ImageData]) -> "NDChecklist":
         return cls(answer=[
-            NDFeature(name=answer.name, schema_id=answer.feature_schema_id)
+            NDFeature(name=answer.name,
+                      schema_id=answer.feature_schema_id,
+                      confidence=answer.confidence)
             for answer in checklist.answer
         ],
                    data_row={'id': data.uid},
@@ -150,7 +159,8 @@ def from_common(cls, radio: Radio, name: str, feature_schema_id: Cuid,
                     extra: Dict[str, Any], data: Union[VideoData, TextData,
                                                        ImageData]) -> "NDRadio":
         return cls(answer=NDFeature(name=radio.answer.name,
-                                    schema_id=radio.answer.feature_schema_id),
+                                    schema_id=radio.answer.feature_schema_id,
+                                    confidence=radio.answer.confidence),
                    data_row={'id': data.uid},
                    name=name,
                    schema_id=feature_schema_id,
@@ -241,6 +251,11 @@ def lookup_classification(
         }.get(type(annotation.value))
 
 
-NDSubclassificationType = Union[NDRadioSubclass, NDChecklistSubclass,
+# Make sure to keep NDChecklistSubclass prior to NDRadioSubclass in the list,
+# otherwise list of answers gets parsed by NDRadio whereas NDChecklist must be used
+NDSubclassificationType = Union[NDChecklistSubclass, NDRadioSubclass,
                                 NDTextSubclass]
-NDClassificationType = Union[NDRadio, NDChecklist, NDText]
+
+# Make sure to keep NDChecklist prior to NDRadio in the list,
+# otherwise list of answers gets parsed by NDRadio whereas NDChecklist must be used
+NDClassificationType = Union[NDChecklist, NDRadio, NDText]