2a2c65bf4f
* Add settings to turn visualization on or off Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Add profiling code to all models Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Refactor and fix profiling codes Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Visualization codes output PNG to debug dir Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Fixes for time logging Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Optimize imports Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Update lockfile Signed-off-by: Christoph Auer <cau@zurich.ibm.com> * Add start_timestamps to ProfilingItem Signed-off-by: Christoph Auer <cau@zurich.ibm.com> --------- Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
from pathlib import Path
|
|
from typing import Iterable, Optional
|
|
|
|
from PIL import ImageDraw
|
|
from pydantic import BaseModel
|
|
|
|
from docling.datamodel.base_models import Page
|
|
from docling.datamodel.document import ConversionResult
|
|
from docling.datamodel.settings import settings
|
|
from docling.models.base_model import BasePageModel
|
|
from docling.utils.profiling import TimeRecorder
|
|
|
|
|
|
class PagePreprocessingOptions(BaseModel):
|
|
images_scale: Optional[float]
|
|
|
|
|
|
class PagePreprocessingModel(BasePageModel):
|
|
def __init__(self, options: PagePreprocessingOptions):
|
|
self.options = options
|
|
|
|
def __call__(
|
|
self, conv_res: ConversionResult, page_batch: Iterable[Page]
|
|
) -> Iterable[Page]:
|
|
for page in page_batch:
|
|
assert page._backend is not None
|
|
if not page._backend.is_valid():
|
|
yield page
|
|
else:
|
|
with TimeRecorder(conv_res, "page_parse"):
|
|
page = self._populate_page_images(page)
|
|
page = self._parse_page_cells(conv_res, page)
|
|
yield page
|
|
|
|
# Generate the page image and store it in the page object
|
|
def _populate_page_images(self, page: Page) -> Page:
|
|
# default scale
|
|
page.get_image(
|
|
scale=1.0
|
|
) # puts the page image on the image cache at default scale
|
|
|
|
images_scale = self.options.images_scale
|
|
# user requested scales
|
|
if images_scale is not None:
|
|
page._default_image_scale = images_scale
|
|
page.get_image(
|
|
scale=images_scale
|
|
) # this will trigger storing the image in the internal cache
|
|
|
|
return page
|
|
|
|
# Extract and populate the page cells and store it in the page object
|
|
def _parse_page_cells(self, conv_res: ConversionResult, page: Page) -> Page:
|
|
assert page._backend is not None
|
|
|
|
page.cells = list(page._backend.get_text_cells())
|
|
|
|
# DEBUG code:
|
|
def draw_text_boxes(image, cells, show: bool = False):
|
|
draw = ImageDraw.Draw(image)
|
|
for c in cells:
|
|
x0, y0, x1, y1 = c.bbox.as_tuple()
|
|
draw.rectangle([(x0, y0), (x1, y1)], outline="red")
|
|
if show:
|
|
image.show()
|
|
else:
|
|
out_path: Path = (
|
|
Path(settings.debug.debug_output_path)
|
|
/ f"debug_{conv_res.input.file.stem}"
|
|
)
|
|
out_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
out_file = out_path / f"cells_page_{page.page_no:05}.png"
|
|
image.save(str(out_file), format="png")
|
|
|
|
if settings.debug.visualize_cells:
|
|
draw_text_boxes(page.get_image(scale=1.0), page.cells)
|
|
|
|
return page
|