mirror of
https://github.com/DS4SD/docling.git
synced 2025-07-26 12:04:31 +00:00
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import json
|
|
import logging
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
from docling.datamodel.base_models import (
|
|
AssembleOptions,
|
|
BoundingBox,
|
|
ConversionStatus,
|
|
CoordOrigin,
|
|
PipelineOptions,
|
|
)
|
|
from docling.datamodel.document import ConvertedDocument, DocumentConversionInput
|
|
from docling.document_converter import DocumentConverter
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
def export_figures(
|
|
converted_docs: Iterable[ConvertedDocument],
|
|
output_dir: Path,
|
|
):
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
success_count = 0
|
|
failure_count = 0
|
|
|
|
for doc in converted_docs:
|
|
if doc.status == ConversionStatus.SUCCESS:
|
|
success_count += 1
|
|
doc_filename = doc.input.file.stem
|
|
|
|
for page in doc.pages:
|
|
page_no = page.page_no + 1
|
|
page_image_filename = output_dir / f"{doc_filename}-{page_no}.png"
|
|
with page_image_filename.open("wb") as fp:
|
|
page.image.save(fp, format="PNG")
|
|
|
|
for fig_ix, fig in enumerate(doc.output.figures):
|
|
page_no = fig.prov[0].page
|
|
page_ix = page_no - 1
|
|
x0, y0, x1, y1 = fig.prov[0].bbox
|
|
crop_bbox = BoundingBox(
|
|
l=x0, b=y0, r=x1, t=y1, coord_origin=CoordOrigin.BOTTOMLEFT
|
|
).to_top_left_origin(page_height=doc.pages[page_ix].size.height)
|
|
|
|
cropped_im = doc.pages[page_ix].image.crop(crop_bbox.as_tuple())
|
|
fig_image_filename = output_dir / f"{doc_filename}-fig{fig_ix+1}.png"
|
|
with fig_image_filename.open("wb") as fp:
|
|
cropped_im.save(fp, "PNG")
|
|
|
|
else:
|
|
_log.info(f"Document {doc.input.file} failed to convert.")
|
|
failure_count += 1
|
|
|
|
_log.info(
|
|
f"Processed {success_count + failure_count} docs, of which {failure_count} failed"
|
|
)
|
|
|
|
|
|
def main():
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
input_doc_paths = [
|
|
Path("./test/data/2206.01062.pdf"),
|
|
]
|
|
|
|
input_files = DocumentConversionInput.from_paths(input_doc_paths)
|
|
|
|
assemble_options = AssembleOptions()
|
|
assemble_options.remove_page_images = False
|
|
|
|
doc_converter = DocumentConverter(assemble_options=assemble_options)
|
|
|
|
start_time = time.time()
|
|
|
|
converted_docs = doc_converter.convert(input_files)
|
|
export_figures(converted_docs, output_dir=Path("./scratch"))
|
|
|
|
end_time = time.time() - start_time
|
|
|
|
_log.info(f"All documents were converted in {end_time:.2f} seconds.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|