Skip to content

Commit 666c0c6

Browse files
authored
Add function to check RKNN quantization in JNI (#26)
* add is quantized check * add non quant model to check * actually load the lib lol * fix issues hopefully? * im hoping this rebase works * resolve merge conflicts * add non quant model to check * Add linting workflow (#27) # Conflicts: # src/main/java/org/photonvision/rknn/RknnJNI.java # src/main/native/include/yolo_common.hpp # src/test/java/org/photonvision/rknn/RknnTest.java * fix member ordering
1 parent c5e111f commit 666c0c6

File tree

6 files changed

+63
-4
lines changed

6 files changed

+63
-4
lines changed
22.9 MB
Binary file not shown.

src/main/java/org/photonvision/rknn/RknnJNI.java

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ public boolean equals(Object obj) {
7474
* @param modelPath Absolute path to the model on disk
7575
* @param numClasses How many classes. MUST MATCH or native code segfaults
7676
* @param modelVer Which model is being used. Detections will be incorrect if not set to
77-
* corrresponding model.
77+
* corresponding model.
7878
* @return Pointer to the detector in native memory
7979
*/
8080
public static native long create(String modelPath, int numClasses, int modelVer, int coreNum);
@@ -89,11 +89,11 @@ public boolean equals(Object obj) {
8989
*/
9090
public static native int setCoreMask(long ptr, int desiredCore);
9191

92-
/** Delete all native resources assocated with a detector */
92+
/** Delete all native resources associated with a detector */
9393
public static native long destroy(long ptr);
9494

9595
/**
96-
* Run detction
96+
* Run detection
9797
*
9898
* @param detectorPtr Pointer to detector created above
9999
* @param imagePtr Pointer to a cv::Mat input image
@@ -102,4 +102,12 @@ public boolean equals(Object obj) {
102102
*/
103103
public static native RknnResult[] detect(
104104
long detectorPtr, long imagePtr, double nmsThresh, double boxThresh);
105+
106+
/**
107+
* Check if a model is quantized (int8)
108+
*
109+
* @param detectorPtr Pointer to a detected created above
110+
* @return true if quantized, false if not
111+
*/
112+
public static native boolean isQuantized(long detectorPtr);
105113
}

src/main/native/cpp/rknn_jni.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,4 +143,16 @@ Java_org_photonvision_rknn_RknnJNI_detect
143143
return jarr;
144144
}
145145

146+
/*
147+
* Class: org_photonvision_rknn_RknnJNI
148+
* Method: isQuantized
149+
* Signature: (J)Z
150+
*/
151+
JNIEXPORT jboolean JNICALL
152+
Java_org_photonvision_rknn_RknnJNI_isQuantized
153+
(JNIEnv *env, jclass, jlong detector_)
154+
{
155+
YoloModel *yolo = reinterpret_cast<YoloModel *>(detector_);
156+
return yolo->is_quant ? JNI_TRUE : JNI_FALSE;
157+
}
146158
} // extern "C"

src/main/native/cpp/rknn_jni.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,14 @@ JNIEXPORT void JNICALL Java_org_photonvision_rknn_RknnJNI_destroy(JNIEnv *,
5757
JNIEXPORT jobjectArray JNICALL Java_org_photonvision_rknn_RknnJNI_detect(
5858
JNIEnv *, jclass, jlong, jlong, jdouble, jdouble);
5959

60+
/*
61+
* Class: org_photonvision_rknn_RknnJNI
62+
* Method: isQuantized
63+
* Signature: (J)Z
64+
*/
65+
JNIEXPORT jboolean JNICALL
66+
Java_org_photonvision_rknn_RknnJNI_isQuantized(JNIEnv *, jclass, jlong);
67+
6068
#ifdef __cplusplus
6169
} // extern "C"
6270
#endif

src/main/native/include/yolo_common.hpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ class YoloModel {
4444

4545
~YoloModel();
4646

47+
bool is_quant;
48+
4749
protected:
4850
virtual detect_result_group_t postProcess(std::vector<rknn_output> output,
4951
DetectionFilterParams params,
@@ -63,7 +65,6 @@ class YoloModel {
6365
rknn_input_output_num io_num;
6466
std::vector<rknn_tensor_attr> input_attrs;
6567
std::vector<rknn_tensor_attr> output_attrs;
66-
bool is_quant;
6768

6869
rknn_input inputs[1];
6970
};

src/test/java/org/photonvision/rknn/RknnTest.java

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
package org.photonvision.rknn;
1919

20+
import static org.junit.jupiter.api.Assertions.*;
21+
2022
import edu.wpi.first.util.CombinedRuntimeLoader;
2123
import java.io.IOException;
2224
import java.util.Arrays;
@@ -68,4 +70,32 @@ public void testBasicBlobs() {
6870
System.out.println("Killing detector");
6971
RknnJNI.destroy(ptr);
7072
}
73+
74+
private boolean checkQuant(String path, int numClasses, ModelVersion modelVer) {
75+
long ptr = RknnJNI.create(path, numClasses, modelVer.ordinal(), 210);
76+
77+
if (ptr <= 0) {
78+
throw new RuntimeException("Failed to create model");
79+
}
80+
81+
boolean isQuantized = RknnJNI.isQuantized(ptr);
82+
83+
RknnJNI.destroy(ptr);
84+
85+
return isQuantized;
86+
}
87+
88+
@Test
89+
public void testQuantizationCheck() {
90+
System.out.println("Loading rknn-jni");
91+
System.load("/home/coolpi/rknn_jni/cmake_build/librknn_jni.so");
92+
93+
assertTrue(
94+
checkQuant("/home/coolpi/rknn_jni/note-640-640-yolov5s.rknn", 1, ModelVersion.YOLO_V5));
95+
assertFalse(
96+
checkQuant(
97+
"/home/coolpi/rknn_jni/coralAlgaeModelNoQuant-640-640-yolov8s.rknn",
98+
2,
99+
ModelVersion.YOLO_V8));
100+
}
71101
}

0 commit comments

Comments
 (0)