add classification and description in PictureData

Signed-off-by: Michele Dolfi <[email protected]>

Author: dolfim-ibm

docling_core/types/experimental/__init__.py

@@ -7,7 +7,6 @@
 
 from .base import BoundingBox, CoordOrigin, Size
 from .document import (
-    BasePictureData,
     BaseTableData,
     DescriptionItem,
     DocItem,
@@ -19,6 +18,7 @@
     KeyValueItem,
     NodeItem,
     PageItem,
+    PictureData,
     PictureItem,
     ProvenanceItem,
     RefItem,

docling_core/types/experimental/document.py

@@ -46,8 +46,26 @@
 }
 
 
-class BasePictureData(BaseModel):  # TBD
-    """BasePictureData."""
+class PictureClassificationData(BaseModel):
+    """PictureClassificationData."""
+
+    provenance: str
+    predicted_class: str
+    confidence: float
+
+
+class PictureDescriptionData(BaseModel):
+    """PictureDescriptionData."""
+
+    text: str
+    provenance: str = ""
+
+
+class PictureData(BaseModel):
+    """PictureData."""
+
+    classification: Optional[PictureClassificationData] = None
+    description: Optional[PictureDescriptionData] = None
 
 
 class TableCell(BaseModel):
@@ -384,7 +402,7 @@ class PictureItem(FloatingItem):
 
     label: typing.Literal[DocItemLabel.PICTURE] = DocItemLabel.PICTURE
 
-    data: BasePictureData
+    data: PictureData
 
     def export_to_document_tokens(
         self,
@@ -863,14 +881,14 @@ def add_table(
 
     def add_picture(
         self,
-        data: BasePictureData,
+        data: PictureData,
         caption: Optional[Union[TextItem, RefItem]] = None,
         prov: Optional[ProvenanceItem] = None,
         parent: Optional[GroupItem] = None,
     ):
         """add_picture.
 
-        :param data: BasePictureData:
+        :param data: PictureData:
         :param caption: Optional[Union[TextItem:
         :param RefItem]]:  (Default value = None)
         :param prov: Optional[ProvenanceItem]:  (Default value = None)

test/data/docling_document/unit/PictureItem.yaml

@@ -1,6 +1,8 @@
 captions: []
 children: []
-data: {}
+data:
+  classification: null
+  description: null
 footnotes: []
 image: null
 label: picture

test/data/experimental/2206.01062.experimental.yaml

@@ -216,7 +216,9 @@ pictures:
 - captions:
   - $ref: '#/texts/12'
   children: []
-  data: {}
+  data:
+    classification: null
+    description: null
   footnotes: []
   image: null
   label: picture
@@ -238,7 +240,9 @@ pictures:
 - captions:
   - $ref: '#/texts/39'
   children: []
-  data: {}
+  data:
+    classification: null
+    description: null
   footnotes: []
   image: null
   label: picture
@@ -260,7 +264,9 @@ pictures:
 - captions:
   - $ref: '#/texts/70'
   children: []
-  data: {}
+  data:
+    classification: null
+    description: null
   footnotes: []
   image: null
   label: picture
@@ -282,7 +288,9 @@ pictures:
 - captions:
   - $ref: '#/texts/76'
   children: []
-  data: {}
+  data:
+    classification: null
+    description: null
   footnotes: []
   image: null
   label: picture
@@ -304,7 +312,9 @@ pictures:
 - captions:
   - $ref: '#/texts/117'
   children: []
-  data: {}
+  data:
+    classification: null
+    description: null
   footnotes: []
   image: null
   label: picture

test/data/experimental/dummy_doc.yaml

@@ -133,10 +133,14 @@ pictures: # All pictures...
       $ref: "#/body"
     captions:
       - $ref: "#/texts/3"
-    data: # BaseFigureData Type
-      classification: "illustration"
-      confidence: 0.78
-      description: "...."
+    data: # PictureData Type
+      classification:
+        provenance: model1
+        predicted_class: "illustration"
+        confidence: 0.78
+      description:
+        text: "..."
+        provenance: "model2"
       # content structure?
     image:
       mimetype: image/png

test/test_docling_doc.py

@@ -6,14 +6,14 @@
 
 from docling_core.types.experimental.document import (
     CURRENT_VERSION,
-    BasePictureData,
     BaseTableData,
     DescriptionItem,
     DocItem,
     DoclingDocument,
     FloatingItem,
     KeyValueItem,
     ListItem,
+    PictureData,
     PictureItem,
     SectionHeaderItem,
     TableCell,
@@ -107,7 +107,7 @@ def verify(dc, obj):
         elif dc is PictureItem:
             obj = dc(
                 self_ref="#",
-                data=BasePictureData(),
+                data=PictureData(),
             )
             verify(dc, obj)
 
@@ -203,7 +203,7 @@ def _test_serialize_and_reload(doc):
     assert doc_reload is not doc  # can't be identical
 
 
-def _test_export_methods(doc):
+def _test_export_methods(doc: DoclingDocument):
     ### Iterate all elements
     doc.print_element_tree()
     ## Export stuff
@@ -341,7 +341,7 @@ def _construct_doc() -> DoclingDocument:
     fig_caption = doc.add_text(
         label=DocItemLabel.CAPTION, text="This is the caption of figure 1."
     )
-    fig_item = doc.add_picture(data=BasePictureData(), caption=fig_caption)
+    fig_item = doc.add_picture(data=PictureData(), caption=fig_caption)
 
     return doc