Yolo duplication fix (#1713)

Somebody wanted a description, so here I am. This PR fixes an error which caused the discoverModels function to be rerun after each upload of a new model, but without clearing the list of available models. This causes any models that were on the list prior to the import to be duplicated. This PR also makes it so that uploading a model automatically updates the list of available models. --------- Co-authored-by: Matt Morley <[email protected]> Co-authored-by: Chris Gerth <[email protected]>
PhotonVision · Jan 13, 2025 · 966b9e8 · 966b9e8
1 parent 5f75619
commit 966b9e8
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 21 deletions.
diff --git a/photon-client/src/components/settings/ObjectDetectionCard.vue b/photon-client/src/components/settings/ObjectDetectionCard.vue
@@ -4,11 +4,12 @@ import axios from "axios";
 import { useStateStore } from "@/stores/StateStore";
 import { useSettingsStore } from "@/stores/settings/GeneralSettingsStore";
 
-const showObjectDetectionImportDialog = ref(false);
+const showImportDialog = ref(false);
 const importRKNNFile = ref<File | null>(null);
 const importLabelsFile = ref<File | null>(null);
 
-const handleObjectDetectionImport = () => {
+// TODO gray out the button when model is uploading
+const handleImport = async () => {
   if (importRKNNFile.value === null || importLabelsFile.value === null) return;
 
   const formData = new FormData();
@@ -50,7 +51,8 @@ const handleObjectDetectionImport = () => {
       }
     });
 
-  showObjectDetectionImportDialog.value = false;
+  showImportDialog.value = false;
+
   importRKNNFile.value = null;
   importLabelsFile.value = null;
 };
