pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reuse loaded modules if possible (#5231)

Browse Source 
### What problem does this PR solve?

Reuse loaded modules if possible

### Type of change

- [x] Refactoring
This commit is contained in:
Zhichang Yu
2025-02-21 17:21:01 +08:00
committed by GitHub
parent 392f28882f
commit 0151d42156
2 changed files with 18 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
deepdoc/vision/ocr.py
View File

@ -31,6 +31,7 @@ import onnxruntime as ort
from .postprocess import build_post_process
loaded_models = {}
def transform(data, ops=None):
""" transform """
@ -67,6 +68,12 @@ def create_operators(op_param_list, global_config=None):
def load_model(model_dir, nm):
model_file_path = os.path.join(model_dir, nm + ".onnx")
global loaded_models
loaded_model = loaded_models.get(model_file_path)
if loaded_model:
logging.info(f"load_model {model_file_path} reuses cached model")
return loaded_model
if not os.path.exists(model_file_path):
raise ValueError("not find model file path {}".format(
model_file_path))
@ -102,15 +109,17 @@ def load_model(model_dir, nm):
provider_options=[cuda_provider_options]
)
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "gpu:0")
logging.info(f"TextRecognizer {nm} uses GPU")
logging.info(f"load_model {model_file_path} uses GPU")
else:
sess = ort.InferenceSession(
model_file_path,
options=options,
providers=['CPUExecutionProvider'])
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu")
logging.info(f"TextRecognizer {nm} uses CPU")
return sess, sess.get_inputs()[0], run_options
logging.info(f"load_model {model_file_path} uses CPU")
loaded_model = (sess, run_options)
loaded_models[model_file_path] = loaded_model
return loaded_model
class TextRecognizer(object):
@ -123,7 +132,8 @@ class TextRecognizer(object):
"use_space_char": True
}
self.postprocess_op = build_post_process(postprocess_params)
self.predictor, self.input_tensor, self.run_options = load_model(model_dir, 'rec')
self.predictor, self.run_options = load_model(model_dir, 'rec')
self.input_tensor = self.predictor.get_inputs()[0]
def resize_norm_img(self, img, max_wh_ratio):
imgC, imgH, imgW = self.rec_image_shape
@ -408,7 +418,8 @@ class TextDetector(object):
"unclip_ratio": 1.5, "use_dilation": False, "score_mode": "fast", "box_type": "quad"}
self.postprocess_op = build_post_process(postprocess_params)
self.predictor, self.input_tensor, self.run_options = load_model(model_dir, 'det')
self.predictor, self.run_options = load_model(model_dir, 'det')
self.input_tensor = self.predictor.get_inputs()[0]
img_h, img_w = self.input_tensor.shape[2:]
if isinstance(img_h, str) or isinstance(img_w, str):