Commit a0eff3be authored by liukaiwen's avatar liukaiwen

feat: table model update with paddle recognition v4

parent 51f56aa3
......@@ -37,4 +37,10 @@ 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"
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
......@@ -271,6 +272,17 @@ class CustomPEKModel:
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):
......@@ -433,7 +445,5 @@ class CustomPEKModel:
logger.info(f"-----page total time: {round(time.time() - page_start, 2)}-----")
return layout_res
if __name__ == '__main__':
print()
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