mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-26 20:14:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add type hints and fix mypy

Browse Source 
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Christoph Auer 2025-06-13 11:08:55 +02:00
parent 9469280802
commit e3cefd0e71
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
docling/models/base_ocr_model.py
View File

@ -7,6 +7,7 @@ from typing import List, Optional, Type
import numpy as np import numpy as np
from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc import BoundingBox, CoordOrigin
from docling_core.types.doc.page import TextCell
from PIL import Image, ImageDraw from PIL import Image, ImageDraw
from rtree import index from rtree import index
from scipy.ndimage import binary_dilation, find_objects, label from scipy.ndimage import binary_dilation, find_objects, label
@ -107,7 +108,9 @@ class BaseOcrModel(BasePageModel, BaseModelWithOptions):
return [] return []
# Filters OCR cells by dropping any OCR cell that intersects with an existing programmatic cell. # Filters OCR cells by dropping any OCR cell that intersects with an existing programmatic cell.
def _filter_ocr_cells(self, ocr_cells, programmatic_cells): def _filter_ocr_cells(
self, ocr_cells: List[TextCell], programmatic_cells: List[TextCell]
) -> List[TextCell]:
# Create R-tree index for programmatic cells # Create R-tree index for programmatic cells
p = index.Property() p = index.Property()
p.dimension = 2 p.dimension = 2
@ -130,7 +133,7 @@ class BaseOcrModel(BasePageModel, BaseModelWithOptions):
] ]
return filtered_ocr_cells return filtered_ocr_cells
def post_process_cells(self, ocr_cells, page): def post_process_cells(self, ocr_cells: List[TextCell], page: Page) -> None:
r""" r"""
Post-process the OCR cells and update the page object. Post-process the OCR cells and update the page object.
Updates parsed_page.textline_cells directly since page.cells is now read-only. Updates parsed_page.textline_cells directly since page.cells is now read-only.
@ -147,7 +150,9 @@ class BaseOcrModel(BasePageModel, BaseModelWithOptions):
page.parsed_page.textline_cells = final_cells page.parsed_page.textline_cells = final_cells
page.parsed_page.has_lines = len(final_cells) > 0 page.parsed_page.has_lines = len(final_cells) > 0
def _combine_cells(self, existing_cells, ocr_cells): def _combine_cells(
self, existing_cells: List[TextCell], ocr_cells: List[TextCell]
) -> List[TextCell]:
"""Combine existing and OCR cells with filtering and re-indexing.""" """Combine existing and OCR cells with filtering and re-indexing."""
if self.options.force_full_page_ocr: if self.options.force_full_page_ocr:
combined = ocr_cells combined = ocr_cells