Add custom models

docs/source/docs/objectDetection/about-object-detection.md

@@ -43,4 +43,4 @@ Coming soon!
 PhotonVision currently ONLY supports YOLOv5 models trained and converted to `.rknn` format for RK3588 CPUs! Other models require different post-processing code and will NOT work. The model conversion process is also highly particular. Proceed with care.
 :::
 
-Use a program like WinSCP or FileZilla to access your coprocessor's filesystem, and copy the new `.rknn` model file into /home/pi. Next, SSH into the coprocessor and `sudo mv /path/to/new/model.rknn /opt/photonvision/photonvision_config/models/NEW-MODEL-NAME.rknn`. Repeat this process with the labels file, which should contain one line per label the model outputs with no training newline. Next, restart PhotonVision via the web UI.
+In the settings, under ``Device Control``, there's an option to upload a new object detection model. When uploading the files, ensure that that the labels file has the same name as the RKNN file, with ``-labels`` appended to the end. For example, if the RKNN file is named ``foo.rknn``, the labels file should be named ``foo-labels.txt``. The labels file should contain one line per label the model outputs with no training newline.

photon-client/src/components/settings/ObjectDetectionCard.vue

@@ -0,0 +1,186 @@
+<script setup lang="ts">
+import { ref, computed } from "vue";
+import axios from "axios";
+import { useStateStore } from "@/stores/StateStore";
+import { useSettingsStore } from "@/stores/settings/GeneralSettingsStore";
+
+const showObjectDetectionImportDialog = ref(false);
+const importRKNNFile = ref<File | null>(null);
+const importLabelsFile = ref<File | null>(null);
+
+const handleObjectDetectionImport = () => {
+  if (importRKNNFile.value === null || importLabelsFile.value === null) return;
+
+  const formData = new FormData();
+  formData.append("rknn", importRKNNFile.value);
+  formData.append("labels", importLabelsFile.value);
+
+  useStateStore().showSnackbarMessage({
+    message: "Importing Object Detection Model...",
+    color: "secondary",
+    timeout: -1
+  });
+
+  axios
+    .post("/utils/importObjectDetectionModel", formData, {
+      headers: { "Content-Type": "multipart/form-data" }
+    })
+    .then((response) => {
+      useStateStore().showSnackbarMessage({
+        message: response.data.text || response.data,
+        color: "success"
+      });
+    })
+    .catch((error) => {
+      if (error.response) {
+        useStateStore().showSnackbarMessage({
+          color: "error",
+          message: error.response.data.text || error.response.data
+        });
+      } else if (error.request) {
+        useStateStore().showSnackbarMessage({
+          color: "error",
+          message: "Error while trying to process the request! The backend didn't respond."
+        });
+      } else {
+        useStateStore().showSnackbarMessage({
+          color: "error",
+          message: "An error occurred while trying to process the request."
+        });
+      }
+    });
+
+  showObjectDetectionImportDialog.value = false;
+  importRKNNFile.value = null;
+  importLabelsFile.value = null;
+};
+
+// Filters out models that are not supported by the current backend, and returns a flattened list.
+const supportedModels = computed(() => {
+  const { availableModels, supportedBackends } = useSettingsStore().general;
+  return supportedBackends.flatMap((backend) => availableModels[backend] || []);
+});
+</script>
+
+<template>
+  <v-card dark class="mb-3" style="background-color: #006492">
+    <v-card-title class="pa-6">Object Detection</v-card-title>
+    <div class="pa-6 pt-0">
+      <v-row>
+        <v-col cols="12 ">
+          <v-btn color="secondary" @click="() => (showObjectDetectionImportDialog = true)" class="justify-center">
+            <v-icon left class="open-icon"> mdi-import </v-icon>
+            <span class="open-label">Import New Model</span>
+          </v-btn>
+          <v-dialog
+            v-model="showObjectDetectionImportDialog"
+            width="600"
+            @input="
+              () => {
+                importRKNNFile = null;
+                importLabelsFile = null;
+              }
+            "
+          >
+            <v-card color="primary" dark>
+              <v-card-title>Import New Object Detection Model</v-card-title>
+              <v-card-text>
+                Upload a new object detection model to this device that can be used in a pipeline. Naming convention is
+                that the labels file ought to have the same name as the RKNN file, with -labels appended to the end. For
+                example, if the RKNN file is named <i>foo.rknn</i>, the labels file should be named
+                <i>foo-labels.txt</i>.
+                <v-row class="mt-6 ml-4 mr-8">
+                  <v-file-input label="RKNN File" v-model="importRKNNFile" accept=".rknn" />
+                </v-row>
+                <v-row class="mt-6 ml-4 mr-8">
+                  <v-file-input label="Labels File" v-model="importLabelsFile" accept=".txt" />
+                </v-row>
+                <v-row
+                  class="mt-12 ml-8 mr-8 mb-1"
+                  style="display: flex; align-items: center; justify-content: center"
+                  align="center"
+                >
+                  <v-btn
+                    color="secondary"
+                    :disabled="importRKNNFile === null || importLabelsFile === null"
+                    @click="handleObjectDetectionImport"
+                  >
+                    <v-icon left class="open-icon"> mdi-import </v-icon>
+                    <span class="open-label">Import Object Detection Model</span>
+                  </v-btn>
+                </v-row>
+              </v-card-text>
+            </v-card>
+          </v-dialog>
+        </v-col>
+      </v-row>
+      <v-row>
+        <v-col cols="12">
+          <v-simple-table fixed-header height="100%" dense dark>
+            <thead style="font-size: 1.25rem">
+              <tr>
+                <th class="text-left">Available Models</th>
+              </tr>
+            </thead>
+            <tbody>
+              <tr v-for="model in supportedModels" :key="model">
+                <td>{{ model }}</td>
+              </tr>
+            </tbody>
+          </v-simple-table>
+        </v-col>
+      </v-row>
+    </div>
+  </v-card>
+</template>
+
+<style scoped lang="scss">
+.v-btn {
+  width: 100%;
+}
+@media only screen and (max-width: 351px) {
+  .open-icon {
+    margin: 0 !important;
+  }
+  .open-label {
+    display: none;
+  }
+}
+.v-data-table {
+  width: 100%;
+  height: 100%;
+  text-align: center;
+  background-color: #006492 !important;
+
+  th,
+  td {
+    background-color: #006492 !important;
+    font-size: 1rem !important;
+    color: white !important;
+  }
+
+  td {
+    font-family: monospace !important;
+  }
+
+  tbody :hover td {
+    background-color: #005281 !important;
+  }
+
+  ::-webkit-scrollbar {
+    width: 0;
+    height: 0.55em;
+    border-radius: 5px;
+  }
+
+  ::-webkit-scrollbar-track {
+    -webkit-box-shadow: inset 0 0 6px rgba(0, 0, 0, 0.3);
+    border-radius: 10px;
+  }
+
+  ::-webkit-scrollbar-thumb {
+    background-color: #ffd843;
+    border-radius: 10px;
+  }
+}
+</style>

