From 95fde2086c899ec17f7262e07e0f84fb4283849a Mon Sep 17 00:00:00 2001
From: jiangxt <1579525634@qq.com>
Date: Fri, 13 Oct 2023 10:16:44 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4RKNN=E6=A8=A1=E5=9E=8B?=
=?UTF-8?q?=E5=8A=A0=E8=BD=BD=E6=96=87=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
.../model_load/model_load_rknn.py | 35 +++++++++++++++++++
1 file changed, 35 insertions(+)
create mode 100644 Bank_second_part/xbank_detect_process/model_load/model_load_rknn.py
diff --git a/Bank_second_part/xbank_detect_process/model_load/model_load_rknn.py b/Bank_second_part/xbank_detect_process/model_load/model_load_rknn.py
new file mode 100644
index 0000000..113b287
--- /dev/null
+++ b/Bank_second_part/xbank_detect_process/model_load/model_load_rknn.py
@@ -0,0 +1,35 @@
+import fastdeploy as fd
+import os.path as op
+
+
+def Load_model(model_file):
+ """
+ 加载rknn模型文件,根据使用模型不同,设置后处理的class_num
+ model_file:rknn转化模型存放的地址
+ """
+
+ model_name = op.basename(model_file).split('_')[1].split('.')[0]
+
+ class_num = 0
+
+ if model_name == "phone":
+ class_num = 5
+ elif model_name == "sleep":
+ class_num = 2
+ elif model_name == "person":
+ class_num = 80
+ else:
+ print("model load failed !")
+
+ runtime_option = fd.RuntimeOption()
+ runtime_option.use_rknpu2()
+
+ model_inference = fd.vision.detection.RKYOLOV5(
+ model_file,
+ runtime_option=runtime_option,
+ model_format=fd.ModelFormat.RKNN)
+
+ model_inference.postprocessor.class_num = class_num
+
+ return model_inference
+
\ No newline at end of file