diff --git a/docker-compose.yml b/docker-compose.yml index 4700cd9..0147748 100644 --- a/docker-compose.yml +++ b/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: diff --git a/paddle_detection/__init__.py b/paddle_detection/__init__.py index e69de29..ba7922b 100644 --- a/paddle_detection/__init__.py +++ b/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}]) diff --git a/paddle_detection/detector.py b/paddle_detection/detector.py index 94d48b6..c49d877 100644 --- a/paddle_detection/detector.py +++ b/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):