add model cache update mechanism

Gin · Gin · commit 05277447d339 · 2025-08-25T00:11:16.000-07:00
diff --git a/SerialPrograms/Source/ML/DataLabeling/ML_SegmentAnythingModel.cpp b/SerialPrograms/Source/ML/DataLabeling/ML_SegmentAnythingModel.cpp
@@ -24,8 +24,7 @@ namespace ML{
 
 
 SAMEmbedderSession::SAMEmbedderSession(const std::string& model_path)
-    : session_options(create_session_option("SAMEmbedder"))
-    , session{env, str_to_onnx_str(model_path).c_str(), session_options}
+    : session{create_session(model_path, "SAMEmbedder")}
     , memory_info{Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU)}
     , input_names{session.GetInputNames()}
     , output_names{session.GetOutputNames()}
@@ -64,8 +63,7 @@ void SAMEmbedderSession::run(cv::Mat& input_image, std::vector<float>& model_out
 
 
 SAMSession::SAMSession(const std::string& model_path)
-    : session_options(create_session_option("SAM"))
-    , session{env, str_to_onnx_str(model_path).c_str(), session_options}
+    : session{create_session(model_path, "SAM")}
     , memory_info{Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU)}
     , input_names{session.GetInputNames()}
     , output_names{session.GetOutputNames()}