@@ -68,12 +70,12 @@ const supportedModels = computed(() => {
     <div class="pa-6 pt-0">
       <v-row>
         <v-col cols="12 ">
-          <v-btn color="secondary" @click="() => (showObjectDetectionImportDialog = true)" class="justify-center">
+          <v-btn color="secondary" @click="() => (showImportDialog = 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"
+            v-model="showImportDialog"
             width="600"
             @input="
               () => {
@@ -105,7 +107,7 @@ const supportedModels = computed(() => {
                   <v-btn
                     color="secondary"
                     :disabled="importRKNNFile === null || importLabelsFile === null"
-                    @click="handleObjectDetectionImport"
+                    @click="handleImport"
                   >
                     <v-icon left class="open-icon"> mdi-import </v-icon>
                     <span class="open-label">Import Object Detection Model</span>

diff --git a/...n-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java b/...n-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java
@@ -221,29 +221,28 @@ private void loadModel(File model) {
     /**
      * Discovers DNN models from the specified folder.
      *
-     * @param modelsFolder The folder where the models are stored
+     * @param modelsDirectory The folder where the models are stored
      */
-    public void discoverModels(File modelsFolder) {
+    public void discoverModels(File modelsDirectory) {
         logger.info("Supported backends: " + supportedBackends);
 
-        if (!modelsFolder.exists()) {
-            logger.error("Models folder " + modelsFolder.getAbsolutePath() + " does not exist.");
+        if (!modelsDirectory.exists()) {
+            logger.error("Models folder " + modelsDirectory.getAbsolutePath() + " does not exist.");
             return;
         }
 
-        if (models == null) {
-            models = new HashMap<>();
-        }
+        models = new HashMap<>();
 
         try {
-            Files.walk(modelsFolder.toPath())
+            Files.walk(modelsDirectory.toPath())
                     .filter(Files::isRegularFile)
                     .forEach(path -> loadModel(path.toFile()));
         } catch (IOException e) {
-            logger.error("Failed to discover models at " + modelsFolder.getAbsolutePath(), e);
+            logger.error("Failed to discover models at " + modelsDirectory.getAbsolutePath(), e);
         }
 
-        // After loading all of the models, sort them by name to ensure a consistent ordering
+        // After loading all of the models, sort them by name to ensure a consistent
+        // ordering
         models.forEach(
                 (backend, backendModels) ->
                         backendModels.sort((a, b) -> a.getName().compareTo(b.getName())));

diff --git a/photon-server/src/main/java/org/photonvision/server/RequestHandler.java b/photon-server/src/main/java/org/photonvision/server/RequestHandler.java
@@ -42,7 +42,9 @@
 import org.photonvision.common.dataflow.DataChangeDestination;
 import org.photonvision.common.dataflow.DataChangeService;
 import org.photonvision.common.dataflow.events.IncomingWebSocketEvent;
+import org.photonvision.common.dataflow.events.OutgoingUIEvent;
 import org.photonvision.common.dataflow.networktables.NetworkTablesManager;
+import org.photonvision.common.dataflow.websocket.UIPhotonConfiguration;
 import org.photonvision.common.hardware.HardwareManager;
 import org.photonvision.common.hardware.Platform;
 import org.photonvision.common.logging.LogGroup;
@@ -546,7 +548,7 @@ public static void onProgramRestartRequest(Context ctx) {
         restartProgram();
     }
 
-    public static void onObjectDetectionModelImportRequest(Context ctx) {
+    public static void onImportObjectDetectionModelRequest(Context ctx) {
         try {
             // Retrieve the uploaded files
             var modelFile = ctx.uploadedFile("rknn");
@@ -579,7 +581,11 @@ public static void onObjectDetectionModelImportRequest(Context ctx) {
                     Pattern.compile("^[a-zA-Z0-9]+-\\d+-\\d+-yolov[58][a-z]*-labels\\.txt$");
 
             if (!modelPattern.matcher(modelFile.filename()).matches()
-                    || !labelsPattern.matcher(labelsFile.filename()).matches()) {
+                    || !labelsPattern.matcher(labelsFile.filename()).matches()
+                    || !(modelFile
+                            .filename()
+                            .substring(0, modelFile.filename().indexOf("-"))
+                            .equals(labelsFile.filename().substring(0, labelsFile.filename().indexOf("-"))))) {
                 ctx.status(400);
                 ctx.result("The uploaded files were not named correctly.");
                 logger.error("The uploaded object detection model files were not named correctly.");
@@ -610,6 +616,12 @@ public static void onObjectDetectionModelImportRequest(Context ctx) {
         } catch (Exception e) {
             ctx.status(500).result("Error processing files: " + e.getMessage());
         }
+
+        DataChangeService.getInstance()
+                .publishEvent(
+                        new OutgoingUIEvent<>(
+                                "fullsettings",
+                                UIPhotonConfiguration.programStateToUi(ConfigManager.getInstance().getConfig())));
     }
 
     public static void onDeviceRestartRequest(Context ctx) {

diff --git a/photon-server/src/main/java/org/photonvision/server/Server.java b/photon-server/src/main/java/org/photonvision/server/Server.java
@@ -102,7 +102,7 @@ private static void start(int port) {
                                     });
                         });
 
-        /*Web Socket Events for Data Exchange */
+        /* Web Socket Events for Data Exchange */
         var dsHandler = DataSocketHandler.getInstance();
         app.ws(
                 "/websocket_data",
@@ -112,7 +112,7 @@ private static void start(int port) {
                     ws.onBinaryMessage(dsHandler::onBinaryMessage);
                 });
 
-        /*API Events*/
+        /* API Events */
         // Settings
         app.post("/api/settings", RequestHandler::onSettingsImportRequest);
         app.get("/api/settings/photonvision_config.zip", RequestHandler::onSettingsExportRequest);
@@ -129,7 +129,7 @@ private static void start(int port) {
         app.post("/api/utils/offlineUpdate", RequestHandler::onOfflineUpdateRequest);
         app.post(
                 "/api/utils/importObjectDetectionModel",
-                RequestHandler::onObjectDetectionModelImportRequest);
+                RequestHandler::onImportObjectDetectionModelRequest);
         app.get("/api/utils/photonvision-journalctl.txt", RequestHandler::onLogExportRequest);
         app.post("/api/utils/restartProgram", RequestHandler::onProgramRestartRequest);
         app.post("/api/utils/restartDevice", RequestHandler::onDeviceRestartRequest);