Add code to export annotation to YOLO dataset

Author: Gin

Common/Cpp/StringTools.cpp

@@ -26,7 +26,25 @@ std::string replace(const std::string& str, const std::string& desired, const st
     return ret;;
 }
 
+std::string strip(const std::string& str){
+    size_t first = str.find_first_not_of(" \t\n\r"); // Find first non-whitespace character
+    if (std::string::npos == first) { // If the string is all whitespace or empty
+        return ""; // Return an empty string
+    }
+    size_t last = str.find_last_not_of(" \t\n\r"); // Find last non-whitespace character
+    return str.substr(first, (last - first + 1)); // Extract the trimmed substring
+}
 
+size_t to_size_t(const std::string& str){
+    try {
+        int num = std::stoi(str);
+        return static_cast<size_t>(num);
+    } catch (const std::invalid_argument& e) {
+        return SIZE_MAX;
+    } catch (const std::out_of_range& e) {
+        return SIZE_MAX;
+    }
+}
 
 
 }

Common/Cpp/StringTools.h

@@ -15,7 +15,11 @@ namespace StringTools{
 
 std::string replace(const std::string& str, const std::string& desired, const std::string& replace_with);
 
+// Trim leading and trailing white spaces
+std::string strip(const std::string& str);
 
+// Parse str to size_t. Return SIZE_MAX if parsing fails
+size_t to_size_t(const std::string& str);
 
 }
 }

SerialPrograms/Source/ML/DataLabeling/ML_AnnotationIO.cpp

@@ -7,11 +7,23 @@
 
 #include <fstream>
 #include <iostream>
+#include <map>
 #include <QDirIterator>
 #include <QDir>
+#include <QMessageBox>
 
+#include "Common/Cpp/Json/JsonTools.h"
+#include "Common/Cpp/Json/JsonArray.h"
+#include "Common/Cpp/Json/JsonObject.h"
+#include "Common/Cpp/Json/JsonValue.h"
+#include "Common/Cpp/StringTools.h"
+#include "Common/Cpp/PrettyPrint.h"
 #include "ML_AnnotationIO.h"
 #include "ML_SegmentAnythingModelConstants.h"
