opensource/MinerU
3
0
Fork 0
mirror of https://github.com/opendatalab/MinerU.git synced 2026-09-24 23:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(model): replace AtomModelSingleton with ocr_model_init for OCR model initialization

Browse Source 
- 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
This commit is contained in:
myhloli
2024-12-06 15:43:41 +08:00
parent 6f636b6e7e
commit 488660dddd
1 changed files with 14 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
magic_pdf/model/pdf_extract_kit.py
+14 -9
View File
@@ -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: