Merge branch 'main' of https://github.com/PokemonAutomation/Arduino-Source

Author: Mysticial

SerialPrograms/Source/CommonFramework/VideoPipeline/VideoSession.cpp

@@ -12,9 +12,9 @@
 #include "VideoSources/VideoSource_Null.h"
 #include "VideoSession.h"
 
-//#include <iostream>
-//using std::cout;
-//using std::endl;
+// #include <iostream>
+// using std::cout;
+// using std::endl;
 
 namespace PokemonAutomation{
 

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingModel.cpp

@@ -9,8 +9,10 @@
 #include <QDirIterator>
 #include <fstream>
 #include <iostream>
+#include <QMessageBox>
 #include <onnxruntime_cxx_api.h>
 #include <opencv2/imgcodecs.hpp>
+#include <opencv2/imgproc.hpp>
 #include "3rdParty/ONNX/OnnxToolsPA.h"
 #include "SegmentAnythingModel.h"
 
@@ -207,7 +209,7 @@ void SAMSession::run(
 
 // 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){
-    std::string embedding_path = image_filepath + ".embedding";
+    const std::string embedding_path = image_filepath + ".embedding";
     std::ofstream fout(embedding_path, std::ios::binary);
     // write embedding shape
     fout.write(reinterpret_cast<const char*>(&SAM_EMBEDDER_OUTPUT_N_CHANNELS), sizeof(SAM_EMBEDDER_OUTPUT_N_CHANNELS));
@@ -248,7 +250,7 @@ bool load_image_embedding(const std::string& image_filepath, std::vector<float>&
 }
 
 
-void compute_embeddings_for_folder(const std::string& image_folder_path){
+void compute_embeddings_for_folder(const std::string& embedding_model_path, 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;
@@ -261,6 +263,48 @@ void compute_embeddings_for_folder(const std::string& image_folder_path){
         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;
+
+    SAMEmbedderSession embedding_session(embedding_model_path);
+    std::vector<float> output_image_embedding;
+    for (size_t i = 0; i < all_image_paths.size(); i++){
+        const auto& image_path = all_image_paths[i];
+        std::cout << (i+1) << "/" << all_image_paths.size() << ": ";
+        const std::string embedding_path = image_path + ".embedding";
+        if (std::filesystem::exists(embedding_path)){
+            std::cout << "skip already computed embedding " << embedding_path << "." << std::endl;
+            continue;
+        }
+        std::cout << "computing embedding for " << image_path << "..." << std::endl;
+        cv::Mat image_bgr = cv::imread(image_path);
+        if (image_bgr.empty()){
+            std::cerr << "Error: image empty. Probably the file is not an image?" << std::endl;
+            QMessageBox box;
+            box.warning(nullptr, "Unable To Open Image",
+                QString::fromStdString("Cannot open image file " + image_path + ". Probably not an actual image?"));
+            return;
+        }
+        cv::Mat image;
+        if (image_bgr.channels() == 4){
+            cv::cvtColor(image_bgr, image, cv::COLOR_BGRA2RGB);
+        } else if (image_bgr.channels() == 3){
+            cv::cvtColor(image_bgr, image, cv::COLOR_BGR2RGB);
+        } else{
+            std::cerr << "Error: wrong image channels. Only work with RGB or RGBA images." << std::endl;
+            QMessageBox box;
+            box.warning(nullptr, "Wrong Image Channels",
+                QString::fromStdString("Image has " + std::to_string(image_bgr.channels()) + " channels. Only support 3 or 4 channels."));
+            return;
+        }
+
+        cv::Mat resized_mat;  // resize to the shape for the ML model input
+        cv::resize(image, resized_mat, cv::Size(SAM_EMBEDDER_INPUT_IMAGE_WIDTH, SAM_EMBEDDER_INPUT_IMAGE_HEIGHT));
+
+        output_image_embedding.clear();
+        embedding_session.run(resized_mat, output_image_embedding);
+        save_image_embedding_to_disk(image_path, output_image_embedding);
+    }
+
+
 }
 
 }

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingModel.h

@@ -29,9 +29,9 @@ bool load_image_embedding(const std::string& image_filepath, std::vector<float>&
 // 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.
+// Compute embeddings for all images in a folder. Only support .png, .jpg and .jpeg filename extensions so far.
 // This can be very slow!
-void compute_embeddings_for_folder(const std::string& image_folder_path);
+void compute_embeddings_for_folder(const std::string& embedding_model_path, const std::string& image_folder_path);
 
 
 class SAMEmbedderSession{

SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp

@@ -227,6 +227,11 @@ JsonValue LabelImages::to_json() const{
     JsonObject obj = std::move(*m_options.to_json().to_object());
     obj["ImageSetup"] = m_display_option.to_json();
 
+    save_annotation_to_file();
+    return obj;
+}
+
+void LabelImages::save_annotation_to_file() const{
     // m_annotation_file_path
     if (m_annotation_file_path.size() > 0 && !m_fail_to_load_annotation_file){
         JsonArray anno_json_arr;
@@ -236,8 +241,19 @@ JsonValue LabelImages::to_json() const{
         cout << "Saving annotation to " << m_annotation_file_path << endl;
         anno_json_arr.dump(m_annotation_file_path);
     }
-    return obj;
 }
+
+void LabelImages::clear_for_new_image(){
+    source_image_height = source_image_height = 0;
+    m_image_embedding.clear();
+    m_output_boolean_mask.clear();
+    m_mask_image = ImageRGB32();
+    m_annotations.clear();
+    m_last_object_idx = 0;
+    m_annotation_file_path = "";
+    m_fail_to_load_annotation_file = false;
+}
+
 QWidget* LabelImages::make_widget(QWidget& parent, PanelHolder& holder){
     return new LabelImages_Widget(parent, *this, holder);
 }
@@ -253,6 +269,7 @@ void LabelImages::load_image_related_data(const std::string& image_path, size_t
     if (!embedding_loaded){
         return; // no embedding, then no way for us to annotate
     }
+    
     // see if we can load the previously created labels
     const std::string anno_filename = std::filesystem::path(image_path).filename().replace_extension(".json").string();
 
@@ -410,14 +427,16 @@ 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);
+    std::string embedding_model_path = RESOURCE_PATH() + "ML/sam_embedder_cpu.onnx";
+    std::cout << "Use SAM Embedding model " << embedding_model_path << std::endl;
+    ML::compute_embeddings_for_folder(embedding_model_path, image_folder_path);
 }
 
 
 
 LabelImages_Widget::~LabelImages_Widget(){
     m_program.FORM_LABEL.remove_listener(*this);
-    delete m_switch_widget;
+    delete m_image_display_widget;
 }
 LabelImages_Widget::LabelImages_Widget(
     QWidget& parent,
@@ -431,6 +450,9 @@ LabelImages_Widget::LabelImages_Widget(
     , m_drawn_box(*this, m_display_session.overlay())
 {
     m_program.FORM_LABEL.add_listener(*this);
+    m_display_session.video_session().add_state_listener(*this);
+
+    m_embedding_info_label = new QLabel(this);
 
     QVBoxLayout* layout = new QVBoxLayout(this);
     layout->setContentsMargins(0, 0, 0, 0);
@@ -445,8 +467,13 @@ LabelImages_Widget::LabelImages_Widget(
     QVBoxLayout* scroll_layout = new QVBoxLayout(scroll_inner);
     scroll_layout->setAlignment(Qt::AlignTop);
 
-    m_switch_widget = new ImageAnnotationDisplayWidget(*this, m_display_session, 0);
-    scroll_layout->addWidget(m_switch_widget);
+    m_image_display_widget = new ImageAnnotationDisplayWidget(*this, m_display_session, 0);
+    scroll_layout->addWidget(m_image_display_widget);
+
+    QHBoxLayout* embedding_info_row = new QHBoxLayout();
+    scroll_layout->addLayout(embedding_info_row);
+    embedding_info_row->addWidget(new QLabel("<b>Image Embedding File:</b> ", this));    
+    embedding_info_row->addWidget(m_embedding_info_label);
 
     QPushButton* button = new QPushButton("Delete Last Mask", scroll_inner);
     scroll_layout->addWidget(button);
@@ -476,17 +503,14 @@ LabelImages_Widget::LabelImages_Widget(
         }
     });
 
-    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;
-        m_program.load_image_related_data(image_path, cur_res.width, cur_res.height);
-        m_program.update_rendered_objects(m_overlay_set);
-    }
-
-
     cout << "LabelImages_Widget built" << endl;
 }
 
+void LabelImages_Widget::clear_for_new_image(){
+    m_overlay_set.clear();
+    m_program.clear_for_new_image();
+}
+
 void LabelImages_Widget::on_config_value_changed(void* object){
     if (m_program.m_annotations.size() > 0 && m_program.m_last_object_idx < m_program.m_annotations.size()){
         std::string& cur_label = m_program.m_annotations[m_program.m_last_object_idx].label;
@@ -495,6 +519,40 @@ void LabelImages_Widget::on_config_value_changed(void* object){
     }
 }
 
+// This callback function will be called whenever the display source (the image source) is loaded or reloaded:
+void LabelImages_Widget::post_startup(VideoSource* source){
+    const std::string& image_path = m_display_session.option().m_image_path;
+
+    m_program.save_annotation_to_file();  // save the current annotation file
+    clear_for_new_image();
+    if (image_path.size() == 0){
+        m_embedding_info_label->setText("");
+        return;
+    }
+
+    const std::string embedding_path = image_path + ".embedding";
+    const std::string embedding_path_display = "<IMAGE_FOLDER>/" + std::filesystem::path(embedding_path).filename().string();
+    if (!std::filesystem::exists(embedding_path)){
+        m_embedding_info_label->setText(QString::fromStdString(embedding_path_display + " Dose Not Exist. Cannot Annotate The Image!"));
+        m_embedding_info_label->setStyleSheet("color: red");
+        return;
+    }
+        
+    m_embedding_info_label->setText(QString::fromStdString(embedding_path_display));
+    m_embedding_info_label->setStyleSheet("color: blue");
+
+    const auto cur_res = m_display_session.video_session().current_resolution();
+    if (cur_res.width == 0 || cur_res.height == 0){
+        QMessageBox box;
+        box.warning(nullptr, "Invalid Image Dimension",
+            QString::fromStdString("Loaded image " + image_path + " has invalid dimension: " + cur_res.to_string()));
+        return;
+    }
+    
+    m_program.load_image_related_data(image_path, cur_res.width, cur_res.height);
+    m_program.update_rendered_objects(m_overlay_set);
+}
+
 
 
 }

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -25,6 +25,7 @@
 
 class QGraphicsView;
 class QGraphicsPixmapItem;
+class QLabel;
 
 namespace PokemonAutomation{
 
@@ -72,6 +73,11 @@ class LabelImages : public PanelInstance{
     virtual void from_json(const JsonValue& json) override;
     virtual JsonValue to_json() const override;
 
+    void save_annotation_to_file() const;
+
+    // called after loading a new image, clean up all internal data 
+    void clear_for_new_image();
+
     // 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
@@ -145,7 +151,7 @@ class DrawnBoundingBox : public ConfigOption::Listener, public VideoOverlay::Mou
 };
 
 
-class LabelImages_Widget : public PanelWidget, public ConfigOption::Listener{
+class LabelImages_Widget : public PanelWidget, public ConfigOption::Listener, public VideoSession::StateListener{
 public:
     ~LabelImages_Widget();
     LabelImages_Widget(
@@ -154,14 +160,24 @@ class LabelImages_Widget : public PanelWidget, public ConfigOption::Listener{
         PanelHolder& holder
     );
 
+    // called after loading a new image, clean up all internal data 
+    void clear_for_new_image();
+
     virtual void on_config_value_changed(void* object) override;
 
+    //  Overwrites VideoSession::StateListener::post_startup().
+    virtual void post_startup(VideoSource* source) override;
+
 private:
     LabelImages& m_program;
     ImageAnnotationDisplaySession& m_display_session;
-    ImageAnnotationDisplayWidget* m_switch_widget;
+
+    ImageAnnotationDisplayWidget* m_image_display_widget;
+
     VideoOverlaySet m_overlay_set;
     DrawnBoundingBox m_drawn_box;
+    
+    QLabel* m_embedding_info_label = nullptr;
     ConfigWidget* m_option_widget;
 
     friend class DrawnBoundingBox;

SerialPrograms/Source/ML/UI/ML_ImageAnnotationCommandRow.cpp

@@ -33,10 +33,6 @@ ImageAnnotationCommandRow::ImageAnnotationCommandRow(
 {
     QHBoxLayout* command_row = new QHBoxLayout(this);
     command_row->setContentsMargins(0, 0, 0, 0);
-
-    command_row->addWidget(new QLabel("<b>Console Type:</b>", this), 2);
-    command_row->addSpacing(5);
-
     QHBoxLayout* row = new QHBoxLayout();
     command_row->addLayout(row, 12);
 
@@ -80,12 +76,6 @@ ImageAnnotationCommandRow::ImageAnnotationCommandRow(
 
     row->addSpacing(5);
 
-    m_load_profile_button = new QPushButton("Load Profile", this);
-    row->addWidget(m_load_profile_button, 2);
-
-    m_save_profile_button = new QPushButton("Save Profile", this);
-    row->addWidget(m_save_profile_button, 2);
-
 //    m_test_button = new QPushButton("Test Button", this);
 //    row->addWidget(m_test_button, 3);
 
@@ -132,30 +122,6 @@ ImageAnnotationCommandRow::ImageAnnotationCommandRow(
         this, [this](Qt::CheckState state){ m_session.set_enabled_stats(state == Qt::Checked); }
     );
 #endif
-    connect(
-        m_load_profile_button, &QPushButton::clicked,
-        this, [this](bool) { emit load_profile(); }
-    );
-    connect(
-        m_save_profile_button, &QPushButton::clicked,
-        this, [this](bool) { emit save_profile(); }
-    );
-    
-#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);
-        command_row->addWidget(m_video_button, 2);
-        if (GlobalSettings::instance().STREAM_HISTORY->enabled()){
-            connect(
-                m_video_button, &QPushButton::clicked,
-                this, [this](bool){ emit video_requested(); }
-            );
-        }else{
-            m_video_button->setEnabled(false);
-            m_video_button->setToolTip("Please turn on Stream History to enable video capture.");
-        }
-    }
-#endif
 
     m_session.add_listener(*this);
 }
@@ -177,11 +143,6 @@ void ImageAnnotationCommandRow::set_focus(bool focused){
 }
 
 void ImageAnnotationCommandRow::update_ui(){
-//    cout << "ImageAnnotationCommandRow::update_ui(): focus = " << m_last_known_focus << endl;
-
-    bool stopped = m_last_known_state == ProgramState::STOPPED;
-    m_load_profile_button->setEnabled(stopped);
-
     if (!m_last_known_focus){
         m_status->setText(
             QString::fromStdString(

SerialPrograms/Source/ML/UI/ML_ImageAnnotationCommandRow.h

@@ -39,12 +39,6 @@ class ImageAnnotationCommandRow :
     void on_key_press(const QKeyEvent& key);
     void on_key_release(const QKeyEvent& key);
 
-signals:
-    void load_profile();
-    void save_profile();
-    void screenshot_requested();
-    void video_requested();
-
 public:
     void set_focus(bool focused);
     void update_ui();
@@ -69,9 +63,6 @@ class ImageAnnotationCommandRow :
     QCheckBox* m_overlay_boxes;
     QCheckBox* m_overlay_stats;
 
-    QPushButton* m_load_profile_button;
-    QPushButton* m_save_profile_button;
-    QPushButton* m_video_button;
     bool m_last_known_focus;
     ProgramState m_last_known_state;
 };