mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-08-01 23:12:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

move expansion_factor to base class

Browse Source 
Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
Michele Dolfi 2025-01-24 16:23:44 +01:00
parent a476d15c2c
commit 9afd8369eb
2 changed files with 18 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
docling/models/base_model.py
View File

@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
from typing import Any, Generic, Iterable, Optional
from docling_core.types.doc import DoclingDocument, NodeItem, TextItem
from docling_core.types.doc import BoundingBox, DoclingDocument, NodeItem, TextItem
from typing_extensions import TypeVar
from docling.datamodel.base_models import ItemAndImageEnrichmentElement, Page
@ -53,6 +53,7 @@ class BaseItemAndImageEnrichmentModel(
):
images_scale: float
expansion_factor: float = 0.0
def prepare_element(
self, conv_res: ConversionResult, element: NodeItem
@ -62,8 +63,22 @@ class BaseItemAndImageEnrichmentModel(
assert isinstance(element, TextItem)
element_prov = element.prov[0]
bbox = element_prov.bbox
width = bbox.r - bbox.l
height = bbox.t - bbox.b
# TODO: move to a utility in the BoundingBox class
expanded_bbox = BoundingBox(
l=bbox.l - width * self.expansion_factor,
t=bbox.t + height * self.expansion_factor,
r=bbox.r + width * self.expansion_factor,
b=bbox.b - height * self.expansion_factor,
coord_origin=bbox.coord_origin,
)
page_ix = element_prov.page_no - 1
cropped_image = conv_res.pages[page_ix].get_image(
scale=self.images_scale, cropbox=element_prov.bbox
scale=self.images_scale, cropbox=expanded_bbox
)
return ItemAndImageEnrichmentElement(item=element, image=cropped_image)

32
docling/models/code_formula_model.py
View File

@ -3,7 +3,6 @@ from pathlib import Path
from typing import Iterable, List, Literal, Optional, Tuple, Union
from docling_core.types.doc import (
BoundingBox,
CodeItem,
DocItemLabel,
DoclingDocument,
@ -15,7 +14,6 @@ from PIL import Image
from pydantic import BaseModel
from docling.datamodel.base_models import ItemAndImageEnrichmentElement
from docling.datamodel.document import ConversionResult
from docling.datamodel.pipeline_options import AcceleratorOptions
from docling.models.base_model import BaseItemAndImageEnrichmentModel
from docling.utils.accelerator_utils import decide_device
@ -64,6 +62,7 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
"""
images_scale = 1.66 # = 120 dpi, aligned with training data resolution
expansion_factor = 0.03
def __init__(
self,
@ -201,35 +200,6 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
except ValueError:
return CodeLanguageLabel.UNKNOWN
def prepare_element(
self, conv_res: ConversionResult, element: NodeItem
) -> Optional[ItemAndImageEnrichmentElement]:
if not self.is_processable(doc=conv_res.document, element=element):
return None
assert isinstance(element, TextItem)
element_prov = element.prov[0]
expansion_factor = 0.03
bbox = element_prov.bbox
width = bbox.r - bbox.l
height = bbox.t - bbox.b
expanded_bbox = BoundingBox(
l=bbox.l - width * expansion_factor,
t=bbox.t + height * expansion_factor,
r=bbox.r + width * expansion_factor,
b=bbox.b - height * expansion_factor,
coord_origin=bbox.coord_origin,
)
page_ix = element_prov.page_no - 1
cropped_image = conv_res.pages[page_ix].get_image(
scale=self.images_scale, cropbox=expanded_bbox
)
return ItemAndImageEnrichmentElement(item=element, image=cropped_image)
def __call__(
self,
doc: DoclingDocument,