diff --git a/src/main/java/org/photonvision/rknn/RknnJNI.java b/src/main/java/org/photonvision/rknn/RknnJNI.java index 4163074..856157d 100644 --- a/src/main/java/org/photonvision/rknn/RknnJNI.java +++ b/src/main/java/org/photonvision/rknn/RknnJNI.java @@ -84,6 +84,15 @@ public boolean equals(Object obj) { * @return Pointer to the detector in native memory */ public static native long create(String modelPath, int numClasses, int modelVer, int coreNum); + + /** + * Given an already running detector, change the bitmask controlling which + * of the 3 cores the model is running on + * @param ptr Pointer to detector in native memory + * @param desiredCore Which of the three cores to operate on + * @return return code of rknn_set_core_mask call, indicating success or failure + */ + public static native int setCoreMask(long ptr, int desiredCore); /** * Delete all native resources assocated with a detector diff --git a/src/main/native/cpp/rknn_jni.cpp b/src/main/native/cpp/rknn_jni.cpp index c06b43a..da05e78 100644 --- a/src/main/native/cpp/rknn_jni.cpp +++ b/src/main/native/cpp/rknn_jni.cpp @@ -77,6 +77,15 @@ Java_org_photonvision_rknn_RknnJNI_create return reinterpret_cast(ret); } +JNIEXPORT jint JNICALL Java_org_photonvision_rknn_RknnJNI_setCoreMask(JNIEnv *env, + jclass, + jlong ptr, + jint coreMask) +{ + YoloModel *yolo = reinterpret_cast(ptr); + return yolo->changeCoreMask(coreMask); +} + /* * Class: org_photonvision_rknn_RknnJNI * Method: destroy diff --git a/src/main/native/cpp/rknn_jni.h b/src/main/native/cpp/rknn_jni.h index b80f876..d3b6808 100644 --- a/src/main/native/cpp/rknn_jni.h +++ b/src/main/native/cpp/rknn_jni.h @@ -35,6 +35,12 @@ JNIEXPORT jlong JNICALL Java_org_photonvision_rknn_RknnJNI_create(JNIEnv *, jstring, jint, jint, jint); + +JNIEXPORT jint JNICALL Java_org_photonvision_rknn_RknnJNI_setCoreMask(JNIEnv *, + jclass, + jlong, + jint); + /* * Class: org_photonvision_rknn_RknnJNI * Method: destroy diff --git a/src/main/native/cpp/yolo_common.cpp b/src/main/native/cpp/yolo_common.cpp index 2b55e31..82ce261 100644 --- a/src/main/native/cpp/yolo_common.cpp +++ b/src/main/native/cpp/yolo_common.cpp @@ -82,37 +82,8 @@ YoloModel::YoloModel(std::string modelPath, int num_classes_, ModelVersion type_ throw std::runtime_error("rknn_init error ret=" + ret); } - // 设置模型绑定的核心/Set the core of the model that needs to be bound - rknn_core_mask core_mask; - switch (coreNumber) - { - case -1: - core_mask = RKNN_NPU_CORE_AUTO; - break; - case 0: - core_mask = RKNN_NPU_CORE_0; - break; - case 1: - core_mask = RKNN_NPU_CORE_1; - break; - case 2: - core_mask = RKNN_NPU_CORE_2; - break; - case 10: - core_mask = RKNN_NPU_CORE_0_1; - break; - case 210: - core_mask = RKNN_NPU_CORE_0_1_2; - break; - default: - throw std::runtime_error("invalid core selection! core selected: " + coreNumber); - break; - } - ret = rknn_set_core_mask(ctx, core_mask); - if (ret < 0) - { - throw std::runtime_error("rknn_init core error ret=" + ret); - } + // hard coded to let npu decide where the model runs + this->changeCoreMask(-1); rknn_sdk_version version; ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version)); @@ -204,6 +175,40 @@ YoloModel::~YoloModel() { free(model_data); } +int YoloModel::changeCoreMask(int coreNumber) { + // 设置模型绑定的核心/Set the core of the model that needs to be bound + rknn_core_mask core_mask; + switch (coreNumber) + { + case -1: + core_mask = RKNN_NPU_CORE_AUTO; + break; + case 0: + core_mask = RKNN_NPU_CORE_0; + break; + case 1: + core_mask = RKNN_NPU_CORE_1; + break; + case 2: + core_mask = RKNN_NPU_CORE_2; + break; + case 10: + core_mask = RKNN_NPU_CORE_0_1; + break; + case 210: + core_mask = RKNN_NPU_CORE_0_1_2; + break; + default: + throw std::runtime_error("invalid core selection! core selected: " + coreNumber); + break; + } + ret = rknn_set_core_mask(ctx, core_mask); + if (ret < 0) + { + throw std::runtime_error("rknn_init core error ret=" + ret); + } + return ret; +} detect_result_group_t YoloModel::forward(cv::Mat &orig_img, DetectionFilterParams params) { cv::Mat img; diff --git a/src/main/native/include/yolo_common.hpp b/src/main/native/include/yolo_common.hpp index 069376d..8810106 100644 --- a/src/main/native/include/yolo_common.hpp +++ b/src/main/native/include/yolo_common.hpp @@ -19,6 +19,8 @@ class YoloModel { public: YoloModel(std::string modelPath, int num_classes_, ModelVersion type_, int coreNumber); + int YoloModel::changeCoreMask(int coreNumber); + detect_result_group_t forward(cv::Mat &orig_img, DetectionFilterParams params); ~YoloModel(); diff --git a/src/test/java/org/photonvision/rknn/RknnTest.java b/src/test/java/org/photonvision/rknn/RknnTest.java index d944bf0..a731073 100644 --- a/src/test/java/org/photonvision/rknn/RknnTest.java +++ b/src/test/java/org/photonvision/rknn/RknnTest.java @@ -55,18 +55,26 @@ public void testBasicBlobs() { System.out.println("Loading bus"); - Mat img = Imgcodecs.imread("silly_notes2.png"); + Mat img = Imgcodecs.imread("silly_notes.png"); + Mat img2 = Imgcodecs.imread("silly_notes2.png"); System.out.println("Loading rknn-jni"); System.load("/home/coolpi/rknn_jni/cmake_build/librknn_jni.so"); - System.out.println("Creating detector"); - long ptr = RknnJNI.create("/home/coolpi/rknn_jni/note-640-640-yolov5s.rknn", 1, ModelVersion.YOLO_V5.ordinal(), 0); + System.out.println("Creating detector on three cores"); + long ptr = RknnJNI.create("/home/coolpi/rknn_jni/note-640-640-yolov5s.rknn", 1, ModelVersion.YOLO_V5.ordinal(), 210); System.out.println("Running detector"); var ret = RknnJNI.detect(ptr, img.getNativeObjAddr(), .45, .25); System.out.println(Arrays.toString(ret)); + System.out.println("Changing detector to run on core 0"); + System.out.println("return code: " + RknnJNI.setCoreMask(ptr, 0)); + + System.out.println("Running detector again"); + ret = RknnJNI.detect(ptr, img2.getNativeObjAddr(), .45, .25); + System.out.println(Arrays.toString(ret)); + System.out.println("Killing detector"); RknnJNI.destroy(ptr); }