save annotation automatically

Author: Gin

SerialPrograms/Source/ML/DataLabeling/SegmentAnythingModel.cpp

@@ -9,6 +9,7 @@
 #include <QDirIterator>
 #include <fstream>
 #include <iostream>
+#include <QMessageBox>
 #include <onnxruntime_cxx_api.h>
 #include <opencv2/imgcodecs.hpp>
 #include <opencv2/imgproc.hpp>
@@ -249,7 +250,6 @@ bool load_image_embedding(const std::string& image_filepath, std::vector<float>&
 }
 
 
-// TODO: skip already computed embedding
 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()){
@@ -268,13 +268,39 @@ void compute_embeddings_for_folder(const std::string& embedding_model_path, cons
     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();
-        std::cout << ": computing embedding for " << image_path << "..." << std::endl;
+        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;
-        cv::cvtColor(image_bgr, image, cv::COLOR_BGRA2RGB);
+        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(image, output_image_embedding);
+        embedding_session.run(resized_mat, output_image_embedding);
         save_image_embedding_to_disk(image_path, output_image_embedding);
     }
 

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,7 +241,6 @@ 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(){
@@ -265,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();
 
@@ -518,6 +523,7 @@ void LabelImages_Widget::on_config_value_changed(void* object){
 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("");

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -73,6 +73,8 @@ 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();