Commit bd1b7677 authored by myhloli's avatar myhloli
Browse files

refactor(magic_pdf): update OCR engine selection in RapidTableModel

- Remove conditional logic for OCR engine selection
- Always use RapidOCR as the OCR engine
- Simplify the __init__ method by removing unused code
parent f5f35633
...@@ -8,7 +8,18 @@ from rapid_table import RapidTable ...@@ -8,7 +8,18 @@ from rapid_table import RapidTable
class RapidTableModel(object): class RapidTableModel(object):
def __init__(self, ocr_engine): def __init__(self, ocr_engine):
self.table_model = RapidTable() self.table_model = RapidTable()
if ocr_engine is None: # if ocr_engine is None:
# self.ocr_model_name = "RapidOCR"
# if torch.cuda.is_available():
# from rapidocr_paddle import RapidOCR
# self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
# else:
# from rapidocr_onnxruntime import RapidOCR
# self.ocr_engine = RapidOCR()
# else:
# self.ocr_model_name = "PaddleOCR"
# self.ocr_engine = ocr_engine
self.ocr_model_name = "RapidOCR" self.ocr_model_name = "RapidOCR"
if torch.cuda.is_available(): if torch.cuda.is_available():
from rapidocr_paddle import RapidOCR from rapidocr_paddle import RapidOCR
...@@ -16,9 +27,6 @@ class RapidTableModel(object): ...@@ -16,9 +27,6 @@ class RapidTableModel(object):
else: else:
from rapidocr_onnxruntime import RapidOCR from rapidocr_onnxruntime import RapidOCR
self.ocr_engine = RapidOCR() self.ocr_engine = RapidOCR()
else:
self.ocr_model_name = "PaddleOCR"
self.ocr_engine = ocr_engine
def predict(self, image): def predict(self, image):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment