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