From 0348cfb9644758fd0e005de1dc2e23c4b9de1140 Mon Sep 17 00:00:00 2001 From: Swaymaw Date: Wed, 27 Nov 2024 10:26:43 +0530 Subject: [PATCH] simplifying rapidocr options so that device can be changed using a single option for all models Signed-off-by: Swaymaw --- docling/datamodel/pipeline_options.py | 17 ++++------- docling/models/rapid_ocr_model.py | 41 +++++++++++++++++++++++---- 2 files changed, 41 insertions(+), 17 deletions(-) diff --git a/docling/datamodel/pipeline_options.py b/docling/datamodel/pipeline_options.py index 50fb4cab..56c743a3 100644 --- a/docling/datamodel/pipeline_options.py +++ b/docling/datamodel/pipeline_options.py @@ -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 diff --git a/docling/models/rapid_ocr_model.py b/docling/models/rapid_ocr_model.py index 1e12a106..7686f784 100644 --- a/docling/models/rapid_ocr_model.py +++ b/docling/models/rapid_ocr_model.py @@ -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,