Merge pull request #987 from Labelbox/VB/document-entity-annotation_AL-5106

Vb/document entity annotation al 5106

Author: vbrodsky

examples/annotation_import/pdf.ipynb

@@ -11,6 +11,8 @@
 from .annotation import VideoObjectAnnotation
 
 from .ner import TextEntity
+from .ner import DocumentEntity
+from .ner import DocumentTextSelection
 
 from .classification import Checklist
 from .classification import ClassificationAnswer

labelbox/data/annotation_types/__init__.py

@@ -6,7 +6,7 @@
 from .classification import Checklist, Dropdown, Radio, Text
 from .feature import FeatureSchema
 from .geometry import Geometry, Rectangle, Point
-from .ner import TextEntity
+from .ner import DocumentEntity, TextEntity
 
 
 class BaseAnnotation(FeatureSchema, abc.ABC):
@@ -51,7 +51,7 @@ class ObjectAnnotation(BaseAnnotation, ConfidenceMixin):
         classifications (Optional[List[ClassificationAnnotation]]): Optional sub classification of the annotation
         extra (Dict[str, Any])
     """
-    value: Union[TextEntity, Geometry]
+    value: Union[TextEntity, DocumentEntity, Geometry]
     classifications: List[ClassificationAnnotation] = []
 
 

labelbox/data/annotation_types/annotation.py

@@ -6,6 +6,7 @@
 
 import labelbox
 from labelbox.data.annotation_types.data.tiled_image import TiledImageData
+from labelbox.data.annotation_types.ner import DocumentEntity
 from labelbox.schema import ontology
 from .annotation import (ClassificationAnnotation, ObjectAnnotation,
                          VideoClassificationAnnotation, VideoObjectAnnotation)

labelbox/data/annotation_types/label.py

@@ -0,0 +1,2 @@
+from .text_entity import TextEntity
+from .document_entity import DocumentEntity, DocumentTextSelection

labelbox/data/annotation_types/ner/__init__.py

@@ -0,0 +1,23 @@
+from typing import List
+
+from pydantic import BaseModel, validator
+
+from labelbox.utils import _CamelCaseMixin
+
+
+class DocumentTextSelection(_CamelCaseMixin, BaseModel):
+    token_ids: List[str]
+    group_id: str
+    page: int
+
+    @validator("page")
+    def validate_page(cls, v):
+        if v < 1:
+            raise ValueError("Page must be greater than 1")
+        return v
+
+
+class DocumentEntity(_CamelCaseMixin, BaseModel):
+    """ Represents a text entity """
+    name: str
+    text_selections: List[DocumentTextSelection]

labelbox/data/annotation_types/ner/document_entity.py

@@ -40,6 +40,7 @@ def serialize(
         Returns:
             A generator for accessing the ndjson representation of the data
         """
+
         for example in NDLabel.from_common(labels):
             res = example.dict(by_alias=True)
             for k, v in list(res.items()):

labelbox/data/annotation_types/ner.py renamed to labelbox/data/annotation_types/ner/text_entity.py

@@ -12,7 +12,7 @@
 from labelbox.data.annotation_types.data.video import VideoData
 
 from ...annotation_types.data import ImageData, TextData, MaskData
-from ...annotation_types.ner import TextEntity
+from ...annotation_types.ner import DocumentEntity, DocumentTextSelection, TextEntity
 from ...annotation_types.types import Cuid
 from ...annotation_types.geometry import Rectangle, Polygon, Line, Point, Mask
 from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoObjectAnnotation
@@ -372,6 +372,33 @@ def from_common(cls,
                    confidence=confidence)
 
 
+class NDDocumentEntity(NDBaseObject, ConfidenceMixin):
+    name: str
+    text_selections: List[DocumentTextSelection]
+
+    def to_common(self) -> DocumentEntity:
+        return DocumentEntity(name=self.name,
+                              text_selections=self.text_selections)
+
+    @classmethod
+    def from_common(cls,
+                    document_entity: DocumentEntity,
+                    classifications: List[ClassificationAnnotation],
+                    name: str,
+                    feature_schema_id: Cuid,
+                    extra: Dict[str, Any],
+                    data: Union[ImageData, TextData],
+                    confidence: Optional[float] = None) -> "NDDocumentEntity":
+
+        return cls(text_selections=document_entity.text_selections,
+                   dataRow=DataRow(id=data.uid),
+                   name=name,
+                   schema_id=feature_schema_id,
+                   uuid=extra.get('uuid'),
+                   classifications=classifications,
+                   confidence=confidence)
+
+
 class NDObject:
 
     @staticmethod
@@ -434,7 +461,8 @@ def lookup_object(
                 Polygon: NDPolygon,
                 Rectangle: NDRectangle,
                 Mask: NDMask,
-                TextEntity: NDTextEntity
+                TextEntity: NDTextEntity,
+                DocumentEntity: NDDocumentEntity,
             }.get(type(annotation.value))
         if result is None:
             raise TypeError(
@@ -444,6 +472,6 @@ def lookup_object(
 
 
 NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
-                     NDTextEntity]
+                     NDTextEntity, NDDocumentEntity]
 
 NDFrameObjectType = NDFrameRectangle, NDFramePoint, NDFrameLine