writing image embedding computation feature

Author: Gin

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingModel.cpp

@@ -5,6 +5,8 @@
  *  Run Segment Anything Model (SAM) to segment objects on images
  */
 
+#include <QDir>
+#include <QDirIterator>
 #include <fstream>
 #include <iostream>
 #include <onnxruntime_cxx_api.h>
@@ -246,5 +248,20 @@ bool load_image_embedding(const std::string& image_filepath, std::vector<float>&
 }
 
 
+void compute_embeddings_for_folder(const std::string& image_folder_path){
+    QDir image_dir(image_folder_path.c_str());
+    if (!image_dir.exists()){
+        std::cerr << "Error: input image folder path " << image_folder_path << " does not exist." << std::endl;
+        return;
+    }
+
+    QDirIterator image_file_iter(image_dir.absolutePath(), {"*.png", "*.jpg", "*.jpeg"}, QDir::Files, QDirIterator::Subdirectories);
+    std::vector<std::string> all_image_paths;
+    while (image_file_iter.hasNext()){
+        all_image_paths.emplace_back(image_file_iter.next().toStdString());
+    }
+    std::cout << "Found " << all_image_paths.size() << " images recursively in folder " << image_folder_path << std::endl;
+}
+
 }
 }

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingModel.h

@@ -21,13 +21,18 @@ namespace PokemonAutomation{
 namespace ML{
 
 
-// load pre-computed image embedding from disk
-// return true if there is the embedding file
+// Load pre-computed image embedding from disk
+// Return true if there is the embedding file.
+// The embedding is stored in a file in the same folder as the image, having the same name but with a suffix ".embedding".
 bool load_image_embedding(const std::string& image_filepath, std::vector<float>& image_embedding);
 
-// save the image embedding as a file with path <image_filepath>.embedding
+// Save the image embedding as a file with path <image_filepath>.embedding.
 void save_image_embedding_to_disk(const std::string& image_filepath, const std::vector<float>& embedding);
 
+// Compute embeddings for all images in a folder.
+// This can be very slow!
+void compute_embeddings_for_folder(const std::string& image_folder_path);
+
 
 class SAMEmbedderSession{
 public:

SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp

@@ -4,8 +4,10 @@
  *
  */
 
+#include <QFileDialog>
 #include <QLabel>
 #include <QDir>
+#include <QDirIterator>
 #include <QVBoxLayout>
 #include <QGraphicsView>
 #include <QGraphicsScene>
@@ -407,6 +409,12 @@ void LabelImages::compute_mask(VideoOverlaySet& overlay_set){
     }
 }
 
+void LabelImages::compute_embeddings_for_folder(const std::string& image_folder_path){
+    ML::compute_embeddings_for_folder(image_folder_path);
+}
+
+
+
 LabelImages_Widget::~LabelImages_Widget(){
     m_program.FORM_LABEL.remove_listener(*this);
     delete m_switch_widget;
@@ -453,9 +461,21 @@ LabelImages_Widget::LabelImages_Widget(
         program.update_rendered_objects(this->m_overlay_set);
     });
 
+    // Add all option UI elements defined by LabelImage program.
     m_option_widget = program.m_options.make_QtWidget(*scroll_inner);
     scroll_layout->addWidget(&m_option_widget->widget());
 
+    button = new QPushButton("Compute Image Embeddings (SLOW!)", scroll_inner);
+    scroll_layout->addWidget(button);
+    connect(button, &QPushButton::clicked, this, [this](bool){
+        std::string folder_path = QFileDialog::getExistingDirectory(
+            nullptr, "Open image folder", ".").toStdString();
+
+        if (folder_path.size() > 0){
+            this->m_program.compute_embeddings_for_folder(folder_path);
+        }
+    });
+
     const auto cur_res = m_display_session.video_session().current_resolution();
     if (cur_res.width > 0 && cur_res.height > 0){
         const std::string& image_path = m_display_session.option().m_image_path;

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -60,23 +60,34 @@ class LabelImages_Descriptor : public PanelDescriptor{
     LabelImages_Descriptor();
 };
 
-// label image program
+
+// Program to annoatation images for training ML models
 class LabelImages : public PanelInstance{
 public:
     LabelImages(const LabelImages_Descriptor& descriptor);
     virtual QWidget* make_widget(QWidget& parent, PanelHolder& holder) override;
 
 public:
-    //  Serialization
+    // Serialization
     virtual void from_json(const JsonValue& json) override;
     virtual JsonValue to_json() const override;
 
+    // Load image related data:
+    // - Image SAM embedding data file, which has the same file path but with a name suffix ".embedding"
+    // - Existing annotation file, which is stored in a pre-defined ML_ANNOTATION_PATH() and with the same filename as
+    //   the image but with name extension replaced to be ".json".
     void load_image_related_data(const std::string& image_path, const size_t source_image_width, const size_t source_image_height);
 
+    // Update rendering data reflect the current annotation
     void update_rendered_objects(VideoOverlaySet& overlayset);
 
+    // Use user currently drawn box to compute per-pixel masks on the image using SAM model
     void compute_mask(VideoOverlaySet& overlay_set);
 
+    // Compute embeddings for all images in a folder.
+    // This can be very slow!
+    void compute_embeddings_for_folder(const std::string& image_folder);
+
 private:
     friend class LabelImages_Widget;
     friend class DrawnBoundingBox;

SerialPrograms/Source/ML/UI/ML_ImageAnnotationCommandRow.cpp

@@ -86,11 +86,6 @@ ImageAnnotationCommandRow::ImageAnnotationCommandRow(
     m_save_profile_button = new QPushButton("Save Profile", this);
     row->addWidget(m_save_profile_button, 2);
 
-    m_screenshot_button = new QPushButton("Screenshot", this);
-//    m_screenshot_button->setToolTip("Take a screenshot of the console and save to disk.");
-    row->addWidget(m_screenshot_button, 2);
-
-
 //    m_test_button = new QPushButton("Test Button", this);
 //    row->addWidget(m_test_button, 3);
 
@@ -145,11 +140,7 @@ ImageAnnotationCommandRow::ImageAnnotationCommandRow(
         m_save_profile_button, &QPushButton::clicked,
         this, [this](bool) { emit save_profile(); }
     );
-    connect(
-        m_screenshot_button, &QPushButton::clicked,
-        this, [this](bool){ emit screenshot_requested(); }
-    );
-
+    
 #if (QT_VERSION_MAJOR == 6) && (QT_VERSION_MINOR >= 8)
     if (IS_BETA_VERSION || PreloadSettings::instance().DEVELOPER_MODE){
         m_video_button = new QPushButton("Video Capture", this);

SerialPrograms/Source/ML/UI/ML_ImageAnnotationCommandRow.h

@@ -71,7 +71,6 @@ class ImageAnnotationCommandRow :
 
     QPushButton* m_load_profile_button;
     QPushButton* m_save_profile_button;
-    QPushButton* m_screenshot_button;
     QPushButton* m_video_button;
     bool m_last_known_focus;
     ProgramState m_last_known_state;

SerialPrograms/Source/ML/UI/ML_ImageAnnotationSourceSelectorWidget.cpp

@@ -55,11 +55,13 @@ ImageAnnotationSourceSelectorWidget::ImageAnnotationSourceSelectorWidget(ImageAn
         m_reset_button, &QPushButton::clicked,
         this, [this](bool){
             std::string path = QFileDialog::getOpenFileName(
-                nullptr, "Open image file", ".", "*.png *.jpg"
+                nullptr, "Open image file", ".", "*.png *.jpg *.jpeg"
             ).toStdString();
-        
-            m_source_file_path_label->setText(QString::fromStdString(path));
-            m_session.set_image_source(path);
+
+            if (path.size() > 0){
+                m_source_file_path_label->setText(QString::fromStdString(path));
+                m_session.set_image_source(path);
+            }
         }
     );