Merge pull request #265 from Labelbox/ms/custom-metric-models

custom scalar metric

Author: msokoloff1

labelbox/data/annotation_types/__init__.py

@@ -29,3 +29,5 @@
 from .collection import LabelGenerator
 
 from .metrics import ScalarMetric
+from .metrics import CustomScalarMetric
+from .metrics import MetricAggregation

labelbox/data/annotation_types/label.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from labelbox.data.annotation_types.metrics.scalar import CustomScalarMetric
 
 from typing import Any, Callable, Dict, List, Union, Optional
 
@@ -21,7 +22,8 @@ class Label(BaseModel):
     data: Union[VideoData, ImageData, TextData]
     annotations: List[Union[ClassificationAnnotation, ObjectAnnotation,
                             VideoObjectAnnotation,
-                            VideoClassificationAnnotation, ScalarMetric]] = []
+                            VideoClassificationAnnotation, CustomScalarMetric,
+                            ScalarMetric]] = []
     extra: Dict[str, Any] = {}
 
     def object_annotations(self) -> List[ObjectAnnotation]:

labelbox/data/annotation_types/metrics.py

@@ -0,0 +1,2 @@
+from .scalar import ScalarMetric, CustomScalarMetric
+from .aggregations import MetricAggregation

labelbox/data/annotation_types/metrics/__init__.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class MetricAggregation(Enum):
+    ARITHMETIC_MEAN = "ARITHMETIC_MEAN"
+    GEOMETRIC_MEAN = "GEOMETRIC_MEAN"
+    HARMONIC_MEAN = "HARMONIC_MEAN"
+    SUM = "SUM"

labelbox/data/annotation_types/metrics/aggregations.py

@@ -0,0 +1,18 @@
+from labelbox.data.annotation_types.metrics.aggregations import MetricAggregation
+from typing import Any, Dict, Optional
+from pydantic import BaseModel
+
+
+class ScalarMetric(BaseModel):
+    """ Class representing metrics """
+    value: float
+    extra: Dict[str, Any] = {}
+
+
+class CustomScalarMetric(BaseModel):
+    metric_name: str
+    value: float
+    feature_name: Optional[str] = None
+    subclass_name: Optional[str] = None
+    aggregation: MetricAggregation = MetricAggregation.ARITHMETIC_MEAN
+    extra: Dict[str, Any] = {}

labelbox/data/annotation_types/metrics/scalar.py

@@ -388,8 +388,8 @@ def create_from_local_file(cls,
 
     def delete(self) -> None:
         """ Deletes the import job and also any annotations created by this import.
-        
-        Returns: 
+
+        Returns:
             None
         """
         id_param = "bulk_request_id"

labelbox/schema/bulk_import_request.py

@@ -1,6 +1,10 @@
+from labelbox.data.annotation_types.metrics.aggregations import MetricAggregation
+from labelbox.data.annotation_types.metrics.scalar import CustomScalarMetric
 from labelbox.data.annotation_types.collection import LabelList
 from labelbox.data.annotation_types import ScalarMetric, Label, ImageData
 
+import pytest
+
 
 def test_scalar_metric():
     value = 10
@@ -27,3 +31,49 @@ def test_scalar_metric():
     }
     assert label.dict() == expected
     next(LabelList([label])).dict() == expected
+
+
+@pytest.mark.parametrize('feature_name,subclass_name,aggregation', [
+    ("cat", "orange", MetricAggregation.ARITHMETIC_MEAN),
+    ("cat", None, MetricAggregation.ARITHMETIC_MEAN),
+    (None, None, MetricAggregation.ARITHMETIC_MEAN),
+    (None, None, None),
+    ("cat", "orange", MetricAggregation.ARITHMETIC_MEAN),
+    ("cat", None, MetricAggregation.HARMONIC_MEAN),
+    (None, None, MetricAggregation.GEOMETRIC_MEAN),
+    (None, None, MetricAggregation.SUM),
+])
+def test_custom_scalar_metric(feature_name, subclass_name, aggregation):
+    value = 0.5
+    kwargs = {'aggregation': aggregation} if aggregation is not None else {}
+    metric = CustomScalarMetric(metric_name="iou",
+                                value=value,
+                                feature_name=feature_name,
+                                subclass_name=subclass_name,
+                                **kwargs)
+    assert metric.value == value
+
+    label = Label(data=ImageData(uid="ckrmd9q8g000009mg6vej7hzg"),
+                  annotations=[metric])
+    expected = {
+        'data': {
+            'external_id': None,
+            'uid': 'ckrmd9q8g000009mg6vej7hzg',
+            'im_bytes': None,
+            'file_path': None,
+            'url': None,
+            'arr': None
+        },
+        'annotations': [{
+            'value': value,
+            'metric_name': 'iou',
+            'feature_name': feature_name,
+            'subclass_name': subclass_name,
+            'aggregation': aggregation or MetricAggregation.ARITHMETIC_MEAN,
+            'extra': {}
+        }],
+        'extra': {},
+        'uid': None
+    }
+    assert label.dict() == expected
+    next(LabelList([label])).dict() == expected