diff --git a/SerialPrograms/Source/ML/DataLabeling/ML_SegmentAnythingModel.h b/SerialPrograms/Source/ML/DataLabeling/ML_SegmentAnythingModel.h
@@ -36,8 +36,6 @@ class SAMEmbedderSession{
     void run(cv::Mat& input_image, std::vector<float>& output_image_embedding);
     
 private:
-    Ort::Env env;
-    Ort::SessionOptions session_options;
     Ort::Session session;
     Ort::MemoryInfo memory_info;
     Ort::RunOptions run_options;
@@ -71,8 +69,6 @@ class SAMSession{
         const std::vector<int>& input_box,
         std::vector<bool>& output_boolean_mask);
 private:
-    Ort::Env env;
-    Ort::SessionOptions session_options;
     Ort::Session session;
     Ort::MemoryInfo memory_info;
     Ort::RunOptions run_options;
diff --git a/SerialPrograms/Source/ML/Models/ML_ONNXRuntimeHelpers.cpp b/SerialPrograms/Source/ML/Models/ML_ONNXRuntimeHelpers.cpp
@@ -5,25 +5,53 @@
  *  Helper functions to work with ONNX Runtime library
  */
 
+#include <QString>
+#include <QFile>
+#include <QCryptographicHash>
+#include <QByteArray>
+
 #include <iostream>
+#include <string>
+#include <filesystem>
+#include <fstream>
 #include <onnxruntime_cxx_api.h>
+#include "3rdParty/ONNX/OnnxToolsPA.h"
 #include "Common/Compiler.h"
 #include "ML_ONNXRuntimeHelpers.h"
 
+namespace fs = std::filesystem;
+
 namespace PokemonAutomation{
 namespace ML{
 
-Ort::SessionOptions create_session_option(const std::string& cache_folder_name){
-    Ort::SessionOptions so;
+const char* MODEL_CACHE_FOLDER = "./ModelCache/";
+
+// Computes the cryptographic hash of a file.
+std::string create_file_hash(const std::string& filepath){
+    QFile file(QString::fromStdString(filepath));
+    if (!file.open(QIODevice::ReadOnly)) {
+        return "";
+    }
+
+    QCryptographicHash hash(QCryptographicHash::Sha256);
+    if (hash.addData(&file)) {
+        return hash.result().toHex(0).toStdString();
+    } else {
+        return "";
+    }
+}
+
 
+Ort::SessionOptions create_session_options(const std::string& model_cache_path){
+    Ort::SessionOptions so;
+    std::cout << "Set potential model cache path in session options: " << model_cache_path << std::endl;
 #if __APPLE__
     // create session using Apple ML acceleration library CoreML
     std::unordered_map<std::string, std::string> provider_options;
     // See for provider options: https://onnxruntime.ai/docs/execution-providers/CoreML-ExecutionProvider.html
     // "NeuralNetwork" is a faster ModelFormat than "MLProgram".
     provider_options["ModelFormat"] = std::string("NeuralNetwork");
-    // TODO: need to make sure the cache works
-    provider_options["ModelCacheDirectory"] = "./ModelCache/" + cache_folder_name;
+    provider_options["ModelCacheDirectory"] = model_cache_path;
     // provider_options["MLComputeUnits"] = "ALL";
     // provider_options["RequireStaticInputShapes"] = "0";
     // provider_options["EnableOnSubgraphs"] = "0";
@@ -34,6 +62,73 @@ Ort::SessionOptions create_session_option(const std::string& cache_folder_name){
     return so;
 }
 
+// Check the model file cache integrity by checking the existence of a flag file and the model hash stored 
+// in the flag file. If the flag does not exist, we assume the file cache does not exist or is broken.
+// If the hash stored in the flag file does not match the model file, the model file is a new model, delete
+// the old cache.
+// Return
+// - bool: whether to write flag file after cache is created
+// - string: the model file hash to write into the flag file after Ort::Session is built and the cache is created.
+//
+// model_cache_path: Folder path to store model cache. This name is better to be unique for each model for
+//   easier file management.
+// model_path: the model path to load the ML model. This is needed to ensure we delete the old model cache
+//   when a new model
+std::pair<bool, std::string> clean_up_old_model_cache(const std::string& model_cache_path, const std::string& model_path){
+    std::string file_hash = create_file_hash(model_path);
+    if (file_hash.size() == 0){
+        // the model file cannot be loaded
+        return {true, ""};
+    }
+
+    if (!fs::exists(fs::path(model_cache_path))){
+        return {true, file_hash};
+    }
+
+    const std::string flag_file_path = model_cache_path + "/HASH.txt";
+    if (fs::exists(fs::path(flag_file_path))){
+        std::ifstream fin(flag_file_path);
+        if (fin){
+            std::string line;
+            fin >> line;
+            if (line == file_hash){
+                // hash match!
+                return {false, file_hash};
+            }
+        }
+    }
+    // remove everything from model_cache_path
+    fs::remove_all(fs::path(model_cache_path));
+    return {true, file_hash};
+}
+
+
+void write_cache_flag_file(const std::string& model_cache_path, const std::string& hash){
+    if (!fs::exists(fs::path(model_cache_path))){
+        return;
+    }
+    const std::string flag_file_path = model_cache_path + "/HASH.txt";
+    std::ofstream fout(flag_file_path);
+    fout << hash;
+}
+
+
+Ort::Session create_session(const std::string& model_path, const std::string& cache_folder_name){
+    const std::string model_cache_path = MODEL_CACHE_FOLDER + cache_folder_name;
+    Ort::SessionOptions so = create_session_options(model_cache_path);
+    bool write_flag_file = true;
+    std::string file_hash;
+    std::tie(write_flag_file, file_hash) = clean_up_old_model_cache(model_cache_path, model_path);
+    
+    Ort::Env env;
+    Ort::Session session{env, str_to_onnx_str(model_path).c_str(), so};
+    // when Ort::Ssssion is created, if possible, it will create a model cache
+    if (write_flag_file){
+        write_cache_flag_file(model_cache_path, file_hash);
+    }
+    return session;
+}
+
 
 void print_model_input_output_info(const Ort::Session& session){
     const auto input_names = session.GetInputNames();
diff --git a/SerialPrograms/Source/ML/Models/ML_ONNXRuntimeHelpers.h b/SerialPrograms/Source/ML/Models/ML_ONNXRuntimeHelpers.h
@@ -16,14 +16,14 @@
 namespace PokemonAutomation{
 namespace ML{
 
-
-// Create an ONNX Runtime session options object.
+// Create an ONNX Session.
+// cache_folder_name: the folder name in under ./ModelCache/ to store model caches. This name is better
+//   to be unique for each model for easier file management.
+//
 // If on macOS, will use CoreML as the backend.
 // Otherwise, use CPU to run the model.
 // TODO: add Cuda backend for Windows machine.
-// cache_folder_name: the folder name in under ./ModelCache/ to store model caches. This name is better
-//   to be unique for each model for easier file management.
-Ort::SessionOptions create_session_option(const std::string& cache_folder_name);
+Ort::Session create_session(const std::string& model_path, const std::string& cache_folder_name);
 
 // Handy function to create an ONNX Runtime tensor view class from a vector-like `buffer` object holding
 // the tensor data and an array-like `shape` object that represents the dimension of the tensor.
diff --git a/SerialPrograms/Source/ML/Models/ML_YOLOv5Model.cpp b/SerialPrograms/Source/ML/Models/ML_YOLOv5Model.cpp
@@ -58,8 +58,7 @@ std::tuple<int, int, double, double> resize_image_with_border(
 
 YOLOv5Session::YOLOv5Session(const std::string& model_path, std::vector<std::string> label_names)
 : m_label_names(std::move(label_names))
-, m_session_options(create_session_option("YOLOv5"))
-, m_session{m_env, str_to_onnx_str(model_path).c_str(), m_session_options}
+, m_session{create_session(model_path, "YOLOv5")}
 , m_memory_info{Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU)}
 , m_input_names{m_session.GetInputNames()}
 , m_output_names{m_session.GetOutputNames()}
diff --git a/SerialPrograms/Source/ML/Models/ML_YOLOv5Model.h b/SerialPrograms/Source/ML/Models/ML_YOLOv5Model.h
@@ -36,8 +36,6 @@ class YOLOv5Session{
 
     std::vector<std::string> m_label_names;
 
-    Ort::Env m_env;
-    Ort::SessionOptions m_session_options;
     Ort::Session m_session;
     Ort::MemoryInfo m_memory_info;
     Ort::RunOptions m_run_options;