photon-client/src/views/GeneralSettingsView.vue

@@ -1,6 +1,7 @@
 <script setup lang="ts">
 import MetricsCard from "@/components/settings/MetricsCard.vue";
 import DeviceControlCard from "@/components/settings/DeviceControlCard.vue";
+import ObjectDetectionCard from "@/components/settings/ObjectDetectionCard.vue";
 import NetworkingCard from "@/components/settings/NetworkingCard.vue";
 import LightingControlCard from "@/components/settings/LEDControlCard.vue";
 import { useSettingsStore } from "@/stores/settings/GeneralSettingsStore";
@@ -12,6 +13,7 @@ import ApriltagControlCard from "@/components/settings/ApriltagControlCard.vue";
     <MetricsCard />
     <DeviceControlCard />
     <NetworkingCard />
+    <ObjectDetectionCard />
     <LightingControlCard v-if="useSettingsStore().lighting.supported" />
     <ApriltagControlCard />
   </div>

photon-server/src/main/java/org/photonvision/server/RequestHandler.java

@@ -98,7 +98,8 @@ public static void onSettingsImportRequest(Context ctx) {
 
         ConfigManager.getInstance().setWriteTaskEnabled(false);
         ConfigManager.getInstance().disableFlushOnShutdown();
-        // We want to delete the -whole- zip file, so we need to teardown loggers for now
+        // We want to delete the -whole- zip file, so we need to teardown loggers for
+        // now
         logger.info("Writing new settings zip (logs may be truncated)...");
         Logger.closeAllLoggers();
         if (ConfigManager.saveUploadedSettingsZip(tempFilePath.get())) {
@@ -543,6 +544,68 @@ public static void onProgramRestartRequest(Context ctx) {
         restartProgram();
     }
 
+    public static void onObjectDetectionModelImportRequest(Context ctx) {
+        try {
+            // Retrieve the uploaded files
+            var modelFile = ctx.uploadedFile("rknn");
+            var labelsFile = ctx.uploadedFile("labels");
+
+            if (modelFile == null || labelsFile == null) {
+                ctx.status(400);
+                ctx.result(
+                        "No File was sent with the request. Make sure that the model and labels files are sent at the keys 'rknn' and 'labels'");
+                logger.error(
+                        "No File was sent with the request. Make sure that the model and labels files are sent at the keys 'rknn' and 'labels'");
+                return;
+            }
+
+            if (!modelFile.extension().contains("rknn") || !labelsFile.extension().contains("txt")) {
+                ctx.status(400);
+                ctx.result(
+                        "The uploaded files were not of type 'rknn' and 'txt'. The uploaded files should be a .rknn and .txt file.");
+                logger.error(
+                        "The uploaded files were not of type 'rknn' and 'txt'. The uploaded files should be a .rknn and .txt file.");
+                return;
+            }
+
+            // verify naming convention
+            // this check will need to be modified if different model types are added
+            String expectedLabelsFileName = modelFile.filename().replace(".rknn", "-labels.txt");
+
+            if (!labelsFile.filename().equals(expectedLabelsFileName)) {
+                ctx.status(400);
+                ctx.result(
+                        "The labels file name does not follow the naming convention. Expected: "
+                                + expectedLabelsFileName);
+                logger.error(
+                        "The labels file name does not follow the naming convention. Expected: "
+                                + expectedLabelsFileName);
+                return;
+            }
+
+            // dump files into models directory
+
+            var modelPath =
+                    Paths.get(
+                            ConfigManager.getInstance().getModelsDirectory().toString(), modelFile.filename());
+            var labelsPath =
+                    Paths.get(
+                            ConfigManager.getInstance().getModelsDirectory().toString(), labelsFile.filename());
+
+            try (FileOutputStream out = new FileOutputStream(modelPath.toFile())) {
+                modelFile.content().transferTo(out);
+            }
+
+            try (FileOutputStream out = new FileOutputStream(labelsPath.toFile())) {
+                labelsFile.content().transferTo(out);
+            }
+
+            ctx.status(200).result("Successfully uploaded object detection model");
+        } catch (Exception e) {
+            ctx.status(500).result("Error processing files: " + e.getMessage());
+        }
+    }
+
     public static void onDeviceRestartRequest(Context ctx) {
         ctx.status(HardwareManager.getInstance().restartDevice() ? 204 : 500);
     }
@@ -602,7 +665,8 @@ public static void onCalibrationSnapshotRequest(Context ctx) {
             return;
         }
 
-        // encode as jpeg to save even more space. reduces size of a 1280p image from 300k to 25k
+        // encode as jpeg to save even more space. reduces size of a 1280p image from
+        // 300k to 25k
         var jpegBytes = new MatOfByte();
         Mat img = null;
         try {

photon-server/src/main/java/org/photonvision/server/Server.java

@@ -127,6 +127,9 @@ private static void start(int port) {
 
         // Utilities
         app.post("/api/utils/offlineUpdate", RequestHandler::onOfflineUpdateRequest);
+        app.post(
+                "/api/utils/importObjectDetectionModel",
+                RequestHandler::onObjectDetectionModelImportRequest);
         app.get("/api/utils/photonvision-journalctl.txt", RequestHandler::onLogExportRequest);
         app.post("/api/utils/restartProgram", RequestHandler::onProgramRestartRequest);
         app.post("/api/utils/restartDevice", RequestHandler::onDeviceRestartRequest);