Merge pull request #836 from Labelbox/sdubinin/al-000-fixed-missing-confidence-attribute

[AL-000] Fixed missing confidence attribute in NDline and NDTextEntity

Author: dubininsergey

labelbox/data/serialization/ndjson/objects.py

@@ -93,10 +93,14 @@ def to_common(self) -> Line:
         return Line(points=[Point(x=pt.x, y=pt.y) for pt in self.line])
 
     @classmethod
-    def from_common(cls, line: Line,
-                    classifications: List[ClassificationAnnotation], name: str,
-                    feature_schema_id: Cuid, extra: Dict[str, Any],
-                    data: Union[ImageData, TextData]) -> "NDLine":
+    def from_common(cls,
+                    line: Line,
+                    classifications: List[ClassificationAnnotation],
+                    name: str,
+                    feature_schema_id: Cuid,
+                    extra: Dict[str, Any],
+                    data: Union[ImageData, TextData],
+                    confidence: Optional[float] = None) -> "NDLine":
         return cls(line=[{
             'x': pt.x,
             'y': pt.y
@@ -105,7 +109,8 @@ def from_common(cls, line: Line,
                    name=name,
                    schema_id=feature_schema_id,
                    uuid=extra.get('uuid'),
-                   classifications=classifications)
+                   classifications=classifications,
+                   confidence=confidence)
 
 
 class NDFrameLine(VideoSupported):
@@ -340,7 +345,7 @@ class Location(BaseModel):
     end: int
 
 
-class NDTextEntity(NDBaseObject):
+class NDTextEntity(NDBaseObject, ConfidenceMixin):
     location: Location
 
     def to_common(self) -> TextEntity:
@@ -350,7 +355,8 @@ def to_common(self) -> TextEntity:
     def from_common(cls, text_entity: TextEntity,
                     classifications: List[ClassificationAnnotation], name: str,
                     feature_schema_id: Cuid, extra: Dict[str, Any],
-                    data: Union[ImageData, TextData]) -> "NDTextEntity":
+                    data: Union[ImageData,
+                                TextData], confidence: float) -> "NDTextEntity":
         return cls(location=Location(
             start=text_entity.start,
             end=text_entity.end,
@@ -359,7 +365,8 @@ def from_common(cls, text_entity: TextEntity,
                    name=name,
                    schema_id=feature_schema_id,
                    uuid=extra.get('uuid'),
-                   classifications=classifications)
+                   classifications=classifications,
+                   confidence=confidence)
 
 
 class NDObject:

tests/data/assets/ndjson/polyline_import.json

@@ -0,0 +1,26 @@
+[
+    {
+        "line": [
+            {
+                "x": 2534.353,
+                "y": 249.471
+            },
+            {
+                "x": 2429.492,
+                "y": 182.092
+            },
+            {
+                "x": 2294.322,
+                "y": 221.962
+            }
+        ],
+        "uuid": "5ad9c52f-058d-49c8-a749-3f20b84f8cd4",
+        "dataRow": {
+            "id": "cl6xnv9h61fv0085yhtoq06ht"
+        },
+        "name": "some-line",
+        "schemaId": "cl6xnuwt95lqq07330tbb3mfd",
+        "classifications": [],
+        "confidence": 0.58
+    }
+]

tests/data/assets/ndjson/polyline_without_confidence_import.json

@@ -0,0 +1,25 @@
+[
+    {
+        "line": [
+            {
+                "x": 2534.353,
+                "y": 249.471
+            },
+            {
+                "x": 2429.492,
+                "y": 182.092
+            },
+            {
+                "x": 2294.322,
+                "y": 221.962
+            }
+        ],
+        "uuid": "5ad9c52f-058d-49c8-a749-3f20b84f8cd4",
+        "dataRow": {
+            "id": "cl6xnv9h61fv0085yhtoq06ht"
+        },
+        "name": "some-line",
+        "schemaId": "cl6xnuwt95lqq07330tbb3mfd",
+        "classifications": []
+    }
+]

tests/data/assets/ndjson/text_entity_import.json

@@ -0,0 +1,16 @@
+[
+    {
+        "location": {
+            "start": 67,
+            "end": 128
+        },
+        "uuid": "5ad9c52f-058d-49c8-a749-3f20b84f8cd4",
+        "dataRow": {
+            "id": "cl6xnv9h61fv0085yhtoq06ht"
+        },
+        "name": "some-text-entity",
+        "schemaId": "cl6xnuwt95lqq07330tbb3mfd",
+        "classifications": [],
+        "confidence": 0.53
+    }
+]

tests/data/assets/ndjson/text_entity_without_confidence_import.json

@@ -0,0 +1,15 @@
+[
+    {
+        "location": {
+            "start": 67,
+            "end": 128
+        },
+        "uuid": "5ad9c52f-058d-49c8-a749-3f20b84f8cd4",
+        "dataRow": {
+            "id": "cl6xnv9h61fv0085yhtoq06ht"
+        },
+        "name": "some-text-entity",
+        "schemaId": "cl6xnuwt95lqq07330tbb3mfd",
+        "classifications": []
+    }
+]

tests/data/serialization/ndjson/test_polyline.py

@@ -0,0 +1,15 @@
+import json
+import pytest
+from labelbox.data.serialization.ndjson.converter import NDJsonConverter
+
+
+@pytest.mark.parametrize("filename", [
+    "tests/data/assets/ndjson/polyline_without_confidence_import.json",
+    "tests/data/assets/ndjson/polyline_import.json"
+])
+def test_polyline_import(filename: str):
+    with open(filename, 'r') as file:
+        data = json.load(file)
+    res = NDJsonConverter.deserialize(data).as_list()
+    res = list(NDJsonConverter.serialize(res))
+    assert res == data

tests/data/serialization/ndjson/test_text_entity.py

@@ -0,0 +1,17 @@
+import json
+
+import pytest
+
+from labelbox.data.serialization.ndjson.converter import NDJsonConverter
+
+
+@pytest.mark.parametrize("filename", [
+    "tests/data/assets/ndjson/text_entity_import.json",
+    "tests/data/assets/ndjson/text_entity_without_confidence_import.json"
+])
+def test_text_entity_import(filename: str):
+    with open('tests/data/assets/ndjson/text_entity_import.json', 'r') as file:
+        data = json.load(file)
+    res = NDJsonConverter.deserialize(data).as_list()
+    res = list(NDJsonConverter.serialize(res))
+    assert res == data