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

simplifying rapidocr options so that device can be changed using a single option for all models

Browse Source 
Signed-off-by: Swaymaw <swaymaw@gmail.com>
This commit is contained in:
Swaymaw 2024-11-27 10:26:43 +05:30
parent 686affe782
commit 0348cfb964
2 changed files with 41 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
docling/datamodel/pipeline_options.py
View File

@ -41,18 +41,13 @@ class RapidOcrOptions(OcrOptions):
# 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
class Device(Enum):
CPU = "CPU"
CUDA = "CUDA"
DIRECTML = "DIRECTML"
AUTO = "AUTO"
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
device: Device = Device.AUTO # Default value is AUTO
print_verbose: bool = False # same default as rapidocr

41
docling/models/rapid_ocr_model.py
View File

@ -29,15 +29,44 @@ class RapidOcrModel(BaseOcrModel):
"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."
)
# Same as Defaults in RapidOCR
cls_use_cuda = False
rec_use_cuda = False
det_use_cuda = False
det_use_dml = False
cls_use_dml = False
rec_use_dml = False
# If we set everything to true onnx-runtime would automatically choose the fastest accelerator
if self.options.device == self.options.Device.AUTO:
cls_use_cuda = True
rec_use_cuda = True
det_use_cuda = True
det_use_dml = True
cls_use_dml = True
rec_use_dml = True
# If we set use_cuda to true onnx would use the cuda device available in runtime if no cuda device is available it would run on CPU.
elif self.options.device == self.options.Device.CUDA:
cls_use_cuda = True
rec_use_cuda = True
det_use_cuda = True
# If we set use_dml to true onnx would use the dml device available in runtime if no dml device is available it would work on CPU.
elif self.options.device == self.options.Device.DIRECTML:
det_use_dml = True
cls_use_dml = True
rec_use_dml = True
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,
cls_use_cuda=cls_use_cuda,
rec_use_cuda=rec_use_cuda,
det_use_cuda=det_use_cuda,
det_use_dml=det_use_dml,
cls_use_dml=cls_use_dml,
rec_use_dml=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,