From 4f1bf344c45b25e6540bf23f1a4331d2dfbafe07 Mon Sep 17 00:00:00 2001
From: MistEO <mistereo@hotmail.com>
Date: Wed, 11 Oct 2023 17:44:18 +0800
Subject: [PATCH] feat: support GPU OCR

---
 src/MaaCore/Assistant.cpp                    | 19 +++++++++++++++++++
 src/MaaCore/Common/AsstTypes.h               |  2 ++
 src/MaaCore/Config/Miscellaneous/OcrPack.cpp |  3 +++
 src/MaaCore/Config/Miscellaneous/OcrPack.h   |  5 +++++
 4 files changed, 29 insertions(+)

diff --git a/src/MaaCore/Assistant.cpp b/src/MaaCore/Assistant.cpp
index 14e749a487..60ede52477 100644
--- a/src/MaaCore/Assistant.cpp
+++ b/src/MaaCore/Assistant.cpp
@@ -5,6 +5,7 @@
 #include <meojson/json.hpp>
 
 #include "Config/GeneralConfig.h"
+#include "Config/Miscellaneous/OcrPack.h"
 #include "Config/ResourceLoader.h"
 #include "Controller/Controller.h"
 #include "Status.h"
@@ -34,6 +35,24 @@ using namespace asst;
 bool ::AsstExtAPI::set_static_option(StaticOptionKey key, const std::string& value)
 {
     Log.info(__FUNCTION__, "| key", static_cast<int>(key), "value", value);
+
+    switch (key) {
+    case StaticOptionKey::CpuOCR: {
+        WordOcr::get_instance().use_cpu();
+        CharOcr::get_instance().use_cpu();
+        return true;
+    } break;
+    case StaticOptionKey::GpuOCR: {
+        int device_id = std::stoi(value);
+        WordOcr::get_instance().use_gpu(device_id);
+        CharOcr::get_instance().use_gpu(device_id);
+        return true;
+    } break;
+    default:
+        Log.error(__FUNCTION__, "| unknown key:", static_cast<int>(key));
+        break;
+    }
+
     return false;
 }
 
diff --git a/src/MaaCore/Common/AsstTypes.h b/src/MaaCore/Common/AsstTypes.h
index 24df4d0608..1732d69314 100644
--- a/src/MaaCore/Common/AsstTypes.h
+++ b/src/MaaCore/Common/AsstTypes.h
@@ -27,6 +27,8 @@ namespace asst
     enum class StaticOptionKey
     {
         Invalid = 0,
+        CpuOCR = 1,             // use CPU to OCR, no value. It does not support switching after the resource is loaded.
+        GpuOCR = 2,             // use GPU to OCR, value is gpu_id int to string. It does not support switching after the resource is loaded.
     };
 
     enum class InstanceOptionKey
diff --git a/src/MaaCore/Config/Miscellaneous/OcrPack.cpp b/src/MaaCore/Config/Miscellaneous/OcrPack.cpp
index 5ffeebcfd9..060c024e97 100644
--- a/src/MaaCore/Config/Miscellaneous/OcrPack.cpp
+++ b/src/MaaCore/Config/Miscellaneous/OcrPack.cpp
@@ -131,6 +131,9 @@ bool asst::OcrPack::check_and_load()
 
     fastdeploy::RuntimeOption option;
     option.UseOrtBackend();
+    if (m_gpu_id) {
+        option.UseGpu(*m_gpu_id);
+    }
 
     auto det_model = asst::utils::read_file<std::string>(m_det_model_path);
     option.SetModelBuffer(det_model.data(), det_model.size(), nullptr, 0, fastdeploy::ModelFormat::ONNX);
diff --git a/src/MaaCore/Config/Miscellaneous/OcrPack.h b/src/MaaCore/Config/Miscellaneous/OcrPack.h
index 2984f86419..120fab20e0 100644
--- a/src/MaaCore/Config/Miscellaneous/OcrPack.h
+++ b/src/MaaCore/Config/Miscellaneous/OcrPack.h
@@ -4,6 +4,7 @@
 #include "Config/AbstractResource.h"
 
 #include <vector>
+#include <optional>
 
 namespace cv
 {
@@ -38,6 +39,8 @@ namespace asst
         virtual ~OcrPack() override;
 
         virtual bool load(const std::filesystem::path& path) override;
+        void use_cpu() { m_gpu_id = std::nullopt; }
+        void use_gpu(int gpu_id) { m_gpu_id = gpu_id; }
 
         ResultsVec recognize(const cv::Mat& image, bool without_det = false);
 
@@ -53,6 +56,8 @@ namespace asst
         std::filesystem::path m_det_model_path;
         std::filesystem::path m_rec_model_path;
         std::filesystem::path m_rec_label_path;
+
+        std::optional<int> m_gpu_id = std::nullopt;
     };
 
     class WordOcr final : public SingletonHolder<WordOcr>, public OcrPack