Change YOLO detector to use custom model path

Author: Gin

Common/Cpp/PrettyPrint.h

@@ -17,7 +17,9 @@ std::string tostr_padded(size_t digits, uint64_t x);
 std::string tostr_u_commas(int64_t x);
 std::string tostr_bytes(uint64_t bytes);
 
+// Convert double to string using the default precision on ostream.
 std::string tostr_default(double x);
+// Convert double to string with fixed precision.
 std::string tostr_fixed(double x, int precision);
 
 // Format current time to a string to be used as filenames.

SerialPrograms/Source/ML/Inference/ML_YOLOv5Detector.cpp

@@ -6,13 +6,16 @@
 
 #include <filesystem>
 #include <iostream>
+#include <fstream>
 #include <QMessageBox>
 #include <opencv2/imgproc.hpp>
 #include <opencv2/imgcodecs.hpp>
 #include "CommonFramework/ImageTypes/ImageViewRGB32.h"
 #include "CommonFramework/VideoPipeline/VideoOverlay.h"
 #include "CommonFramework/VideoPipeline/VideoOverlayScopes.h"
 #include "CommonFramework/Globals.h"
+#include "Common/Cpp/StringTools.h"
+#include "Common/Cpp/PrettyPrint.h"
 #include "ML_YOLOv5Detector.h"
 
 //#include <iostream>
