mirror of
https://github.com/DS4SD/docling.git
synced 2025-07-27 04:24:45 +00:00
Properly propagating image data per page, together with predicted tags in VLM pipeline. This enables correct figure extraction and page numbers in provenances
Signed-off-by: Maksym Lysak <mly@zurich.ibm.com>
This commit is contained in:
Maksym Lysak
parent
01c46e24b1
commit
61bb9dbba2
@ -73,7 +73,7 @@ class VlmPipeline(PaginatedPipeline):
|
||||
|
||||
disable_progress_bars()
|
||||
|
||||
# TODO download the correct model (private repo)
|
||||
# TODO: download the correct model (private repo)
|
||||
download_path = snapshot_download(
|
||||
repo_id="ds4sd/xxx",
|
||||
force_download=force,
|
||||
@ -95,30 +95,17 @@ class VlmPipeline(PaginatedPipeline):
|
||||
with TimeRecorder(conv_res, "doc_assemble", scope=ProfilingScope.DOCUMENT):
|
||||
|
||||
# Read and concatenate the page doctags:
|
||||
document_tags = ""
|
||||
# document_tags = ""
|
||||
page_tags = []
|
||||
page_images = []
|
||||
|
||||
for page in conv_res.pages:
|
||||
if page.predictions.doctags is not None:
|
||||
document_tags += page.predictions.doctags.tag_string
|
||||
page_tags.append(page.predictions.doctags.tag_string)
|
||||
page_images.append(page.image)
|
||||
|
||||
conv_res.document = self._turn_tags_into_doc(document_tags, page.image)
|
||||
"""
|
||||
image_bytes = BytesIO()
|
||||
if page.image:
|
||||
page.image.save(image_bytes, format="PNG")
|
||||
# TODO implement this function
|
||||
conv_res.document = self._turn_tags_into_doc(
|
||||
document_tags, image_bytes.getvalue()
|
||||
)
|
||||
conv_res.document = self._turn_tags_into_doc(page_tags, page_images)
|
||||
|
||||
# Generate page images in the output
|
||||
if self.pipeline_options.generate_page_images:
|
||||
for page in conv_res.pages:
|
||||
assert page.image is not None
|
||||
page_no = page.page_no + 1
|
||||
conv_res.document.pages[page_no].image = ImageRef.from_pil(
|
||||
page.image, dpi=int(72 * self.pipeline_options.images_scale)
|
||||
)
|
||||
"""
|
||||
# Generate images of the requested element types
|
||||
if (
|
||||
self.pipeline_options.generate_picture_images
|
||||
@ -153,9 +140,8 @@ class VlmPipeline(PaginatedPipeline):
|
||||
|
||||
return conv_res
|
||||
|
||||
# def _turn_tags_into_doc(self, xml_content: str, image_bytes: bytes) -> (DoclingDocument, list):
|
||||
def _turn_tags_into_doc(
|
||||
self, xml_content: str, pil_image: Optional[Image] = None
|
||||
self, full_doc_xml_content: list[str], pil_images: list[Image | None]
|
||||
) -> DoclingDocument:
|
||||
def extract_text(tag_content: str) -> str:
|
||||
return re.sub(r"<.*?>", "", tag_content).strip()
|
||||
@ -251,7 +237,7 @@ class VlmPipeline(PaginatedPipeline):
|
||||
cell_text = texts[i + 1]
|
||||
right_offset = 2
|
||||
|
||||
# TODO: Check next element(s) for lcel / ucel / xcel, set properly row_span, col_span
|
||||
# Check next element(s) for lcel / ucel / xcel, set properly row_span, col_span
|
||||
next_right_cell = texts[i + right_offset]
|
||||
|
||||
next_bottom_cell = ""
|
||||
@ -333,6 +319,10 @@ class VlmPipeline(PaginatedPipeline):
|
||||
|
||||
doc = DoclingDocument(name="Example Document")
|
||||
current_group = None
|
||||
|
||||
for pg_idx, xml_content in enumerate(full_doc_xml_content):
|
||||
pil_image = pil_images[pg_idx]
|
||||
page_no = pg_idx + 1
|
||||
lines = xml_content.split("\n")
|
||||
# pil_image = input_image #Image.open(BytesIO(image_bytes))
|
||||
bounding_boxes = []
|
||||
@ -351,7 +341,9 @@ class VlmPipeline(PaginatedPipeline):
|
||||
parent=current_group,
|
||||
prov=(
|
||||
# [ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)]
|
||||
ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)
|
||||
ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
if prov_item
|
||||
else None
|
||||
),
|
||||
@ -361,14 +353,18 @@ class VlmPipeline(PaginatedPipeline):
|
||||
prov_item = extract_bounding_box(line)
|
||||
if prov_item:
|
||||
bounding_boxes.append((prov_item, "blue"))
|
||||
current_group = doc.add_group(label=GroupLabel.SECTION, name=content)
|
||||
current_group = doc.add_group(
|
||||
label=GroupLabel.SECTION, name=content
|
||||
)
|
||||
doc.add_text(
|
||||
label=DocItemLabel.TITLE,
|
||||
text=content,
|
||||
parent=current_group,
|
||||
prov=(
|
||||
# [ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)]
|
||||
ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)
|
||||
ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
if prov_item
|
||||
else None
|
||||
),
|
||||
@ -379,14 +375,18 @@ class VlmPipeline(PaginatedPipeline):
|
||||
prov_item = extract_bounding_box(line)
|
||||
if prov_item:
|
||||
bounding_boxes.append((prov_item, "green"))
|
||||
current_group = doc.add_group(label=GroupLabel.SECTION, name=content)
|
||||
current_group = doc.add_group(
|
||||
label=GroupLabel.SECTION, name=content
|
||||
)
|
||||
doc.add_text(
|
||||
label=DocItemLabel.SECTION_HEADER,
|
||||
text=content,
|
||||
parent=current_group,
|
||||
prov=(
|
||||
# [ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)]
|
||||
ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)
|
||||
ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
if prov_item
|
||||
else None
|
||||
),
|
||||
@ -411,7 +411,9 @@ class VlmPipeline(PaginatedPipeline):
|
||||
parent=current_group,
|
||||
prov=(
|
||||
# [ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)]
|
||||
ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)
|
||||
ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
if prov_item
|
||||
else None
|
||||
),
|
||||
@ -428,7 +430,9 @@ class VlmPipeline(PaginatedPipeline):
|
||||
parent=current_group,
|
||||
prov=(
|
||||
# [ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)]
|
||||
ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)
|
||||
ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
if prov_item
|
||||
else None
|
||||
),
|
||||
@ -445,7 +449,9 @@ class VlmPipeline(PaginatedPipeline):
|
||||
parent=current_group,
|
||||
prov=(
|
||||
# [ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)]
|
||||
ProvenanceItem(bbox=prov_item, charspan=(0, 0), page_no=1)
|
||||
ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
if prov_item
|
||||
else None
|
||||
),
|
||||
@ -469,12 +475,16 @@ class VlmPipeline(PaginatedPipeline):
|
||||
doc.add_picture(
|
||||
parent=current_group,
|
||||
image=ImageRef.from_pil(image=cropped_image, dpi=300),
|
||||
prov=ProvenanceItem(bbox=bbox, charspan=(0, 0), page_no=1),
|
||||
prov=ProvenanceItem(
|
||||
bbox=bbox, charspan=(0, 0), page_no=page_no
|
||||
),
|
||||
)
|
||||
else:
|
||||
doc.add_picture(
|
||||
parent=current_group,
|
||||
prov=ProvenanceItem(bbox=bbox, charspan=(0, 0), page_no=1),
|
||||
prov=ProvenanceItem(
|
||||
bbox=bbox, charspan=(0, 0), page_no=page_no
|
||||
),
|
||||
)
|
||||
elif line.startswith("<list>"):
|
||||
content = extract_text(line)
|
||||
@ -483,7 +493,7 @@ class VlmPipeline(PaginatedPipeline):
|
||||
if prov_item:
|
||||
bounding_boxes.append((prov_item, "brown"))
|
||||
prov_item_inst = ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=1
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
doc.add_text(
|
||||
label=DocItemLabel.LIST_ITEM,
|
||||
@ -499,7 +509,7 @@ class VlmPipeline(PaginatedPipeline):
|
||||
if prov_item:
|
||||
bounding_boxes.append((prov_item, "magenta"))
|
||||
prov_item_inst = ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=1
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
doc.add_text(
|
||||
label=DocItemLabel.PARAGRAPH,
|
||||
@ -514,7 +524,7 @@ class VlmPipeline(PaginatedPipeline):
|
||||
if prov_item:
|
||||
bounding_boxes.append((prov_item, "gray"))
|
||||
prov_item_inst = ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=1
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
doc.add_text(
|
||||
label=DocItemLabel.CHECKBOX_UNSELECTED,
|
||||
@ -530,7 +540,7 @@ class VlmPipeline(PaginatedPipeline):
|
||||
if prov_item:
|
||||
bounding_boxes.append((prov_item, "black"))
|
||||
prov_item_inst = ProvenanceItem(
|
||||
bbox=prov_item, charspan=(0, 0), page_no=1
|
||||
bbox=prov_item, charspan=(0, 0), page_no=page_no
|
||||
)
|
||||
doc.add_text(
|
||||
label=DocItemLabel.CHECKBOX_SELECTED,
|
||||
|
||||
@ -6,10 +6,11 @@ from docling.datamodel.pipeline_options import PdfPipelineOptions
|
||||
from docling.document_converter import DocumentConverter, PdfFormatOption
|
||||
from docling.pipeline.vlm_pipeline import VlmPipeline
|
||||
|
||||
# source = "https://arxiv.org/pdf/2408.09869" # document per local path or URL
|
||||
source = "https://arxiv.org/pdf/2408.09869" # document per local path or URL
|
||||
# source = "tests/data/2305.03393v1-pg9-img.png"
|
||||
source = "tests/data/2305.03393v1-pg9.pdf"
|
||||
# source = "page.png"
|
||||
# source = "tests/data/2305.03393v1-pg9.pdf"
|
||||
# source = "demo_data/page.png"
|
||||
# source = "demo_data/original_tables.pdf"
|
||||
|
||||
pipeline_options = PdfPipelineOptions()
|
||||
pipeline_options.generate_page_images = True
|
||||
|
||||
Loading…
Reference in New Issue
Block a user