Added get_image to DocItem and FloatingItem

Signed-off-by: Shubham Gupta <[email protected]>

docling_core/types/doc/document.py

@@ -551,6 +551,23 @@ def get_location_tokens(
 
         return location
 
+    def get_image(self, doc: "DoclingDocument") -> Optional[PILImage.Image]:
+        """Returns the image of this DocItem if the document stores page images."""
+        if not len(self.prov):
+            return None
+
+        page = doc.pages.get(self.prov[0].page_no)
+        if page is None or page.size is None or page.image is None:
+            return None
+
+        page_image = page.image.pil_image
+        crop_bbox = (
+            self.prov[0]
+            .bbox.to_top_left_origin(page_height=page.size.height)
+            .scaled(scale=page_image.height / page.size.height)
+        )
+        return page_image.crop(crop_bbox.as_tuple())
+
 
 class TextItem(DocItem):
     """TextItem."""
@@ -633,6 +650,12 @@ def caption_text(self, doc: "DoclingDocument") -> str:
             text += cap.resolve(doc).text
         return text
 
+    def get_image(self, doc: "DoclingDocument") -> Optional[PILImage.Image]:
+        """Returns the image corresponding to FloatingItem."""
+        if self.image is not None:
+            return self.image.pil_image
+        return super().get_image(doc=doc)
+
 
 class PictureItem(FloatingItem):
     """PictureItem."""