fix: only save applicable page images (#226)

fix(document): save document images page by page

Signed-off-by: Clément Doumouro <[email protected]>

Author: ClemDoum

docling_core/types/doc/document.py

@@ -4098,7 +4098,10 @@ def _with_embedded_pictures(self) -> "DoclingDocument":
         return result
 
     def _with_pictures_refs(
-        self, image_dir: Path, reference_path: Optional[Path] = None
+        self,
+        image_dir: Path,
+        page_no: Optional[int],
+        reference_path: Optional[Path] = None,
     ) -> "DoclingDocument":
         """Document with images as refs.
 
@@ -4111,7 +4114,7 @@ def _with_pictures_refs(
         image_dir.mkdir(parents=True, exist_ok=True)
 
         if image_dir.is_dir():
-            for item, level in result.iterate_items(with_groups=False):
+            for item, level in result.iterate_items(page_no=page_no, with_groups=False):
                 if isinstance(item, PictureItem):
 
                     if (
@@ -4211,7 +4214,7 @@ def save_as_json(
             os.makedirs(artifacts_dir, exist_ok=True)
 
         new_doc = self._make_copy_with_refmode(
-            artifacts_dir, image_mode, reference_path=reference_path
+            artifacts_dir, image_mode, page_no=None, reference_path=reference_path
         )
 
         out = new_doc.export_to_dict(
@@ -4254,7 +4257,7 @@ def save_as_yaml(
             os.makedirs(artifacts_dir, exist_ok=True)
 
         new_doc = self._make_copy_with_refmode(
-            artifacts_dir, image_mode, reference_path=reference_path
+            artifacts_dir, image_mode, page_no=None, reference_path=reference_path
         )
 
         out = new_doc.export_to_dict(
@@ -4327,7 +4330,7 @@ def save_as_markdown(
             os.makedirs(artifacts_dir, exist_ok=True)
 
         new_doc = self._make_copy_with_refmode(
-            artifacts_dir, image_mode, reference_path=reference_path
+            artifacts_dir, image_mode, page_no, reference_path=reference_path
         )
 
         md_out = new_doc.export_to_markdown(
@@ -4503,7 +4506,7 @@ def save_as_html(
             os.makedirs(artifacts_dir, exist_ok=True)
 
         new_doc = self._make_copy_with_refmode(
-            artifacts_dir, image_mode, reference_path=reference_path
+            artifacts_dir, image_mode, page_no, reference_path=reference_path
         )
 
         html_out = new_doc.export_to_html(
@@ -4542,14 +4545,15 @@ def _make_copy_with_refmode(
         self,
         artifacts_dir: Path,
         image_mode: ImageRefMode,
+        page_no: Optional[int],
         reference_path: Optional[Path] = None,
     ):
         new_doc = None
         if image_mode == ImageRefMode.PLACEHOLDER:
             new_doc = self
         elif image_mode == ImageRefMode.REFERENCED:
             new_doc = self._with_pictures_refs(
-                image_dir=artifacts_dir, reference_path=reference_path
+                image_dir=artifacts_dir, page_no=page_no, reference_path=reference_path
             )
         elif image_mode == ImageRefMode.EMBEDDED:
             new_doc = self._with_embedded_pictures()

test/test_docling_doc.py

@@ -1349,12 +1349,51 @@ def test_save_pictures():
 
     doc: DoclingDocument = _construct_doc()
 
-    new_doc = doc._with_pictures_refs(image_dir=Path("./test/data/constructed_images/"))
+    new_doc = doc._with_pictures_refs(
+        image_dir=Path("./test/data/constructed_images/"), page_no=None
+    )
 
     img_paths = new_doc._list_images_on_disk()
     assert len(img_paths) == 1, "len(img_paths)!=1"
 
 
+def test_save_pictures_with_page():
+    # Given
+    doc = DoclingDocument(name="Dummy")
+
+    doc.add_page(page_no=1, size=Size(width=2000, height=4000), image=None)
+    doc.add_page(
+        page_no=2,
+        size=Size(width=2000, height=4000),
+    )
+    image = PILImage.new(mode="RGB", size=(200, 400), color=(0, 0, 0))
+    doc.add_picture(
+        image=ImageRef.from_pil(image=image, dpi=72),
+        prov=ProvenanceItem(
+            page_no=2,
+            bbox=BoundingBox(
+                b=0, l=0, r=200, t=400, coord_origin=CoordOrigin.BOTTOMLEFT
+            ),
+            charspan=(1, 2),
+        ),
+    )
+
+    # When
+    with_ref = doc._with_pictures_refs(
+        image_dir=Path("./test/data/constructed_images/"), page_no=1
+    )
+    # Then
+    n_images = len(with_ref._list_images_on_disk())
+    assert n_images == 0
+    # When
+    with_ref = with_ref._with_pictures_refs(
+        image_dir=Path("./test/data/constructed_images/"), page_no=2
+    )
+    n_images = len(with_ref._list_images_on_disk())
+    # Then
+    assert n_images == 1
+
+
 def _normalise_string_wrt_filepaths(instr: str, paths: List[Path]):
 
     for p in paths:
@@ -1406,7 +1445,8 @@ def test_save_to_disk():
     image_dir = Path("./test/data/doc/constructed_images/")
 
     doc_with_references = doc._with_pictures_refs(
-        image_dir=image_dir  # Path("./test/data/constructed_images/")
+        image_dir=image_dir,  # Path("./test/data/constructed_images/")
+        page_no=None,
     )
 
     # paths will be different on different machines, so needs to be kept!