From 74e005df63dc544ede755ec8fd6ca885960035fd Mon Sep 17 00:00:00 2001 From: Swaymaw Date: Wed, 27 Nov 2024 10:52:25 +0530 Subject: [PATCH] fix styling issues and small bug in rapidOcrOptions Signed-off-by: Swaymaw --- docling/datamodel/pipeline_options.py | 7 +++++- docling/models/rapid_ocr_model.py | 34 +++++++++++++-------------- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/docling/datamodel/pipeline_options.py b/docling/datamodel/pipeline_options.py index 56c743a3..0a5f75a0 100644 --- a/docling/datamodel/pipeline_options.py +++ b/docling/datamodel/pipeline_options.py @@ -41,7 +41,12 @@ 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): + + 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 + + class Device(Enum): CPU = "CPU" CUDA = "CUDA" DIRECTML = "DIRECTML" diff --git a/docling/models/rapid_ocr_model.py b/docling/models/rapid_ocr_model.py index 7686f784..f36291ce 100644 --- a/docling/models/rapid_ocr_model.py +++ b/docling/models/rapid_ocr_model.py @@ -29,36 +29,36 @@ 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 + 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 + 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: + # 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 - + rec_use_dml = True + self.reader = RapidOCR( text_score=self.options.text_score, cls_use_cuda=cls_use_cuda,