优化文档检测

docker-compose.yml

@@ -1,6 +1,6 @@
 x-env:
   &template
-  image: fcb_photo_review:1.13.2
+  image: fcb_photo_review:1.13.4
   restart: always
 
 services:

paddle_detection/__init__.py

@@ -0,0 +1,4 @@
+from onnxruntime import InferenceSession
+
+PADDLE_DET = InferenceSession("model/object_det_model/ppyoloe_plus_crn_l_80e_coco_w_nms.onnx",
+                              providers=["CUDAExecutionProvider"], provider_options=[{"device_id": 0}])

paddle_detection/detector.py

@@ -3,8 +3,8 @@ from collections import defaultdict
 
 import cv2
 import numpy as np
-from onnxruntime import InferenceSession
 
+from paddle_detection import PADDLE_DET
 from paddle_detection.deploy.third_engine.onnx.infer import PredictConfig
 from paddle_detection.deploy.third_engine.onnx.preprocess import Compose
 from util import image_util, util
@@ -30,14 +30,11 @@ def predict_image(infer_config, predictor, img_path):
 
 
 def detect_image(img_path):
-    onnx_file = "model/object_det_model/ppyoloe_plus_crn_x_80e_coco_w_nms.onnx"
     infer_cfg = "model/object_det_model/infer_cfg.yml"
-    # load predictor
-    predictor = InferenceSession(onnx_file)
     # load infer config
     infer_config = PredictConfig(infer_cfg)
 
-    return predict_image(infer_config, predictor, img_path)
+    return predict_image(infer_config, PADDLE_DET, img_path)
 
 
 def get_book_areas(image):