@@ -25,18 +28,45 @@ namespace ML{
 
 YOLOv5Detector::~YOLOv5Detector() = default;
 
-YOLOv5Detector::YOLOv5Detector()
+
+YOLOv5Detector::YOLOv5Detector(const std::string& model_path)
 {
-    const std::string sam_model_path = RESOURCE_PATH() + "ML/yolov5.onnx";
-    std::vector<std::string> labels = {"Bidoof"};
-    if (std::filesystem::exists(sam_model_path)){
-        m_yolo_session = std::make_unique<YOLOv5Session>(sam_model_path, std::move(labels));
-    } else{
-        std::cerr << "Error: no such YOLOv5 model path " << sam_model_path << "." << std::endl;
+    if (!model_path.ends_with(".onnx")){
+        std::cerr << "Error: wrong model path extension: " << model_path << ". It must be .onnx" << std::endl;
+        QMessageBox box;
+        box.critical(nullptr, "Wrong Model Extension",
+            QString::fromStdString("YOLOv5 model path must end with .onnx. But got " + model_path + "."));
+        return;
+    }
+
+    std::string label_file_path = model_path.substr(0, model_path.size() - 5) + "_label.txt";
+    std::vector<std::string> labels;
+    if (!std::filesystem::exists(label_file_path)){
+        std::cerr << "Error: no such YOLOv5 label file path " << label_file_path << "." << std::endl;
         QMessageBox box;
-        box.critical(nullptr, "YOLOv5 Model Does Not Exist",
-            QString::fromStdString("YOLOv5 model path" + sam_model_path + " does not exist."));
+        box.critical(nullptr, "YOLOv5 Label File Does Not Exist",
+            QString::fromStdString("YOLOv5 label file path " + label_file_path + " does not exist."));
+        return;
     }
+    std::ifstream label_file(label_file_path);
+    if (!label_file.is_open()){
+        std::cerr << "Error: failed to open YOLOv5 label file " << label_file_path << "." << std::endl;
+        QMessageBox box;
+        box.critical(nullptr, "Cannot Open YOLOv5 Label File",
+            QString::fromStdString("YOLOv5 label file " + label_file_path + " cannot be opened."));
+        return;
+    }
+    std::string line;
+    while (std::getline(label_file, line)){
+        line = StringTools::strip(line);
+        if (line.empty() || line[0] == '#'){
+            continue;
+        }
+        labels.push_back(line);
+    }
+    label_file.close();
+
+    m_yolo_session = std::make_unique<YOLOv5Session>(model_path, std::move(labels));
 }
 
 bool YOLOv5Detector::detect(const ImageViewRGB32& screen){
@@ -55,22 +85,24 @@ bool YOLOv5Detector::detect(const ImageViewRGB32& screen){
 }
 
 
-YOLOv5Watcher::YOLOv5Watcher(VideoOverlay& overlay)
+YOLOv5Watcher::YOLOv5Watcher(VideoOverlay& overlay, const std::string& model_path)
     : VisualInferenceCallback("YOLOv5")
     , m_overlay_set(overlay)
+    , m_detector(model_path)
 {
 }
 
 bool YOLOv5Watcher::process_frame(const ImageViewRGB32& frame, WallClock timestamp){
-    if (!m_detector.session()){
+    if (!m_detector.model_loaded()){
         return false;
     }
 
     m_detector.detect(frame);
 
     m_overlay_set.clear();
     for(const auto& box : m_detector.detected_boxes()){
-        m_overlay_set.add(COLOR_RED, box.box, m_detector.session()->label_name(box.label_idx));
+        std::string text = m_detector.session()->label_name(box.label_idx) + ": " + tostr_fixed(box.score, 2);
+        m_overlay_set.add(COLOR_RED, box.box, text);
     }
     return false;
 }

SerialPrograms/Source/ML/Inference/ML_YOLOv5Detector.h

@@ -23,9 +23,17 @@ namespace ML{
 
 class YOLOv5Detector : public StaticScreenDetector{
 public:
-    YOLOv5Detector();
+    // - model_path: path to the onnx model file. The label name file should be the same
+    //   file path and basename and with _label.txt suffix.
+    //   e.g. .../yolo.onnx, .../yolo_label.txt
+    // If model loading fails, no exception is thrown but you can call `model_loaded()` to
+    // check.
+    YOLOv5Detector(const std::string& model_path);
     virtual ~YOLOv5Detector();
 
+    // If it loads the model successfully
+    bool model_loaded() const { return m_yolo_session != nullptr; }
+
     virtual void make_overlays(VideoOverlaySet& items) const override {}
     virtual bool detect(const ImageViewRGB32& screen) override;
 
@@ -42,7 +50,12 @@ class YOLOv5Detector : public StaticScreenDetector{
 
 class YOLOv5Watcher : public VisualInferenceCallback{
 public:
-    YOLOv5Watcher(VideoOverlay& overlay);
+    // - model_path: path to the onnx model file. The label name file should be the same
+    //   file path and basename and with _label.txt suffix.
+    //   e.g. .../yolo.onnx, .../yolo_label.txt
+    // If model loading fails, no exception is thrown but you can call `model_loaded()` to
+    // check.
+    YOLOv5Watcher(VideoOverlay& overlay, const std::string& model_path);
     virtual ~YOLOv5Watcher() {}
 
     virtual void make_overlays(VideoOverlaySet& items) const override {}

SerialPrograms/Source/ML/Models/ML_YOLOv5Model.h

@@ -19,8 +19,12 @@ namespace ML{
 class YOLOv5Session{
 public:
     struct DetectionBox{
+        // Confidence value the model predicts on the detection. Range: [0.0, 1.0].
+        // The higher the value, the more confident the model thinks the prediction is.
         double score;
+        // Bounding box of the detected object.
         ImageFloatBox box;
+        // Object label ID.
         size_t label_idx;
     };
 

SerialPrograms/Source/ML/Programs/ML_RunYOLO.cpp

@@ -37,10 +37,21 @@ RunYOLO_Descriptor::RunYOLO_Descriptor()
 
 
 
-RunYOLO::RunYOLO() {}
+RunYOLO::RunYOLO()
+    : MODEL_PATH(
+        "<b>YOLO Model Path:</b>",
+        LockMode::UNLOCK_WHILE_RUNNING,
+        RESOURCE_PATH() + "ML/yolov5.onnx",
+        "*.onnx",
+        "Path to YOLO .onnx model file"
+    )
+{
+    PA_ADD_OPTION(MODEL_PATH);
+}
 
 void RunYOLO::program(NintendoSwitch::SingleSwitchProgramEnvironment& env, NintendoSwitch::ProControllerContext& context){
-    YOLOv5Watcher watcher(env.console.overlay());
+    std::string model_path = MODEL_PATH;
+    YOLOv5Watcher watcher(env.console.overlay(), model_path);
 
     wait_until(env.console, context, WallClock::max(), {watcher});
 }

SerialPrograms/Source/ML/Programs/ML_RunYOLO.h

@@ -8,6 +8,7 @@
 #ifndef PokemonAutomation_ML_RunYOLO_H
 #define PokemonAutomation_ML_RunYOLO_H
 
+#include "Common/Cpp/Options/PathOption.h"
 #include "NintendoSwitch/NintendoSwitch_SingleSwitchProgram.h"
 
 namespace PokemonAutomation{
@@ -27,6 +28,7 @@ class RunYOLO : public NintendoSwitch::SingleSwitchProgramInstance{
     virtual void program(NintendoSwitch::SingleSwitchProgramEnvironment& env, NintendoSwitch::ProControllerContext& context) override;
 
 private:
+    PathOption MODEL_PATH;
 };