提交RKNN模型加载文件

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
+