From 488660dddd6c4248a569ef902cd16aa76fd1e85f Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 15:43:41 +0800 Subject: [PATCH] refactor(model): replace AtomModelSingleton with ocr_model_init for OCR model initialization - Remove usage of AtomModelSingleton for OCR model initialization - Add import of ocr_model_init from model_init module - Update OCR model initialization process to use ocr_model_init function - Remove lock for OCR processing as it's no longer needed --- magic_pdf/model/pdf_extract_kit.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 45371a35..0970154f 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -22,7 +22,7 @@ except ImportError: from magic_pdf.config.constants import * from magic_pdf.model.model_list import AtomicModel -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init from magic_pdf.model.sub_modules.model_utils import ( clean_vram, crop_img, get_res_list_from_layout_res) from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( @@ -153,9 +153,14 @@ class CustomPEKModel: device=self.device, ) # 初始化ocr - self.ocr_model = atom_model_manager.get_atom_model( - atom_model_name=AtomicModel.OCR, - ocr_show_log=show_log, + # self.ocr_model = atom_model_manager.get_atom_model( + # atom_model_name=AtomicModel.OCR, + # ocr_show_log=show_log, + # det_db_box_thresh=0.3, + # lang=self.lang + # ) + self.ocr_model = ocr_model_init( + show_log=show_log, det_db_box_thresh=0.3, lang=self.lang ) @@ -218,11 +223,11 @@ class CustomPEKModel: # OCR recognition new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) - with self._lock: - if self.apply_ocr: - ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] - else: - ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res, rec=False)[0] + # with self._lock: + if self.apply_ocr: + ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] + else: + ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res, rec=False)[0] # Integration results if ocr_res: