From d0f3f783ce46f5f0d37ac715098f6a372c261772 Mon Sep 17 00:00:00 2001
From: liuyebo <1515783401@qq.com>
Date: Wed, 28 Aug 2024 10:53:37 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=96=87=E6=A1=A3=E6=A3=80?=
 =?UTF-8?q?=E6=B5=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 docker-compose.yml           | 2 +-
 paddle_detection/__init__.py | 4 ++++
 paddle_detection/detector.py | 7 ++-----
 3 files changed, 7 insertions(+), 6 deletions(-)

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):