diff --git a/magic_pdf/libs/Constants.py b/magic_pdf/libs/Constants.py index 2a51b2dc..0799f6fd 100644 --- a/magic_pdf/libs/Constants.py +++ b/magic_pdf/libs/Constants.py @@ -23,14 +23,20 @@ TABLE_MASTER_DIR = "table_structure_tablemaster_infer/" # pp detect model dir -DETECT_MODEL_DIR = "ch_PP-OCRv3_det_infer" +DETECT_MODEL_DIR = "ch_PP-OCRv4_det_infer" # pp rec model dir -REC_MODEL_DIR = "ch_PP-OCRv3_rec_infer" +REC_MODEL_DIR = "ch_PP-OCRv4_rec_infer" # pp rec char dict path REC_CHAR_DICT = "ppocr_keys_v1.txt" +# pp rec copy rec directory +PP_REC_DIRECTORY = ".paddleocr/whl/rec/ch/ch_PP-OCRv4_rec_infer" + +# pp rec copy det directory +PP_DET_DIRECTORY = ".paddleocr/whl/det/ch/ch_PP-OCRv4_det_infer" + class MODEL_NAME: # pp table structure algorithm diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 145b3e8f..f1478b10 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -1,7 +1,8 @@ from loguru import logger import os import time - +from pathlib import Path +import shutil from magic_pdf.libs.Constants import * from magic_pdf.libs.clean_memory import clean_memory from magic_pdf.model.model_list import AtomicModel @@ -297,6 +298,17 @@ def __init__(self, ocr: bool = False, show_log: bool = False, **kwargs): device=self.device ) + home_directory = Path.home() + det_source = os.path.join(models_dir, table_model_dir, DETECT_MODEL_DIR) + rec_source = os.path.join(models_dir, table_model_dir, REC_MODEL_DIR) + det_dest_dir = os.path.join(home_directory, PP_DET_DIRECTORY) + rec_dest_dir = os.path.join(home_directory, PP_REC_DIRECTORY) + + if not os.path.exists(det_dest_dir): + shutil.copytree(det_source, det_dest_dir) + if not os.path.exists(rec_dest_dir): + shutil.copytree(rec_source, rec_dest_dir) + logger.info('DocAnalysis init done!') def __call__(self, image): @@ -472,3 +484,5 @@ def __call__(self, image): logger.info(f"-----page total time: {round(time.time() - page_start, 2)}-----") return layout_res + +