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

adding rapidocr engine for ocr in docling

Browse Source 
Signed-off-by: swayam-singhal <swayam.singhal@inito.com>
This commit is contained in:
swayam-singhal 2024-11-22 12:45:06 +05:30
parent 2a1d3fd221
commit 9bb2e58e59
9 changed files with 767 additions and 630 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docling/cli/main.py
View File

@ -30,6 +30,7 @@ from docling.datamodel.pipeline_options import (
TableFormerMode,
TesseractCliOcrOptions,
TesseractOcrOptions,
RapidOcrOptions
)
from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption
@ -76,6 +77,7 @@ class OcrEngine(str, Enum):
TESSERACT_CLI = "tesseract_cli"
TESSERACT = "tesseract"
OCRMAC = "ocrmac"
RAPIDOCR = "rapidocr"
def export_documents(
@ -262,6 +264,8 @@ def convert(
ocr_options = TesseractOcrOptions(force_full_page_ocr=force_ocr)
elif ocr_engine == OcrEngine.OCRMAC:
ocr_options = OcrMacOptions(force_full_page_ocr=force_ocr)
elif ocr_engine == OcrEngine.RAPIDOCR:
ocr_options = RapidOcrOptions(force_full_page_ocr=force_ocr)
else:
raise RuntimeError(f"Unexpected OCR engine type {ocr_engine}")

31
docling/datamodel/pipeline_options.py
View File

@ -28,6 +28,37 @@ class OcrOptions(BaseModel):
0.05 # percentage of the area for a bitmap to processed with OCR
)
class RapidOcrOptions(OcrOptions):
kind: Literal["rapidocr"] = "rapidocr"
# English and chinese are the most commly used models and have been tested with RapidOCR.
lang: List[str] = ["english", "chinese"] # However, language as a parameter is not supported by rapidocr yet and hence changing this options doesn't affect anything.
# For more details on supported languages by RapidOCR visit https://rapidai.github.io/RapidOCRDocs/blog/2022/09/28/%E6%94%AF%E6%8C%81%E8%AF%86%E5%88%AB%E8%AF%AD%E8%A8%80/
# For more details on the following options visit https://rapidai.github.io/RapidOCRDocs/install_usage/api/RapidOCR/
text_score: float = 0.5 # same default as rapidocr
use_det: Optional[bool] = None # same default as rapidocr
use_cls: Optional[bool] = None # same default as rapidocr
use_rec: Optional[bool] = None # same default as rapidocr
det_use_cuda: bool = False # same default as rapidocr
cls_use_cuda: bool = False # same default as rapidocr
rec_use_cuda: bool = False # same default as rapidocr
det_use_dml: bool = False # same default as rapidocr
cls_use_dml: bool = False # same default as rapidocr
rec_use_dml: bool = False # same default as rapidocr
print_verbose: bool = False # same default as rapidocr
det_model_path: Optional[str] = None # same default as rapidocr
cls_model_path: Optional[str] = None # same default as rapidocr
rec_model_path: Optional[str] = None # same default as rapidocr
model_config = ConfigDict(
extra="forbid",
)
class EasyOcrOptions(OcrOptions):
kind: Literal["easyocr"] = "easyocr"

104
docling/models/rapid_ocr_model.py Normal file
View File

@ -0,0 +1,104 @@
import logging
from typing import Iterable
import numpy
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 RapidOcrOptions
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 RapidOcrModel(BaseOcrModel):
def __init__(self, enabled: bool, options: RapidOcrOptions):
super().__init__(enabled=enabled, options=options)
self.options: RapidOcrOptions
self.scale = 3 # multiplier for 72 dpi == 216 dpi.
if self.enabled:
try:
from rapidocr_onnxruntime import RapidOCR
except ImportError:
raise ImportError(
"RapidOCR is not installed. Please install it via `pip install rapidocr_onnxruntime` to use this OCR engine. "
"Alternatively, Docling has support for other OCR engines. See the documentation."
)
self.reader = RapidOCR(
text_score = self.options.text_score,
cls_use_cuda = self.options.cls_use_cuda,
rec_use_cuda = self.options.rec_use_cuda,
det_use_cuda = self.options.det_use_cuda,
det_use_dml = self.options.det_use_dml,
cls_use_dml = self.options.cls_use_dml,
rec_use_dml = self.options.rec_use_dml,
print_verbose = self.options.print_verbose,
det_model_path = self.options.det_model_path,
cls_model_path = self.options.cls_model_path,
rec_model_path = self.options.rec_model_path,
)
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:
# Skip zero area boxes
if ocr_rect.area() == 0:
continue
high_res_image = page._backend.get_page_image(
scale=self.scale, cropbox=ocr_rect
)
im = numpy.array(high_res_image)
result, _ = self.reader(im, use_det=self.options.use_det, use_cls=self.options.use_cls, use_rec=self.options.use_rec)
del high_res_image
del im
cells = [
OcrCell(
id=ix,
text=line[1],
confidence=line[2],
bbox=BoundingBox.from_tuple(
coord=(
(line[0][0][0] / self.scale) + ocr_rect.l,
(line[0][0][1] / self.scale) + ocr_rect.t,
(line[0][2][0] / self.scale) + ocr_rect.l,
(line[0][2][1] / self.scale) + ocr_rect.t,
),
origin=CoordOrigin.TOPLEFT,
),
)
for ix, line in enumerate(result)
]
all_ocr_cells.extend(cells)
# Post-process the cells
page.cells = self.post_process_cells(all_ocr_cells, page.cells)
# DEBUG code:
if settings.debug.visualize_ocr:
self.draw_ocr_rects_and_cells(conv_res, page, ocr_rects)
yield page

7
docling/pipeline/standard_pdf_pipeline.py
View File

@ -15,6 +15,7 @@ from docling.datamodel.pipeline_options import (
PdfPipelineOptions,
TesseractCliOcrOptions,
TesseractOcrOptions,
RapidOcrOptions
)
from docling.models.base_ocr_model import BaseOcrModel
from docling.models.ds_glm_model import GlmModel, GlmOptions
@ -29,6 +30,7 @@ from docling.models.page_preprocessing_model import (
from docling.models.table_structure_model import TableStructureModel
from docling.models.tesseract_ocr_cli_model import TesseractOcrCliModel
from docling.models.tesseract_ocr_model import TesseractOcrModel
from docling.models.rapid_ocr_model import RapidOcrModel
from docling.pipeline.base_pipeline import PaginatedPipeline
from docling.utils.profiling import ProfilingScope, TimeRecorder
@ -121,6 +123,11 @@ class StandardPdfPipeline(PaginatedPipeline):
enabled=self.pipeline_options.do_ocr,
options=self.pipeline_options.ocr_options,
)
elif isinstance(self.pipeline_options.ocr_options, RapidOcrOptions):
return RapidOcrModel(
enabled=self.pipeline_options.do_ocr,
options=self.pipeline_options.ocr_options,
)
elif isinstance(self.pipeline_options.ocr_options, OcrMacOptions):
if "darwin" != sys.platform:
raise RuntimeError(

4
docs/examples/full_page_ocr.py
View File

@ -8,6 +8,7 @@ from docling.datamodel.pipeline_options import (
PdfPipelineOptions,
TesseractCliOcrOptions,
TesseractOcrOptions,
RapidOcrOptions
)
from docling.document_converter import DocumentConverter, PdfFormatOption
@ -20,10 +21,11 @@ def main():
pipeline_options.do_table_structure = True
pipeline_options.table_structure_options.do_cell_matching = True
# Any of the OCR options can be used:EasyOcrOptions, TesseractOcrOptions, TesseractCliOcrOptions, OcrMacOptions(Mac only)
# Any of the OCR options can be used:EasyOcrOptions, TesseractOcrOptions, TesseractCliOcrOptions, OcrMacOptions(Mac only), RapidOcrOptions
# ocr_options = EasyOcrOptions(force_full_page_ocr=True)
# ocr_options = TesseractOcrOptions(force_full_page_ocr=True)
# ocr_options = OcrMacOptions(force_full_page_ocr=True)
# ocr_options = RapidOcrOptions(force_full_page_ocr=True)
ocr_options = TesseractCliOcrOptions(force_full_page_ocr=True)
pipeline_options.ocr_options = ocr_options

1
docs/installation.md
View File

@ -31,6 +31,7 @@ Works on macOS, Linux, and Windows, with support for both x86_64 and arm64 archi
| Tesseract | System dependency. See description for Tesseract and Tesserocr below. | `TesseractOcrOptions` |
| Tesseract CLI | System dependency. See description below. | `TesseractCliOcrOptions` |
| OcrMac | System dependency. See description below. | `OcrMacOptions` |
| [RapidOCR](https://github.com/RapidAI/RapidOCR) | Extra feature not included in Default Docling installation can be installed via `pip install rapidocr_onnxruntime` | `RapidOcrOptions` |
The Docling `DocumentConverter` allows to choose the OCR engine with the `ocr_options` settings. For example

1241
poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

2
pyproject.toml
View File

@ -49,6 +49,7 @@ pandas = "^2.1.4"
marko = "^2.1.2"
openpyxl = "^3.1.5"
ocrmac = { version = "^1.0.0", markers = "sys_platform == 'darwin'", optional = true }
rapidocr-onnxruntime = { version = "^1.4.0", optional = true, markers = "python_version < '3.13'" }
[tool.poetry.group.dev.dependencies]
black = {extras = ["jupyter"], version = "^24.4.2"}
@ -103,6 +104,7 @@ torchvision = [
[tool.poetry.extras]
tesserocr = ["tesserocr"]
ocrmac = ["ocrmac"]
rapidocr = ["rapidocr-onnxruntime"]
[tool.poetry.scripts]
docling = "docling.cli.main:app"

3
tests/test_e2e_ocr_conversion.py
View File

@ -12,6 +12,7 @@ from docling.datamodel.pipeline_options import (
PdfPipelineOptions,
TesseractCliOcrOptions,
TesseractOcrOptions,
RapidOcrOptions
)
from docling.document_converter import DocumentConverter, PdfFormatOption
@ -56,9 +57,11 @@ def test_e2e_conversions():
EasyOcrOptions(),
TesseractOcrOptions(),
TesseractCliOcrOptions(),
RapidOcrOptions(),
EasyOcrOptions(force_full_page_ocr=True),
TesseractOcrOptions(force_full_page_ocr=True),
TesseractCliOcrOptions(force_full_page_ocr=True),
RapidOcrOptions(force_full_page_ocr=True)
]
# only works on mac