add handling for missing ML model files

Author: Gin

SerialPrograms/Source/ML/DataLabeling/ML_SegmentAnythingModel.cpp

@@ -199,6 +199,23 @@ void compute_embeddings_for_folder(const std::string& embedding_model_path, cons
         return;
     }
 
+    if (!std::filesystem::exists(embedding_model_path)){
+        std::cerr << "Error: no such embedding model path " << embedding_model_path << "." << std::endl;
+        QMessageBox box;
+        box.critical(nullptr, "Embedding Model Does Not Exist",
+            QString::fromStdString("Embedding model path" + embedding_model_path + " does not exist."));
+        return;
+    }
+    // since the embedding model has too many weights, onnx created a .data file to contain weights.
+    auto embedding_model_data_path = embedding_model_path + ".data";
+    if (!std::filesystem::exists(embedding_model_data_path)){
+        std::cerr << "Error: no such embedding model data path " << embedding_model_data_path << "." << std::endl;
+        QMessageBox box;
+        box.critical(nullptr, "Embedding Model Data File Does Not Exist",
+            QString::fromStdString("Embedding model data file path" + embedding_model_data_path + " does not exist."));
+        return;
+    }
+
     SAMEmbedderSession embedding_session(embedding_model_path);
     std::vector<float> output_image_embedding;
     for (size_t i = 0; i < all_image_paths.size(); i++){
@@ -238,7 +255,7 @@ void compute_embeddings_for_folder(const std::string& embedding_model_path, cons
         embedding_session.run(resized_mat, output_image_embedding);
         save_image_embedding_to_disk(image_path, output_image_embedding);
     }
-
+    std::cout << "Done computing embeddings for images in folder " << image_folder_path << "." << std::endl;
 
 }
 

SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp

@@ -205,14 +205,25 @@ LabelImages::LabelImages(const LabelImages_Descriptor& descriptor)
     , WIDTH("<b>Width:</b>", LockMode::UNLOCK_WHILE_RUNNING, 0.4, 0.0, 1.0)
     , HEIGHT("<b>Height:</b>", LockMode::UNLOCK_WHILE_RUNNING, 0.4, 0.0, 1.0)
     , FORM_LABEL("bulbasaur")
-    , m_sam_session{RESOURCE_PATH() + "ML/sam_cpu.onnx"}
 {
     ADD_OPTION(X);
     ADD_OPTION(Y);
     ADD_OPTION(WIDTH);
     ADD_OPTION(HEIGHT);
     ADD_OPTION(FORM_LABEL);
+
+    // , m_sam_session{RESOURCE_PATH() + "ML/sam_cpu.onnx"}
+    const std::string sam_model_path = RESOURCE_PATH() + "ML/sam_cpu.onnx";
+    if (std::filesystem::exists(sam_model_path)){
+        m_sam_session = std::make_unique<SAMSession>(sam_model_path);
+    } else{
+        std::cerr << "Error: no such SAM model path " << sam_model_path << "." << std::endl;
+        QMessageBox box;
+        box.critical(nullptr, "SAM Model Does Not Exist",
+            QString::fromStdString("SAM model path" + sam_model_path + " does not exist."));
+    }
 }
+
 void LabelImages::from_json(const JsonValue& json){
     const JsonObject* obj = json.to_object();
     if (obj == nullptr){
@@ -370,11 +381,11 @@ void LabelImages::compute_mask(VideoOverlaySet& overlay_set){
         return;
     }
 
-    if (m_image_embedding.size() == 0){
+    if (!m_sam_session || m_image_embedding.size() == 0){
         // no embedding file loaded
         return;
     }
-    m_sam_session.run(
+    m_sam_session->run(
         m_image_embedding,
         (int)source_height, (int)source_width, {}, {},
         {box_x, box_y, box_x + box_width, box_y + box_height},

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -8,6 +8,7 @@
 #define PokemonAutomation_ML_LabelImages_H
 
 #include <QGraphicsScene>
+#include <memory>
 #include "Common/Cpp/Options/BatchOption.h"
 #include "Common/Cpp/Options/FloatingPointOption.h"
 #include "CommonFramework/Panels/PanelInstance.h"
@@ -120,7 +121,7 @@ class LabelImages : public PanelInstance{
     // buffer to compute SAM mask on
     ImageRGB32 m_mask_image;
 
-    SAMSession m_sam_session;
+    std::unique_ptr<SAMSession> m_sam_session;
     std::vector<ObjectAnnotation> m_annotations;
     size_t m_last_object_idx = 0;
     std::string m_annotation_file_path;

SerialPrograms/Source/ML/UI/ML_ImageAnnotationDisplayWidget.cpp

@@ -45,6 +45,7 @@ ImageAnnotationDisplayWidget::ImageAnnotationDisplayWidget(
     : QWidget(&parent)
     , m_session(session)
 {
+    
     QVBoxLayout* layout = new QVBoxLayout(this);
     layout->setContentsMargins(0, 0, 0, 0);
     layout->setAlignment(Qt::AlignTop);

SerialPrograms/Source/ML/UI/ML_ImageAnnotationSourceSelectorWidget.cpp

@@ -73,7 +73,7 @@ ImageAnnotationSourceSelectorWidget::ImageAnnotationSourceSelectorWidget(ImageAn
     folder_info_row->addSpacing(2);
     folder_info_row->addWidget(next_image_button, 2);
     folder_info_row->addSpacing(10);
-    folder_info_row->addWidget(new QLabel("     ", this), 10);
+    folder_info_row->addWidget(new QLabel("     ", this), 10);  // empty label to push the buttons above to the left
 
 
     //  Set the action for the video reset button