feat: add table annotations (#304)

* feat: add table annotations

Signed-off-by: Panos Vagenas <[email protected]>

* refactor annotation types

Signed-off-by: Panos Vagenas <[email protected]>

* expand to HTML

Signed-off-by: Panos Vagenas <[email protected]>

* introduce annotation serializer

Signed-off-by: Panos Vagenas <[email protected]>

* Update dummy_doc.yaml

Signed-off-by: Panos Vagenas <[email protected]>

---------

Signed-off-by: Panos Vagenas <[email protected]>
Signed-off-by: Panos Vagenas <[email protected]>

Author: vagenas

docling_core/transforms/serializer/base.py

@@ -239,6 +239,15 @@ def serialize_captions(
         """Serialize the item's captions."""
         ...
 
+    @abstractmethod
+    def serialize_annotations(
+        self,
+        item: DocItem,
+        **kwargs: Any,
+    ) -> SerializationResult:
+        """Serialize the item's annotations."""
+        ...
+
     @abstractmethod
     def get_excluded_refs(self, **kwargs: Any) -> set[str]:
         """Get references to excluded items."""
@@ -257,3 +266,18 @@ class BaseSerializerProvider(ABC):
     def get_serializer(self, doc: DoclingDocument) -> BaseDocSerializer:
         """Get a the associated serializer."""
         ...
+
+
+class BaseAnnotationSerializer(ABC):
+    """Base class for annotation serializers."""
+
+    @abstractmethod
+    def serialize(
+        self,
+        *,
+        item: DocItem,
+        doc: DoclingDocument,
+        **kwargs: Any,
+    ) -> SerializationResult:
+        """Serializes the passed annotation."""
+        ...

docling_core/transforms/serializer/common.py

@@ -15,6 +15,7 @@
 from typing_extensions import Self, override
 
 from docling_core.transforms.serializer.base import (
+    BaseAnnotationSerializer,
     BaseDocSerializer,
     BaseFallbackSerializer,
     BaseFormSerializer,
@@ -30,6 +31,7 @@
 from docling_core.types.doc.document import (
     DOCUMENT_TOKENS_EXPORT_LABELS,
     ContentLayer,
+    DescriptionAnnotation,
     DocItem,
     DoclingDocument,
     FloatingItem,
@@ -41,9 +43,9 @@
     OrderedList,
     PictureClassificationData,
     PictureDataType,
-    PictureDescriptionData,
     PictureItem,
     PictureMoleculeData,
+    TableAnnotationType,
     TableItem,
     TextItem,
     UnorderedList,
@@ -122,7 +124,9 @@ def _iterate_items(
         yield item
 
 
-def _get_picture_annotation_text(annotation: PictureDataType) -> Optional[str]:
+def _get_annotation_text(
+    annotation: Union[PictureDataType, TableAnnotationType],
+) -> Optional[str]:
     result = None
     if isinstance(annotation, PictureClassificationData):
         predicted_class = (
@@ -132,7 +136,7 @@ def _get_picture_annotation_text(annotation: PictureDataType) -> Optional[str]:
         )
         if predicted_class is not None:
             result = predicted_class.replace("_", " ")
-    elif isinstance(annotation, PictureDescriptionData):
+    elif isinstance(annotation, DescriptionAnnotation):
         result = annotation.text
     elif isinstance(annotation, PictureMoleculeData):
         result = annotation.smi
@@ -211,6 +215,8 @@ class DocSerializer(BaseModel, BaseDocSerializer):
     list_serializer: BaseListSerializer
     inline_serializer: BaseInlineSerializer
 
+    annotation_serializer: BaseAnnotationSerializer
+
     params: CommonParams = CommonParams()
 
     _excluded_refs_cache: dict[str, set[str]] = {}
@@ -505,6 +511,19 @@ def serialize_captions(
             text_res = ""
         return create_ser_result(text=text_res, span_source=results)
 
+    @override
+    def serialize_annotations(
+        self,
+        item: DocItem,
+        **kwargs: Any,
+    ) -> SerializationResult:
+        """Serialize the item's annotations."""
+        return self.annotation_serializer.serialize(
+            item=item,
+            doc=self.doc,
+            **kwargs,
+        )
+
     def _get_applicable_pages(self) -> Optional[list[int]]:
         pages = {
             item.prov[0].page_no: ...

docling_core/transforms/serializer/doctags.py

@@ -7,6 +7,7 @@
 from typing_extensions import override
 
 from docling_core.transforms.serializer.base import (
+    BaseAnnotationSerializer,
     BaseDocSerializer,
     BaseFallbackSerializer,
     BaseFormSerializer,
@@ -460,6 +461,15 @@ def serialize(
         return create_ser_result()
 
 
+class DocTagsAnnotationSerializer(BaseAnnotationSerializer):
+    """DocTags-specific annotation serializer."""
+
+    @override
+    def serialize(self, *, item: DocItem, **kwargs: Any) -> SerializationResult:
+        """Serializes the item's annotations."""
+        return create_ser_result()
+
+
 class DocTagsDocSerializer(DocSerializer):
     """DocTags-specific document serializer."""
 
@@ -473,6 +483,8 @@ class DocTagsDocSerializer(DocSerializer):
     list_serializer: BaseListSerializer = DocTagsListSerializer()
     inline_serializer: BaseInlineSerializer = DocTagsInlineSerializer()
 
+    annotation_serializer: BaseAnnotationSerializer = DocTagsAnnotationSerializer()
+
     params: DocTagsParams = DocTagsParams()
 
     @override

docling_core/transforms/serializer/html.py

@@ -21,6 +21,7 @@
 from typing_extensions import override
 
 from docling_core.transforms.serializer.base import (
+    BaseAnnotationSerializer,
     BaseDocSerializer,
     BaseFallbackSerializer,
     BaseFormSerializer,
@@ -35,7 +36,7 @@
 from docling_core.transforms.serializer.common import (
     CommonParams,
     DocSerializer,
-    _get_picture_annotation_text,
+    _get_annotation_text,
     create_ser_result,
 )
 from docling_core.transforms.serializer.html_styles import (
@@ -47,6 +48,7 @@
 from docling_core.types.doc.document import (
     CodeItem,
     ContentLayer,
+    DescriptionAnnotation,
     DocItem,
     DoclingDocument,
     FloatingItem,
@@ -59,7 +61,9 @@
     ListItem,
     NodeItem,
     OrderedList,
+    PictureClassificationData,
     PictureItem,
+    PictureMoleculeData,
     PictureTabularChartData,
     SectionHeaderItem,
     TableCell,
@@ -758,14 +762,7 @@ class HTMLFallbackSerializer(BaseFallbackSerializer):
     """HTML-specific fallback serializer."""
 
     @override
-    def serialize(
-        self,
-        *,
-        item: NodeItem,
-        doc_serializer: "BaseDocSerializer",
-        doc: DoclingDocument,
-        **kwargs: Any,
-    ) -> SerializationResult:
+    def serialize(self, *, item: NodeItem, **kwargs: Any) -> SerializationResult:
         """Fallback serializer for items not handled by other serializers."""
         if isinstance(item, DocItem):
             return create_ser_result(
@@ -777,6 +774,42 @@ def serialize(
             return create_ser_result()
 
 
+class HTMLAnnotationSerializer(BaseModel, BaseAnnotationSerializer):
+    """HTML-specific annotation serializer."""
+
+    def serialize(
+        self,
+        *,
+        item: DocItem,
+        doc: DoclingDocument,
+        **kwargs: Any,
+    ) -> SerializationResult:
+        """Serializes the passed annotation to HTML format."""
+        res_parts: list[SerializationResult] = []
+        for ann in item.get_annotations():
+            if isinstance(
+                ann,
+                (PictureClassificationData, DescriptionAnnotation, PictureMoleculeData),
+            ):
+                if ann_text := _get_annotation_text(ann):
+                    text_dir = get_text_direction(ann_text)
+                    dir_str = f' dir="{text_dir}"' if text_dir == "rtl" else ""
+                    ann_ser_res = create_ser_result(
+                        text=(
+                            f'<div data-annotation-kind="{ann.kind}"{dir_str}>'
+                            f"{html.escape(ann_text)}"
+                            f"</div>"
+                        ),
+                        span_source=item,
+                    )
+                    res_parts.append(ann_ser_res)
+
+        return create_ser_result(
+            text=" ".join([r.text for r in res_parts if r.text]),
+            span_source=res_parts,
+        )
+
+
 class HTMLDocSerializer(DocSerializer):
     """HTML-specific document serializer."""
 
@@ -790,6 +823,8 @@ class HTMLDocSerializer(DocSerializer):
     list_serializer: BaseListSerializer = HTMLListSerializer()
     inline_serializer: BaseInlineSerializer = HTMLInlineSerializer()
 
+    annotation_serializer: BaseAnnotationSerializer = HTMLAnnotationSerializer()
+
     params: HTMLParams = HTMLParams()
 
     @override
@@ -968,20 +1003,13 @@ def serialize_captions(
                     results.append(cap_ser_res)
 
         if params.include_annotations and item.self_ref not in excluded_refs:
-            if isinstance(item, PictureItem):
-                for ann in item.annotations:
-                    if ann_text := _get_picture_annotation_text(annotation=ann):
-                        text_dir = get_text_direction(ann_text)
-                        dir_str = f' dir="{text_dir}"' if text_dir == "rtl" else ""
-                        ann_ser_res = create_ser_result(
-                            text=(
-                                f'<div data-annotation-kind="{ann.kind}"{dir_str}>'
-                                f"{html.escape(ann_text)}"
-                                f"</div>"
-                            ),
-                            span_source=item,
-                        )
-                        results.append(ann_ser_res)
+            if isinstance(item, (PictureItem, TableItem)):
+                ann_res = self.serialize_annotations(
+                    item=item,
+                    **kwargs,
+                )
+                if ann_res.text:
+                    results.append(ann_res)
 
         text_res = params.caption_delim.join([r.text for r in results])
         if text_res: