| | import logging |
| | import tempfile |
| | from typing import Iterable, Optional, Tuple |
| |
|
| | from docling_core.types.doc import BoundingBox, CoordOrigin |
| |
|
| | from docling.datamodel.base_models import OcrCell, Page |
| | from docling.datamodel.document import ConversionResult |
| | from docling.datamodel.pipeline_options import OcrMacOptions |
| | from docling.datamodel.settings import settings |
| | from docling.models.base_ocr_model import BaseOcrModel |
| | from docling.utils.profiling import TimeRecorder |
| |
|
| | _log = logging.getLogger(__name__) |
| |
|
| |
|
| | class OcrMacModel(BaseOcrModel): |
| | def __init__(self, enabled: bool, options: OcrMacOptions): |
| | super().__init__(enabled=enabled, options=options) |
| | self.options: OcrMacOptions |
| |
|
| | self.scale = 3 |
| |
|
| | if self.enabled: |
| | install_errmsg = ( |
| | "ocrmac is not correctly installed. " |
| | "Please install it via `pip install ocrmac` to use this OCR engine. " |
| | "Alternatively, Docling has support for other OCR engines. See the documentation: " |
| | "https://ds4sd.github.io/docling/installation/" |
| | ) |
| | try: |
| | from ocrmac import ocrmac |
| | except ImportError: |
| | raise ImportError(install_errmsg) |
| |
|
| | self.reader_RIL = ocrmac.OCR |
| |
|
| | def __call__( |
| | self, conv_res: ConversionResult, page_batch: Iterable[Page] |
| | ) -> Iterable[Page]: |
| |
|
| | if not self.enabled: |
| | yield from page_batch |
| | return |
| |
|
| | for page in page_batch: |
| | assert page._backend is not None |
| | if not page._backend.is_valid(): |
| | yield page |
| | else: |
| | with TimeRecorder(conv_res, "ocr"): |
| |
|
| | ocr_rects = self.get_ocr_rects(page) |
| |
|
| | all_ocr_cells = [] |
| | for ocr_rect in ocr_rects: |
| | |
| | if ocr_rect.area() == 0: |
| | continue |
| | high_res_image = page._backend.get_page_image( |
| | scale=self.scale, cropbox=ocr_rect |
| | ) |
| |
|
| | with tempfile.NamedTemporaryFile( |
| | suffix=".png", mode="w" |
| | ) as image_file: |
| | fname = image_file.name |
| | high_res_image.save(fname) |
| |
|
| | boxes = self.reader_RIL( |
| | fname, |
| | recognition_level=self.options.recognition, |
| | framework=self.options.framework, |
| | language_preference=self.options.lang, |
| | ).recognize() |
| |
|
| | im_width, im_height = high_res_image.size |
| | cells = [] |
| | for ix, (text, confidence, box) in enumerate(boxes): |
| | x = float(box[0]) |
| | y = float(box[1]) |
| | w = float(box[2]) |
| | h = float(box[3]) |
| |
|
| | x1 = x * im_width |
| | y2 = (1 - y) * im_height |
| |
|
| | x2 = x1 + w * im_width |
| | y1 = y2 - h * im_height |
| |
|
| | left = x1 / self.scale |
| | top = y1 / self.scale |
| | right = x2 / self.scale |
| | bottom = y2 / self.scale |
| |
|
| | cells.append( |
| | OcrCell( |
| | id=ix, |
| | text=text, |
| | confidence=confidence, |
| | bbox=BoundingBox.from_tuple( |
| | coord=(left, top, right, bottom), |
| | origin=CoordOrigin.TOPLEFT, |
| | ), |
| | ) |
| | ) |
| |
|
| | |
| | all_ocr_cells.extend(cells) |
| |
|
| | |
| | page.cells = self.post_process_cells(all_ocr_cells, page.cells) |
| |
|
| | |
| | if settings.debug.visualize_ocr: |
| | self.draw_ocr_rects_and_cells(conv_res, page, ocr_rects) |
| |
|
| | yield page |
| |
|
