Figure export

In [ ]:

import logging
import time
from pathlib import Path

import logging import time from pathlib import Path

In [ ]:

from docling_core.types.doc import ImageRefMode, PictureItem, TableItem

from docling_core.types.doc import ImageRefMode, PictureItem, TableItem

In [ ]:

from docling.datamodel.base_models import FigureElement, InputFormat, Table
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption

from docling.datamodel.base_models import FigureElement, InputFormat, Table from docling.datamodel.pipeline_options import PdfPipelineOptions from docling.document_converter import DocumentConverter, PdfFormatOption

In [ ]:

_log = logging.getLogger(__name__)

_log = logging.getLogger(__name__)

In [ ]:

IMAGE_RESOLUTION_SCALE = 2.0

IMAGE_RESOLUTION_SCALE = 2.0

In [ ]:

def main():
    logging.basicConfig(level=logging.INFO)

    input_doc_path = Path("./tests/data/2206.01062.pdf")
    output_dir = Path("scratch")

    # Important: For operating with page images, we must keep them, otherwise the DocumentConverter
    # will destroy them for cleaning up memory.
    # This is done by setting PdfPipelineOptions.images_scale, which also defines the scale of images.
    # scale=1 correspond of a standard 72 DPI image
    # The PdfPipelineOptions.generate_* are the selectors for the document elements which will be enriched
    # with the image field
    pipeline_options = PdfPipelineOptions()
    pipeline_options.images_scale = IMAGE_RESOLUTION_SCALE
    pipeline_options.generate_page_images = True
    pipeline_options.generate_picture_images = True

    doc_converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
        }
    )

    start_time = time.time()

    conv_res = doc_converter.convert(input_doc_path)

    output_dir.mkdir(parents=True, exist_ok=True)
    doc_filename = conv_res.input.file.stem

    # Save page images
    for page_no, page in conv_res.document.pages.items():
        page_no = page.page_no
        page_image_filename = output_dir / f"{doc_filename}-{page_no}.png"
        with page_image_filename.open("wb") as fp:
            page.image.pil_image.save(fp, format="PNG")

    # Save images of figures and tables
    table_counter = 0
    picture_counter = 0
    for element, _level in conv_res.document.iterate_items():
        if isinstance(element, TableItem):
            table_counter += 1
            element_image_filename = (
                output_dir / f"{doc_filename}-table-{table_counter}.png"
            )
            with element_image_filename.open("wb") as fp:
                element.get_image(conv_res.document).save(fp, "PNG")

        if isinstance(element, PictureItem):
            picture_counter += 1
            element_image_filename = (
                output_dir / f"{doc_filename}-picture-{picture_counter}.png"
            )
            with element_image_filename.open("wb") as fp:
                element.get_image(conv_res.document).save(fp, "PNG")

    # Save markdown with embedded pictures
    content_md = conv_res.document.export_to_markdown(image_mode=ImageRefMode.EMBEDDED)
    md_filename = output_dir / f"{doc_filename}-with-images.md"
    with md_filename.open("w") as fp:
        fp.write(content_md)

    end_time = time.time() - start_time

    _log.info(f"Document converted and figures exported in {end_time:.2f} seconds.")

def main(): logging.basicConfig(level=logging.INFO) input_doc_path = Path("./tests/data/2206.01062.pdf") output_dir = Path("scratch") # Important: For operating with page images, we must keep them, otherwise the DocumentConverter # will destroy them for cleaning up memory. # This is done by setting PdfPipelineOptions.images_scale, which also defines the scale of images. # scale=1 correspond of a standard 72 DPI image # The PdfPipelineOptions.generate_* are the selectors for the document elements which will be enriched # with the image field pipeline_options = PdfPipelineOptions() pipeline_options.images_scale = IMAGE_RESOLUTION_SCALE pipeline_options.generate_page_images = True pipeline_options.generate_picture_images = True doc_converter = DocumentConverter( format_options={ InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options) } ) start_time = time.time() conv_res = doc_converter.convert(input_doc_path) output_dir.mkdir(parents=True, exist_ok=True) doc_filename = conv_res.input.file.stem # Save page images for page_no, page in conv_res.document.pages.items(): page_no = page.page_no page_image_filename = output_dir / f"{doc_filename}-{page_no}.png" with page_image_filename.open("wb") as fp: page.image.pil_image.save(fp, format="PNG") # Save images of figures and tables table_counter = 0 picture_counter = 0 for element, _level in conv_res.document.iterate_items(): if isinstance(element, TableItem): table_counter += 1 element_image_filename = ( output_dir / f"{doc_filename}-table-{table_counter}.png" ) with element_image_filename.open("wb") as fp: element.get_image(conv_res.document).save(fp, "PNG") if isinstance(element, PictureItem): picture_counter += 1 element_image_filename = ( output_dir / f"{doc_filename}-picture-{picture_counter}.png" ) with element_image_filename.open("wb") as fp: element.get_image(conv_res.document).save(fp, "PNG") # Save markdown with embedded pictures content_md = conv_res.document.export_to_markdown(image_mode=ImageRefMode.EMBEDDED) md_filename = output_dir / f"{doc_filename}-with-images.md" with md_filename.open("w") as fp: fp.write(content_md) end_time = time.time() - start_time _log.info(f"Document converted and figures exported in {end_time:.2f} seconds.")

In [ ]:

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()