+#include "ML_ObjectAnnotation.h"
+
+namespace fs = std::filesystem;
+using std::cout, std::endl;
 
 namespace PokemonAutomation{
 namespace ML{
@@ -77,6 +89,206 @@ std::vector<std::string> find_images_in_folder(const std::string& folder_path, b
     return all_image_paths;
 }
 
+void export_image_annotations_to_yolo_dataset(
+    const std::string& image_folder_path,
+    const std::string& annotation_folder_path,
+    const std::string& yolo_dataset_path
+){
+    const bool recursive = true;
+    const std::vector<std::string>& image_paths = find_images_in_folder(image_folder_path, recursive);
+    if (image_paths.size() == 0){
+        QMessageBox box;
+        box.critical(nullptr, "Empty Image Folder",
+            QString::fromStdString("No images found in " + image_folder_path + "."));
+        return;
+    }
+
+    // TODO for simplicity we will parse this YAML file. In future we should use a proper YAML library
+    std::ifstream fin(yolo_dataset_path.c_str());
+    if (!fin){
+        QMessageBox box;
+        box.critical(nullptr, "Cannot Open YOLO Dataset Config File",
+            QString::fromStdString("Cannot open " + yolo_dataset_path + "."));
+        return;
+    }
+
+    std::vector<std::string> label_names;
+    bool reading_labels = false;
+
+    std::string line;
+    int line_id = 0;
+    while(std::getline(fin, line)){
+        line_id++;
+        // remove "#" comments
+        size_t pound_idx = line.find_first_of("#");
+        if (pound_idx != std::string::npos){
+            line = line.substr(0, pound_idx);
+        }
+        line = StringTools::strip(line);
+        if (line.size() == 0){
+            continue;
+        }
+
+        // cout << "Line: " << line << endl;
+        if (line.starts_with("names:")){
+            reading_labels = true;
+            line.clear();
+            continue;
+        }
+
+        if (reading_labels){
+            // cout << "start reading labels" << endl;
+            size_t colon_idx = line.find_first_of(":");
+            if (colon_idx == std::string::npos){
+                QMessageBox box;
+                box.critical(nullptr, "Error Parsing Dataset YAML",
+                QString::fromStdString("YAML file " + yolo_dataset_path + " line " + std::to_string(line_id) + " contains no colon for labels."));
+                return;
+            }
+            std::string number = StringTools::strip(line.substr(0, colon_idx));
+            std::string label = StringTools::strip(line.substr(colon_idx+1));
+            
+            // cout << "found number " << number << " label " << label << endl;
+            size_t num = StringTools::to_size_t(number);
+            if (num != label_names.size()){
+                QMessageBox box;
+                box.critical(nullptr, "Error Parsing Dataset YAML",
+                QString::fromStdString("YAML file " + yolo_dataset_path + " line " + std::to_string(line_id) + " has no label index."));
+                return;
+            }
+            label_names.push_back(label);
+        }
+
+        line.clear();
+    }
+
+    std::map<std::string, size_t> label_indices;
+    cout << "Load dataset labels: " << endl;
+    for(size_t i = 0; i < label_names.size(); i++){
+        cout << "- " << label_names[i] << endl;
+        label_indices[label_names[i]] = i;
+    }
+    
+
+    // convert images and annotations into new subfolders in the folder of the dataset config
+    auto converted_folder_name = "exported-" + now_to_filestring();
+
+    const auto yolo_dataset_config_file = fs::path(yolo_dataset_path);
+    const fs::path yolo_dataset_folder = yolo_dataset_config_file.parent_path();
+    const auto target_folder = yolo_dataset_folder / converted_folder_name;
+    if (fs::exists(target_folder)){
+        QMessageBox box;
+        box.critical(nullptr, "Export Destination Folder Already Exists",
+            QString::fromStdString("Folder " + target_folder.string() + " already exists."));
+        return;
+    }
+    const auto target_image_folder = target_folder / "images";
+    const auto target_label_folder = target_folder / "labels";
+    cout << "Export to image folder: " << target_image_folder << endl;
+    cout << "Export to label folder: " << target_label_folder << endl;
+
+    fs::create_directories(target_image_folder);
+    fs::create_directories(target_label_folder);
+
+    fs::path anno_folder(annotation_folder_path);
+    for(size_t i = 0; i < image_paths.size(); i++){
+        const auto& image_path = image_paths[i];
+        const auto image_file = fs::path(image_path);
+
+        const std::string anno_filename = image_file.filename().replace_extension(".json").string();
+        fs::path anno_file = anno_folder / anno_filename;
+        if (!fs::exists(anno_file)){
+            QMessageBox box;
+            box.critical(nullptr, "Cannot Find Annotation File",
+                QString::fromStdString("No annotation for " + image_path + "."));
+            return;
+        }
+
+        const auto target_image_file = target_image_folder / image_file.filename();
+        try{
+            fs::copy_file(image_file, target_image_file);
+        } catch (fs::filesystem_error& e)
+        {
+            QMessageBox box;
+            box.critical(nullptr, "Cannot Copy File",
+                QString::fromStdString(
+                    "Cannot copy from " + image_file.string() + " to " + target_image_file.string() + 
+                    ". Probably permission issue, source image is broken or target image path already exists due to image folder having same image filenames"
+                ));
+            return;
+        }
+
+        std::string json_content;
+        const bool anno_loaded = file_to_string(anno_file.string(), json_content);
+        if (!anno_loaded){
+            QMessageBox box;
+            box.warning(nullptr, "Unable to Load Annotation",
+                QString::fromStdString("Cannot open annotation file " + anno_file.string() + ". Probably wrong permission?"));
+            return;
+        }
+
+        const JsonValue loaded_json = parse_json(json_content);
+        const JsonObject* json_obj = loaded_json.to_object();
+        if (!json_obj){
+            QMessageBox box;
+            box.warning(nullptr, "Wrong JSON content",
+                QString::fromStdString("Wong JSON content in annotation file " + anno_file.string() + 
+                ". Probably older annotataion? Try loading and saving this annotation file."));
+            return;
+        }
+
+        std::vector<std::string> label_file_lines;
+        try{
+            const int64_t image_width = json_obj->get_integer_throw("IMAGE_WIDTH");
+            const int64_t image_height = json_obj->get_integer_throw("IMAGE_HEIGHT");
+            const JsonArray& json_array = json_obj->get_array_throw("ANNOTATION");
+            for(size_t i = 0; i < json_array.size(); i++){
+                const ObjectAnnotation anno_obj = ObjectAnnotation::from_json((json_array)[i]);
+                const std::string& label = anno_obj.label;
+
+                auto it = label_indices.find(label);
+                if (it == label_indices.end()){
+                    if (label.ends_with("-male")){
+                        it = label_indices.find(label.substr(0, label.size()-5));
+                    } else if(label.ends_with("-female")){
+                        it = label_indices.find(label.substr(0, label.size()-7));
+                    }
+                }
+                if (it == label_indices.end()){
+                    continue; // label not part of the YOLO dataset. Ignored.
+                }
+
+                const size_t label_id = it->second;
+
+                 // TODO: once we implement the user checkbox on mask reliability, we should change this line
+                const auto& box = anno_obj.mask_box;
+                const double center_x = (box.min_x + box.max_x) / (2.0 * image_width);
+                const double center_y = (box.min_y + box.max_y) / (2.0 * image_height);
+                const double width = box.width() / (double)image_width;
+                const double height = box.height() / (double)image_height;
+                
+                // each row in the YOLO dataclass label file is: class_index x_center y_center width height
+                // https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#12-leverage-models-for-automated-labeling
+                std::ostringstream os;
+                os << label_id << " " << center_x << " " << center_y << " " << width << " " << height;
+                label_file_lines.push_back(os.str());
+            }
+        } catch(JsonParseException& ){
+            QMessageBox box;
+            box.warning(nullptr, "Wrong JSON content",
+                QString::fromStdString("Wong JSON content in annotation file " + anno_file.string() + "."));
+            return;
+        }
+
+        const auto target_label_file = target_label_folder / image_file.filename().replace_extension(".txt");
+        std::ofstream fout(target_label_file.string());
+        for(const auto& line : label_file_lines){
+            fout << line << "\n";
+        }
+    }
+    cout << "Done exporting " << image_paths.size() << " annotations to YOLOv5 dataset" << endl;
+}
+
 
 }
 }

