writing image embedding loading code

Author: Gin

SerialPrograms/CMakeLists.txt

@@ -871,6 +871,8 @@ file(GLOB MAIN_SOURCES
     Source/Kernels/Waterfill/Kernels_Waterfill_Session.h
     Source/Kernels/Waterfill/Kernels_Waterfill_Session.tpp
     Source/Kernels/Waterfill/Kernels_Waterfill_Types.h
+    Source/ML/DataLabeling/SegmentAnythingEmbedding.cpp
+    Source/ML/DataLabeling/SegmentAnythingEmbedding.h
     Source/ML/ML_Panels.cpp
     Source/ML/ML_Panels.h
     Source/ML/Programs/ML_LabelImages.cpp

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingEmbedding.cpp

@@ -0,0 +1,39 @@
+/*  ML Segment Anything Embedding
+ *
+ *  From: https://github.com/PokemonAutomation/
+ *
+ *  Creating and loading image embedding used by Segment Anything Model
+ */
+
+#include <fstream>
+#include <iostream>
+#include "SegmentAnythingEmbedding.h"
+
+
+bool load_image_embedding(const std::string& image_filepath, std::vector<float>& image_embedding){
+    std::string emebdding_path = image_filepath + ".embedding";
+    std::ifstream fin(emebdding_path, std::ios::binary);
+    if (!fin.is_open()){
+        std::cout << "No embedding for image " << image_filepath << std::endl;
+        return false;
+    }
+
+    int embedding_n_channels = 0, embedding_height = 0, emebedding_width = 0;
+    fin.read(reinterpret_cast<char*>(&embedding_n_channels), sizeof(int));
+    fin.read(reinterpret_cast<char*>(&embedding_height), sizeof(int));
+    fin.read(reinterpret_cast<char*>(&emebedding_width), sizeof(int));
+
+    std::cout << "Image embedding shape [" << embedding_n_channels << ", " << embedding_height
+              << ", " << emebedding_width << "]" << std::endl;
+    if (embedding_n_channels <= 0 || embedding_height <= 0 || emebedding_width <= 0){
+        std::string err_msg = "Image embedding wrong dimension from " + emebdding_path;
+        std::cerr << err_msg << std::endl;
+        throw std::runtime_error(err_msg);
+    }
+
+    const int size = embedding_n_channels * embedding_height * emebedding_width;
+    image_embedding.resize(size);
+    fin.read(reinterpret_cast<char*>(image_embedding.data()), sizeof(float) * size);
+    std::cout << "Loaded image embedding from " << emebdding_path << std::endl;
+    return true;
+}

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingEmbedding.h

@@ -0,0 +1,19 @@
+/*  ML Segment Anything Embedding
+ *
+ *  From: https://github.com/PokemonAutomation/
+ *
+ *  Creating and loading image embedding used by Segment Anything Model
+ */
+
+#ifndef PokemonAutomation_ML_SEGMENTANYTHINGEMBEDDING_H
+#define PokemonAutomation_ML_SEGMENTANYTHINGEMBEDDING_H
+
+#include <string>
+#include <vector>
+
+// load pre-computed image embedding from disk
+// return true if there is the embedding file
+bool load_image_embedding(const std::string& image_filepath, std::vector<float>& image_embedding);
+
+
+#endif

SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp

@@ -14,6 +14,8 @@
 #include <QPushButton>
 #include <QResizeEvent>
 #include <iostream>
+#include <fstream>
+#include <filesystem>
 #include "Common/Cpp/Json/JsonObject.h"
 #include "Common/Cpp/Json/JsonValue.h"
 #include "Common/Qt/CollapsibleGroupBox.h"
@@ -24,6 +26,7 @@
 #include "ML_LabelImages.h"
 #include "Pokemon/Pokemon_Strings.h"
 #include "Common/Qt/Options/ConfigWidget.h"
+#include "ML/DataLabeling/SegmentAnythingEmbedding.h"
 
 
 using std::cout;
@@ -52,7 +55,7 @@ DrawnBoundingBox::~DrawnBoundingBox(){
 void DrawnBoundingBox::on_config_value_changed(void* object){
     std::lock_guard<std::mutex> lg(m_lock);
     m_overlay_set.clear();
-    m_overlay_set.add(COLOR_RED, {m_parent.X, m_parent.Y, m_parent.WIDTH, m_parent.HEIGHT});
+    m_overlay_set.add(COLOR_RED, {m_parent.X, m_parent.Y, m_parent.WIDTH, m_parent.HEIGHT}, "Unknown");
 }
 void DrawnBoundingBox::on_mouse_press(double x, double y){
     m_parent.WIDTH.set(0);
@@ -109,6 +112,7 @@ LabelImages_Descriptor::LabelImages_Descriptor()
 {}
 
 
+
 #define ADD_OPTION(x)    m_options.add_option(x, #x)
 
 LabelImages::LabelImages(const LabelImages_Descriptor& descriptor)
@@ -177,20 +181,23 @@ LabelImages_Widget::LabelImages_Widget(
     QPushButton* button = new QPushButton("This is a button", scroll_inner);
     scroll_layout->addWidget(button);
     connect(button, &QPushButton::clicked, this, [&instance](bool){
-        const VideoSourceDescriptor* videoSource = instance.m_switch_control_option.m_video.descriptor().get();
-        auto imageSource = dynamic_cast<const VideoSourceDescriptor_StillImage*>(videoSource);
-        if (imageSource != nullptr){
-            cout << "Image source: " << imageSource->path() << endl;
+        const VideoSourceDescriptor* video_source = instance.m_switch_control_option.m_video.descriptor().get();
+        auto image_source = dynamic_cast<const VideoSourceDescriptor_StillImage*>(video_source);
+        if (image_source != nullptr){
+            cout << "Image source: " << image_source->path() << endl;
         }
     });
 
     m_option_widget = instance.m_options.make_QtWidget(*scroll_inner);
     scroll_layout->addWidget(&m_option_widget->widget());
 
-    const VideoSourceDescriptor* videoSource = instance.m_switch_control_option.m_video.descriptor().get();
-    auto imageSource = dynamic_cast<const VideoSourceDescriptor_StillImage*>(videoSource);
-    if (imageSource != nullptr){
-        cout << "Image source: " << imageSource->path() << endl;
+    const VideoSourceDescriptor* video_source = instance.m_switch_control_option.m_video.descriptor().get();
+    auto image_source = dynamic_cast<const VideoSourceDescriptor_StillImage*>(video_source);
+    if (image_source != nullptr){
+        std::string image_path = image_source->path();
+        cout << "Image source: " << image_path << endl;
+        // if no such embedding file, m_iamge_embedding will be empty
+        load_image_embedding(image_path, m_image_embedding);
     }
     cout << "LabelImages_Widget built" << endl;
 }

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -100,6 +100,7 @@ class LabelImages_Widget : public PanelWidget{
     NintendoSwitch::SwitchSystemWidget* m_switch_widget;
     DrawnBoundingBox m_drawn_box;
     ConfigWidget* m_option_widget;
+    std::vector<float> m_image_embedding;
 };