added in test cases and tested against subclasses as well as accounted for divide by 0

Author: jtsodapop

labelbox/data/metrics/confusion_matrix/calculation.py

@@ -69,7 +69,8 @@ def feature_confusion_matrix(
         return vector_confusion_matrix(ground_truths, predictions,
                                        include_subclasses, iou)
     elif isinstance(predictions[0].value, TextEntity):
-        pass  #TODO
+        return ner_confusion_matrix(ground_truths, predictions,
+                                    include_subclasses, iou)
     elif isinstance(predictions[0], ClassificationAnnotation):
         return classification_confusion_matrix(ground_truths, predictions)
     else:
@@ -293,12 +294,10 @@ def mask_confusion_matrix(ground_truths: List[ObjectAnnotation],
 
 
 def ner_confusion_matrix(ground_truths: List[ObjectAnnotation],
-                         predictions: list[ObjectAnnotation],
+                         predictions: List[ObjectAnnotation],
                          include_subclasses: bool,
                          iou: float) -> Optional[ConfusionMatrixMetricValue]:
     """Computes confusion matrix metric between two lists of TextEntity objects
-    
-    TODO: work on include_subclasses logic
 
     Args:
         ground_truths: List of ground truth mask annotations

labelbox/data/metrics/iou/calculation.py

@@ -281,8 +281,6 @@ def _get_ner_pairs(
     for ground_truth, prediction in product(ground_truths, predictions):
         score = _ner_iou(ground_truth.value, prediction.value)
         pairs.append((ground_truth, prediction, score))
-        # print(ground_truth.value.start, ground_truth.value.end,
-        #       prediction.value.start, prediction.value.end)
     return pairs
 
 
@@ -292,4 +290,7 @@ def _ner_iou(ner1: TextEntity, ner2: TextEntity):
         ner1.end, ner2.end)
     union_start, union_end = min(ner1.start,
                                  ner2.start), max(ner1.end, ner2.end)
+    #edge case of only one character in text
+    if union_start == union_end:
+        return 1
     return (intersection_end - intersection_start) / (union_end - union_start)

tests/data/metrics/confusion_matrix/conftest.py

@@ -6,6 +6,8 @@
 from labelbox.data.annotation_types import Polygon, Point, Rectangle, Mask, MaskData, Line, Radio, Text, Checklist, ClassificationAnswer
 import numpy as np
 
+from labelbox.data.annotation_types.ner import TextEntity
+
 
 class NameSpace(SimpleNamespace):
 
@@ -84,6 +86,13 @@ def get_checklist(name, answer_names):
                                     ]))
 
 
+def get_ner(name, start, end, subclasses=None):
+    return ObjectAnnotation(
+        name=name,
+        value=TextEntity(start=start, end=end),
+        classifications=[] if subclasses is None else subclasses)
+
+
 def get_object_pairs(tool_fn, **kwargs):
     return [
         NameSpace(predictions=[tool_fn("cat", **kwargs)],
@@ -326,6 +335,11 @@ def point_pairs():
     return get_object_pairs(get_point, x=0, y=0)
 
 
+@pytest.fixture
+def ner_pairs():
+    return get_object_pairs(get_ner, start=0, end=10)
+
+
 @pytest.fixture()
 def pair_iou_thresholds():
     return [

tests/data/metrics/confusion_matrix/test_confusion_matrix_data_row.py

@@ -9,7 +9,8 @@
     fixture_ref('rectangle_pairs'),
     fixture_ref('mask_pairs'),
     fixture_ref('line_pairs'),
-    fixture_ref('point_pairs')
+    fixture_ref('point_pairs'),
+    fixture_ref('ner_pairs')
 ])
 def test_overlapping_objects(tool_examples):
     for example in tool_examples:

tests/data/metrics/confusion_matrix/test_confusion_matrix_feature.py

@@ -9,7 +9,8 @@
     fixture_ref('rectangle_pairs'),
     fixture_ref('mask_pairs'),
     fixture_ref('line_pairs'),
-    fixture_ref('point_pairs')
+    fixture_ref('point_pairs'),
+    fixture_ref('ner_pairs')
 ])
 def test_overlapping_objects(tool_examples):
     for example in tool_examples: