PR: add ndson test

Labelbox · Mar 14, 2023 · 0b4590f · 0b4590f
1 parent 0c7be25
commit 0b4590f
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 6 deletions.
diff --git a/labelbox/data/serialization/ndjson/objects.py b/labelbox/data/serialization/ndjson/objects.py
@@ -377,11 +377,8 @@ class NDDocumentEntity(NDBaseObject, ConfidenceMixin):
     text_selections: List[DocumentTextSelection]
 
     def to_common(self) -> DocumentEntity:
-        return TextEntity(name=self.name, text_selections=self.text_selections)
-
-        return obj.from_common(annotation.value, subclasses, annotation.name,
-                               annotation.feature_schema_id, annotation.extra,
-                               data, **optional_kwargs)
+        return DocumentEntity(name=self.name,
+                              text_selections=self.text_selections)
 
     @classmethod
     def from_common(cls,

diff --git a/tests/data/assets/ndjson/pdf_document_entity_import.json b/tests/data/assets/ndjson/pdf_document_entity_import.json
@@ -0,0 +1,25 @@
+[
+    {
+        "uuid": "f6879f59-d2b5-49c2-aceb-d9e8dc478673",
+        "dataRow": {
+            "id": "ckrb1sf1i1g7i0ybcdc6oc8ct"
+        },
+        "name": "named_entity",
+        "classifications": [],
+        "textSelections": [
+            {
+                "groupId": "2f4336f4-a07e-4e0a-a9e1-5629b03b719b",
+                "tokenIds": [
+                    "3f984bf3-1d61-44f5-b59a-9658a2e3440f",
+                    "3bf00b56-ff12-4e52-8cc1-08dbddb3c3b8",
+                    "6e1c3420-d4b7-4c5a-8fd6-ead43bf73d80",
+                    "87a43d32-af76-4a1d-b262-5c5f4d5ace3a",
+                    "e8606e8a-dfd9-4c49-a635-ad5c879c75d0",
+                    "67c7c19e-4654-425d-bf17-2adb8cf02c30",
+                    "149c5e80-3e07-49a7-ab2d-29ddfe6a38fa",
+                    "b0e94071-2187-461e-8e76-96c58738a52c"
+                ],
+                "page": 1            }
+        ]
+    }
+]
diff --git a/tests/data/serialization/ndjson/test_document.py b/tests/data/serialization/ndjson/test_document.py
@@ -6,7 +6,7 @@
 def round_dict(data):
     if isinstance(data, dict):
         for key in data:
-            if isinstance(data[key], float):
+            if isinstance(data[key], (int, float)):
                 data[key] = int(data[key])
             elif isinstance(data[key], dict):
                 data[key] = round_dict(data[key])
@@ -28,6 +28,20 @@ def test_pdf():
     f.close()
 
 
+def test_pdf_document_entity():
+    """
+    Tests a pdf file with bbox annotations only
+    """
+    with open('tests/data/assets/ndjson/pdf_document_entity_import.json',
+              'r') as f:
+        data = json.load(f)
+    res = list(NDJsonConverter.deserialize(data))
+    res = list(NDJsonConverter.serialize(res))
+    assert [round_dict(x) for x in res] == [round_dict(x) for x in data]
+
+    f.close()
+
+
 def test_pdf_with_name_only():
     """
     Tests a pdf file with bbox annotations only