add annotation save load feature

Author: Gin

SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp

@@ -13,13 +13,16 @@
 #include <QScrollArea>
 #include <QPushButton>
 #include <QResizeEvent>
+#include <QMessageBox>
 #include <iostream>
 #include <fstream>
 #include <filesystem>
 #include <cmath>
 #include "CommonFramework/Globals.h"
+#include "Common/Cpp/Json/JsonArray.h"
 #include "Common/Cpp/Json/JsonObject.h"
 #include "Common/Cpp/Json/JsonValue.h"
+#include "Common/Cpp/Json/JsonTools.h"
 #include "Common/Qt/CollapsibleGroupBox.h"
 #include "Pokemon/Resources/Pokemon_PokemonForms.h"
 #include "CommonFramework/VideoPipeline/VideoOverlayScopes.h"
@@ -40,7 +43,66 @@ namespace PokemonAutomation{
 namespace ML{
 
 
-ObjectAnnotation::ObjectAnnotation() {}
+ObjectAnnotation::ObjectAnnotation(): user_box(0,0,0,0), mask_box(0,0,0,0) {}
+
+// if failed to pass, will throw JsonParseException
+ObjectAnnotation json_to_object_annotation(const JsonValue& value){
+    ObjectAnnotation anno_obj;
+
+    const JsonObject& json_obj = value.to_object_throw();
+    const JsonArray& user_box_array = json_obj.get_array_throw("UserBox");
+    anno_obj.user_box = ImagePixelBox(
+        size_t(user_box_array[0].to_integer_throw()),
+        size_t(user_box_array[1].to_integer_throw()),
+        size_t(user_box_array[2].to_integer_throw()),
+        size_t(user_box_array[3].to_integer_throw())
+    );
+    const JsonArray& mask_box_array = json_obj.get_array_throw("MaskBox");
+    anno_obj.mask_box = ImagePixelBox(
+        size_t(mask_box_array[0].to_integer_throw()),
+        size_t(mask_box_array[1].to_integer_throw()),
+        size_t(mask_box_array[2].to_integer_throw()),
+        size_t(mask_box_array[3].to_integer_throw())
+    );
+    size_t mask_width = anno_obj.mask_box.width(), mask_height = anno_obj.mask_box.height();
+    anno_obj.mask.resize(mask_width * mask_height);
+    const JsonArray& mask_values = json_obj.get_array_throw("Mask");
+    for(size_t i = 0; i < anno_obj.mask.size(); i++){
+        anno_obj.mask[i] = bool(mask_values[i].to_integer_throw());
+    }
+
+    anno_obj.label = json_obj.get_string_throw("Label");
+    
+    return anno_obj;
+}
+
+JsonObject object_annotation_to_json(const ObjectAnnotation& object_annotation){
+    JsonObject json_obj;
+    JsonArray user_box_arr;
+    user_box_arr.push_back(int64_t(object_annotation.user_box.min_x));
+    user_box_arr.push_back(int64_t(object_annotation.user_box.min_y));
+    user_box_arr.push_back(int64_t(object_annotation.user_box.max_x));
+    user_box_arr.push_back(int64_t(object_annotation.user_box.max_y));
+    json_obj["UserBox"] = std::move(user_box_arr);
+
+    JsonArray mask_box_arr;
+    mask_box_arr.push_back(int64_t(object_annotation.mask_box.min_x));
+    mask_box_arr.push_back(int64_t(object_annotation.mask_box.min_y));
+    mask_box_arr.push_back(int64_t(object_annotation.mask_box.max_x));
+    mask_box_arr.push_back(int64_t(object_annotation.mask_box.max_y));
+    json_obj["MaskBox"] = std::move(mask_box_arr);
+
+    JsonArray mask_arr;
+    for(size_t i = 0; i < object_annotation.mask.size(); i++){
+        mask_arr.push_back(int64_t(object_annotation.mask[i]));
+    }
+    json_obj["Mask"] = std::move(mask_arr);
+
+    json_obj["Label"] = object_annotation.label;
+
+    return json_obj;
+}
+
 
 DrawnBoundingBox::DrawnBoundingBox(LabelImages_Widget& widget, VideoOverlay& overlay)
     : m_widget(widget)
@@ -62,7 +124,7 @@ DrawnBoundingBox::~DrawnBoundingBox(){
 void DrawnBoundingBox::on_config_value_changed(void* object){
     auto& program = m_widget.m_program;
     std::lock_guard<std::mutex> lg(m_lock);
-    program.set_rendered_objects(m_widget.m_overlay_set);
+    program.update_rendered_objects(m_widget.m_overlay_set);
 }
 void DrawnBoundingBox::on_mouse_press(double x, double y){
     auto& program = m_widget.m_program;
@@ -161,31 +223,99 @@ void LabelImages::from_json(const JsonValue& json){
 JsonValue LabelImages::to_json() const{
     JsonObject obj = std::move(*m_options.to_json().to_object());
     obj["SwitchSetup"] = m_switch_control_option.to_json();
+
+    // m_annotation_file_path
+    if (m_annotation_file_path.size() > 0 && !m_fail_to_load_annotation_file){
+        JsonArray anno_json_arr;
+        for(const auto& anno_obj: m_annotated_objects){
+            anno_json_arr.push_back(object_annotation_to_json(anno_obj));
+        }
+        cout << "Saving annotation to " << m_annotation_file_path << endl;
+        anno_json_arr.dump(m_annotation_file_path);
+    }
     return obj;
 }
 QWidget* LabelImages::make_widget(QWidget& parent, PanelHolder& holder){
     return new LabelImages_Widget(parent, *this, holder);
 }
 
-void LabelImages::set_rendered_objects(VideoOverlaySet& overlay_set){
+void LabelImages::load_image_related_data(const std::string& image_path, size_t source_image_width, size_t source_image_height){
+    this->source_image_height = source_image_height;
+    this->source_image_width = source_image_width;
+
+    m_mask_image = ImageRGB32(source_image_width, source_image_height);
+    cout << "Image source: " << image_path << ", " << source_image_width << " x " << source_image_height << endl;
+    // if no such embedding file, m_iamge_embedding will be empty
+    const bool embedding_loaded = load_image_embedding(image_path, m_image_embedding);
+    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");
+
+    // ensure the folder exists
+    std::filesystem::create_directory(ML_ANNOTATION_PATH());
+    m_annotation_file_path = ML_ANNOTATION_PATH() + anno_filename;
+    if (!std::filesystem::exists(m_annotation_file_path)){
+        cout << "Annotataion output path, " << m_annotation_file_path << " does not exist yet" << endl;
+        return;
+    }
+    std::string json_content;
+    const bool anno_loaded = file_to_string(m_annotation_file_path, json_content);
+    if (!anno_loaded){
+        m_fail_to_load_annotation_file = true;
+        QMessageBox box;
+        box.warning(nullptr, "Unable to Load Annotation",
+            QString::fromStdString("Cannot open annotation file " + m_annotation_file_path + ". Probably wrong permission?"));
+        return;
+    }
+
+    JsonValue loaded_json = parse_json(json_content);
+    const JsonArray* json_array = loaded_json.to_array();
+    if (json_array == nullptr){
+        m_fail_to_load_annotation_file = true;
+        QMessageBox box;
+        box.warning(nullptr, "Unable to Load Annotation",
+            QString::fromStdString("Cannot load annotation file " + m_annotation_file_path + ". Loaded json is not an array"));
+        return;
+    }
+
+    for(size_t i = 0; i < json_array->size(); i++){
+        try{
+            ObjectAnnotation anno_obj = json_to_object_annotation((*json_array)[i]);
+            m_annotated_objects.emplace_back(std::move(anno_obj));
+        } catch(JsonParseException & e){
+            m_fail_to_load_annotation_file = true;
+            QMessageBox box;
+            box.warning(nullptr, "Unable to Load Annotation",
+                QString::fromStdString("Cannot load annotation file " + m_annotation_file_path + 
+                    ". Parsing object " + std::to_string(i) + " failed."));
+        }
+    }
+    m_last_object_idx = m_annotated_objects.size(); 
+    cout << "Loaded existing annotation file " << m_annotation_file_path << endl;
+}
+
+void LabelImages::update_rendered_objects(VideoOverlaySet& overlay_set){
     overlay_set.clear();
     overlay_set.add(COLOR_RED, {X, Y, WIDTH, HEIGHT});
 
-    for(const auto& obj : m_annotated_objects){
+    for(size_t i_obj = 0; i_obj < m_annotated_objects.size(); i_obj++){
+        const auto& obj = m_annotated_objects[i_obj];
         // overlayset.add(COLOR_RED, pixelbox_to_floatbox(source_image_width, source_image_height, obj.user_box));
         const auto mask_float_box = pixelbox_to_floatbox(source_image_width, source_image_height, obj.mask_box);
         std::string label = obj.label;
         const Pokemon::PokemonForm* form = Pokemon::get_pokemon_form(label);
         if (form != nullptr){
             label = form->display_name();
         }
-        overlay_set.add(COLOR_BLUE, mask_float_box, label);
+        Color mask_box_color = (i_obj == m_last_object_idx) ? COLOR_BLACK : COLOR_BLUE;
+        overlay_set.add(mask_box_color, mask_float_box, label);
         size_t mask_width = obj.mask_box.width();
         size_t mask_height = obj.mask_box.height();
         ImageRGB32 mask_image(mask_width, mask_height);
         // cout << "in render, mask_box " << obj.mask_box.min_x << " " << obj.mask_box.min_y << " " << obj.mask_box.max_x << " " << obj.mask_box.max_y << endl;
 
-        // int count = 0;
         for (size_t y = 0; y < mask_height; y++){
             for (size_t x = 0; x < mask_width; x++){
                 const bool mask = obj.mask[y*mask_width + x];
@@ -195,7 +325,6 @@ void LabelImages::set_rendered_objects(VideoOverlaySet& overlay_set){
                 uint32_t color = 0;
                 if (mask){
                     color = (std::abs(int(x) - int(y)) % 4 <= 1) ? combine_argb(150, 30, 144, 255) : combine_argb(150, 0, 0, 60);
-                    // count++;
                 }
                 pixel = color;
             }
@@ -267,9 +396,10 @@ void LabelImages::compute_mask(VideoOverlaySet& overlay_set){
         }
 
         annotation.label = label;
+        m_last_object_idx = m_annotated_objects.size();
         m_annotated_objects.emplace_back(std::move(annotation));
 
-        set_rendered_objects(overlay_set);
+        update_rendered_objects(overlay_set);
     }
 }
 
@@ -309,10 +439,14 @@ LabelImages_Widget::LabelImages_Widget(
     QPushButton* button = new QPushButton("Delete Last Mask", scroll_inner);
     scroll_layout->addWidget(button);
     connect(button, &QPushButton::clicked, this, [this](bool){
-        if (this->m_program.m_annotated_objects.size() > 0){
-            this->m_program.m_annotated_objects.pop_back();
+        auto& program = this->m_program;
+        if (program.m_annotated_objects.size() > 0){
+            program.m_annotated_objects.pop_back();
         }
-        this->m_program.set_rendered_objects(this->m_overlay_set);
+        if (program.m_annotated_objects.size() > 0){
+            program.m_last_object_idx = program.m_annotated_objects.size() - 1;
+        }
+        program.update_rendered_objects(this->m_overlay_set);
     });
 
     m_option_widget = instance.m_options.make_QtWidget(*scroll_inner);
@@ -324,24 +458,18 @@ LabelImages_Widget::LabelImages_Widget(
         const std::string image_path = image_source_desc->path();
         const size_t source_image_height = image_source_desc->source_image_height();
         const size_t source_image_width = image_source_desc->source_image_width();
-        m_program.source_image_height = source_image_height;
-        m_program.source_image_width = source_image_width;
-        m_program.m_mask_image = ImageRGB32(source_image_width, source_image_height);
-        cout << "Image source: " << image_path << ", " << source_image_width << " x " << source_image_height << endl;
-        // if no such embedding file, m_iamge_embedding will be empty
-        load_image_embedding(image_path, m_program.m_image_embedding);
+        m_program.load_image_related_data(image_path, source_image_width, source_image_height);
+        m_program.update_rendered_objects(m_overlay_set);
     }
 
     cout << "LabelImages_Widget built" << endl;
-
-    // TODO: create a custom table to display the annotated bounding boxes
 }
 
 void LabelImages_Widget::on_config_value_changed(void* object){
-    if (m_program.m_annotated_objects.size() > 0){
-        std::string& cur_label = m_program.m_annotated_objects.back().label;
+    if (m_program.m_annotated_objects.size() > 0 && m_program.m_last_object_idx < m_program.m_annotated_objects.size()){
+        std::string& cur_label = m_program.m_annotated_objects[m_program.m_last_object_idx].label;
         cur_label = m_program.FORM_LABEL.slug();
-        m_program.set_rendered_objects(m_overlay_set);
+        m_program.update_rendered_objects(m_overlay_set);
     }
 }
 

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -70,7 +70,9 @@ class LabelImages : public PanelInstance{
     virtual void from_json(const JsonValue& json) override;
     virtual JsonValue to_json() const override;
 
-    void set_rendered_objects(VideoOverlaySet& overlayset);
+    void load_image_related_data(const std::string& image_path, const size_t source_image_width, const size_t source_image_height);
+
+    void update_rendered_objects(VideoOverlaySet& overlayset);
 
     void compute_mask(VideoOverlaySet& overlay_set);
 
@@ -94,9 +96,17 @@ class LabelImages : public PanelInstance{
     std::vector<float> m_image_embedding;
     std::vector<bool> m_output_boolean_mask;
 
+    // buffer to compute SAM mask on
     ImageRGB32 m_mask_image;
+
     SAMSession m_sam_session;
     std::vector<ObjectAnnotation> m_annotated_objects;
+    size_t m_last_object_idx = 0;
+    std::string m_annotation_file_path;
+    // if we find an annotation file that is supposed to be created by user in a previous session, but
+    // we fail to load it, then we shouldn't overwrite this file to possibly erase the previous work.
+    // so this flag is used to denote if we fail to load an annotation file
+    bool m_fail_to_load_annotation_file = false;
 };
 
 
@@ -141,8 +151,6 @@ class LabelImages_Widget : public PanelWidget, public ConfigOption::Listener{
     ConfigWidget* m_option_widget;
 
     friend class DrawnBoundingBox;
-    // std::unique_ptr<ImageRGB32> m_image_mask;
-    // std::unique_ptr<OverlayImage> m_overlay_image;
 };