feat: allow images in doctags deserializer to be optional and support multipage (#225)

* feat: allow images in doctags deserializer to be optional and support multipage

Signed-off-by: Yusik Kim <[email protected]>

* chore: remove commented out lines

Signed-off-by: Yusik Kim <[email protected]>

* Avoid typecasting

Signed-off-by: Christoph Auer <[email protected]>

---------

Signed-off-by: Yusik Kim <[email protected]>
Signed-off-by: Christoph Auer <[email protected]>
Co-authored-by: Christoph Auer <[email protected]>

Author: kmyusk

docling_core/types/doc/document.py

@@ -542,32 +542,58 @@ class DocTagsDocument(BaseModel):
 
     @classmethod
     def from_doctags_and_image_pairs(
-        cls, doctags: List[Union[Path, str]], images: List[Union[Path, PILImage.Image]]
+        cls,
+        doctags: typing.Sequence[Union[Path, str]],
+        images: Optional[List[Union[Path, PILImage.Image]]],
     ):
         """from_doctags_and_image_pairs."""
-        if len(doctags) != len(images):
+        if images is not None and len(doctags) != len(images):
             raise ValueError("Number of page doctags must be equal to page images!")
         doctags_doc = cls()
 
         pages = []
-        for dt, img in zip(doctags, images):
+
+        for ix, dt in enumerate(doctags):
             if isinstance(dt, Path):
                 with dt.open("r") as fp:
                     dt = fp.read()
             elif isinstance(dt, str):
                 pass
 
-            if isinstance(img, Path):
-                img = PILImage.open(img)
-            elif isinstance(dt, PILImage.Image):
-                pass
+            img = None
+            if images is not None:
+                img = images[ix]
+
+                if isinstance(img, Path):
+                    img = PILImage.open(img)
+                elif isinstance(img, PILImage.Image):
+                    pass
 
             page = DocTagsPage(tokens=dt, image=img)
             pages.append(page)
 
         doctags_doc.pages = pages
         return doctags_doc
 
+    @classmethod
+    def from_multipage_doctags_and_images(
+        cls,
+        doctags: Union[Path, str],
+        images: Optional[List[Union[Path, PILImage.Image]]],
+    ):
+        """From doctags with `<page_break>` and corresponding list of page images."""
+        if isinstance(doctags, Path):
+            with doctags.open("r") as fp:
+                doctags = fp.read()
+        dt_list = (
+            doctags.removeprefix(f"<{DocumentToken.DOCUMENT.value}>")
+            .removesuffix(f"</{DocumentToken.DOCUMENT.value}>")
+            .split(f"<{DocumentToken.PAGE_BREAK.value}>")
+        )
+        dt_list = [el.strip() for el in dt_list]
+
+        return cls.from_doctags_and_image_pairs(dt_list, images)
+
 
 class ProvenanceItem(BaseModel):
     """ProvenanceItem."""

test/test_doctags_load.py

@@ -30,10 +30,26 @@ def test_doctags_load_from_memory():
     # print(doc.export_to_html())
 
 
+def test_doctags_load_without_image():
+    doc = DoclingDocument(name="Document")
+    doctags = Path("test/data/doc/page_with_pic.dt").open("r").read()
+    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], None)
+    doc.load_from_doctags(doctags_doc)
+    # print(doc.export_to_html())
+
+
 def test_doctags_load_for_kv_region():
     doc = DoclingDocument(name="Document")
     doctags = Path("test/data/doc/doc_with_kv.dt").open("r").read()
     image = PILImage.open(Path("test/data/doc/doc_with_kv.png"))
     doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
     doc.load_from_doctags(doctags_doc)
     # print(doc.export_to_html())
+
+
+def test_multipage_doctags_load():
+    doc = DoclingDocument(name="Document")
+    doctags = Path("test/data/doc/2206.01062.yaml.dt").open("r").read()
+    doctags_doc = DocTagsDocument.from_multipage_doctags_and_images(doctags, None)
+    doc.load_from_doctags(doctags_doc)
+    # print(doc.export_to_html())