SerialPrograms/Source/ML/DataLabeling/ML_AnnotationIO.h

@@ -25,6 +25,13 @@ void save_image_embedding_to_disk(const std::string& image_filepath, const std::
 // Find image paths stored in a folder. The search can be recursive into child folders or not.
 std::vector<std::string> find_images_in_folder(const std::string& folder_path, bool recursive);
 
+void export_image_annotations_to_yolo_dataset(
+    const std::string& image_folder_path,
+    const std::string& annotation_folder_path,
+    const std::string& yolo_dataset_path
+);
+
+
 }
 }
 

SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp

@@ -118,26 +118,38 @@ void LabelImages::from_json(const JsonValue& json){
     if (file_path){
         load_custom_label_set(*file_path);
     }
+
+    file_path = obj->get_string("YOLO_CONFIG_FILE_PATH");
+    if (file_path){
+        m_yolo_config_file_path = *file_path;
+    }
 }
 JsonValue LabelImages::to_json() const{
     JsonObject obj = std::move(*m_options.to_json().to_object());
     obj["ImageSetup"] = m_display_option.to_json();
     obj["CUSTOM_LABEL_SET_FILE_PATH"] = m_custom_label_set_file_path;
+    obj["YOLO_CONFIG_FILE_PATH"] = m_yolo_config_file_path;
 
     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;
-        for(const auto& anno_obj: m_annotations){
-            anno_json_arr.push_back(anno_obj.to_json());
-        }
-        cout << "Saving annotation to " << m_annotation_file_path << endl;
-        anno_json_arr.dump(m_annotation_file_path);
+    if (m_annotation_file_path.size() == 0 || m_fail_to_load_annotation_file){
+        return;
+    }
+    JsonObject json;
+    json["IMAGE_WIDTH"] = source_image_width;
+    json["IMAGE_HEIGHT"] = source_image_height;
+
+    JsonArray anno_json_arr;
+    for(const auto& anno_obj: m_annotations){
+        anno_json_arr.push_back(anno_obj.to_json());
     }
+    json["ANNOTATION"] = std::move(anno_json_arr);
+
+    cout << "Saving annotation to " << m_annotation_file_path << endl;
+    json.dump(m_annotation_file_path);
 }
 
 void LabelImages::clear_for_new_image(){
@@ -188,7 +200,14 @@ void LabelImages::load_image_related_data(const std::string& image_path, size_t
     }
 
     const JsonValue loaded_json = parse_json(json_content);
-    const JsonArray* json_array = loaded_json.to_array();
+    const JsonObject* json_obj = loaded_json.to_object();
+    const JsonArray* json_array = nullptr;
+    if (json_obj == nullptr){
+        // legacy format, load as an array
+        json_array = loaded_json.to_array();
+    } else{
+        json_array = json_obj->get_array("ANNOTATION");
+    }
     if (json_array == nullptr){
         m_fail_to_load_annotation_file = true;
         QMessageBox box;
@@ -606,17 +625,26 @@ void LabelImages::load_custom_label_set(const std::string& json_path){
     set_selected_label(selected_label());
 }
 
+
 std::pair<size_t, size_t> LabelImages::float_to_pixel(double x, double y) const{
     const size_t px = (size_t)std::max<double>(source_image_width * x + 0.5, 0);
     const size_t py = (size_t)std::max<double>(source_image_height * y + 0.5, 0);
     return std::make_pair(px, py);
 }
 
+
 std::pair<double, double> LabelImages::pixel_to_float(size_t x, size_t y) const{
     return std::make_pair(x / (double)source_image_width, y / (double)source_image_height);
 }
 
 
+void LabelImages::export_to_yolov5_dataset(const std::string& image_folder_path, const std::string& dataset_path){
+    m_yolo_config_file_path = dataset_path;
+
+    export_image_annotations_to_yolo_dataset(image_folder_path, ML_ANNOTATION_PATH(), dataset_path);
+}
+
+
 }
 }
 

SerialPrograms/Source/ML/Programs/ML_LabelImages.h

@@ -104,6 +104,8 @@ class LabelImages : public PanelInstance, public ConfigOption::Listener {
 
     void load_custom_label_set(const std::string& json_path);
 
+    void export_to_yolov5_dataset(const std::string& image_folder_path, const std::string& dataset_path);
+
 private:
     void on_config_value_changed(void* object) override;
 
@@ -175,6 +177,9 @@ class LabelImages : public PanelInstance, public ConfigOption::Listener {
 
     // the file path to load custom label set
     std::string m_custom_label_set_file_path;
+    // the path to get the YOLOv5 YAML config file to export images and finished annotations to
+    // YOLO dataset.
+    std::string m_yolo_config_file_path;
 };