Add support for different RKNN YOLO models in the backend (#1205)

Author: james20902

build.gradle

@@ -30,7 +30,7 @@ ext {
     joglVersion = "2.4.0-rc-20200307"
     javalinVersion = "5.6.2"
     photonGlDriverLibVersion = "dev-v2023.1.0-9-g75fc678"
-    rknnVersion = "dev-v2024.0.0-30-g001b5ec"
+    rknnVersion = "dev-v2024.0.0-64-gc0836a6"
     frcYear = "2024"
     mrcalVersion = "dev-v2024.0.0-7-gc976aaa";
 

photon-core/build.gradle

@@ -37,9 +37,8 @@ dependencies {
     implementation 'org.zeroturnaround:zt-zip:1.14'
 
     implementation "org.xerial:sqlite-jdbc:3.41.0.0"
-    def rknnjniversion = "dev-v2024.0.0-44-g8022c40"
-    implementation "org.photonvision:rknn_jni-jni:$rknnjniversion:linuxarm64"
-    implementation "org.photonvision:rknn_jni-java:$rknnjniversion"
+    implementation "org.photonvision:rknn_jni-jni:$rknnVersion:linuxarm64"
+    implementation "org.photonvision:rknn_jni-java:$rknnVersion"
     implementation "org.photonvision:photon-libcamera-gl-driver-jni:$photonGlDriverLibVersion:linuxarm64"
     implementation "org.photonvision:photon-libcamera-gl-driver-java:$photonGlDriverLibVersion"
 

photon-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java

@@ -25,12 +25,14 @@
 import java.util.List;
 import org.photonvision.common.logging.LogGroup;
 import org.photonvision.common.logging.Logger;
+import org.photonvision.rknn.RknnJNI;
 
 public class NeuralNetworkModelManager {
     private static NeuralNetworkModelManager INSTANCE;
     private static final Logger logger = new Logger(NeuralNetworkModelManager.class, LogGroup.Config);
 
     private final String MODEL_NAME = "note-640-640-yolov5s.rknn";
+    private final RknnJNI.ModelVersion modelVersion = RknnJNI.ModelVersion.YOLO_V5;
     private File defaultModelFile;
     private List<String> labels;
 
@@ -51,7 +53,7 @@ public void initialize(File modelsFolder) {
         this.defaultModelFile = new File(modelsFolder, MODEL_NAME);
         extractResource(modelResourcePath, defaultModelFile);
 
-        File labelsFile = new File(modelsFolder, "labels.txt");
+        File labelsFile = new File(modelsFolder, "labels_v5.txt");
         var labelResourcePath = "/models/" + labelsFile.getName();
         extractResource(labelResourcePath, labelsFile);
 
@@ -95,4 +97,8 @@ public File getDefaultRknnModel() {
     public List<String> getLabels() {
         return labels;
     }
+
+    public RknnJNI.ModelVersion getModelVersion() {
+        return modelVersion;
+    }
 }

photon-core/src/main/java/org/photonvision/jni/RknnDetectorJNI.java

@@ -65,12 +65,11 @@ public static class RknnObjectDetector {
         long objPointer = -1;
         private List<String> labels;
         private final Object lock = new Object();
-
         private static final CopyOnWriteArrayList<Long> detectors = new CopyOnWriteArrayList<>();
 
-        public RknnObjectDetector(String modelPath, List<String> labels) {
+        public RknnObjectDetector(String modelPath, List<String> labels, RknnJNI.ModelVersion version) {
             synchronized (lock) {
-                objPointer = RknnJNI.create(modelPath, labels.size());
+                objPointer = RknnJNI.create(modelPath, labels.size(), version.ordinal(), -1);
                 detectors.add(objPointer);
                 System.out.println(
                         "Created " + objPointer + "! Detectors: " + Arrays.toString(detectors.toArray()));

photon-core/src/main/java/org/photonvision/vision/pipe/impl/RknnDetectionPipe.java

@@ -35,7 +35,8 @@ public RknnDetectionPipe() {
         this.detector =
                 new RknnObjectDetector(
                         NeuralNetworkModelManager.getInstance().getDefaultRknnModel().getAbsolutePath(),
-                        NeuralNetworkModelManager.getInstance().getLabels());
+                        NeuralNetworkModelManager.getInstance().getLabels(),
+                        NeuralNetworkModelManager.getInstance().getModelVersion());
     }
 
     @Override