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

fix styling issues and small bug in rapidOcrOptions

Browse Source 
Signed-off-by: Swaymaw <swaymaw@gmail.com>
This commit is contained in:
Swaymaw 2024-11-27 10:52:25 +05:30
parent 0348cfb964
commit 74e005df63
2 changed files with 23 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
docling/datamodel/pipeline_options.py
View File

@ -41,7 +41,12 @@ class RapidOcrOptions(OcrOptions):
# For more details on the following options visit https://rapidai.github.io/RapidOCRDocs/install_usage/api/RapidOCR/ # 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 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" CPU = "CPU"
CUDA = "CUDA" CUDA = "CUDA"
DIRECTML = "DIRECTML" DIRECTML = "DIRECTML"

34
docling/models/rapid_ocr_model.py
View File

@ -29,36 +29,36 @@ class RapidOcrModel(BaseOcrModel):
"RapidOCR is not installed. Please install it via `pip install rapidocr_onnxruntime` to use this OCR engine. " "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." "Alternatively, Docling has support for other OCR engines. See the documentation."
) )
# Same as Defaults in RapidOCR # Same as Defaults in RapidOCR
cls_use_cuda = False cls_use_cuda = False
rec_use_cuda = False rec_use_cuda = False
det_use_cuda = False det_use_cuda = False
det_use_dml = False det_use_dml = False
cls_use_dml = False cls_use_dml = False
rec_use_dml = False rec_use_dml = False
# If we set everything to true onnx-runtime would automatically choose the fastest accelerator # If we set everything to true onnx-runtime would automatically choose the fastest accelerator
if self.options.device == self.options.Device.AUTO: if self.options.device == self.options.Device.AUTO:
cls_use_cuda = True cls_use_cuda = True
rec_use_cuda = True rec_use_cuda = True
det_use_cuda = True det_use_cuda = True
det_use_dml = True det_use_dml = True
cls_use_dml = True cls_use_dml = True
rec_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. # 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.DIRECTML: 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 det_use_dml = True
cls_use_dml = True cls_use_dml = True
rec_use_dml = True rec_use_dml = True
self.reader = RapidOCR( self.reader = RapidOCR(
text_score=self.options.text_score, text_score=self.options.text_score,
cls_use_cuda=cls_use_cuda, cls_use_cuda=cls_use_cuda,