From 0f46eb0af036861c8cf2db329da7933d5c2e7da1 Mon Sep 17 00:00:00 2001 From: MistEO Date: Fri, 28 Apr 2023 00:41:23 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20OCR=20=E6=83=B0=E6=80=A7=E5=8A=A0?= =?UTF-8?q?=E8=BD=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/MaaCore/Config/Miscellaneous/OcrPack.cpp | 67 +++++++++++++++----- src/MaaCore/Config/Miscellaneous/OcrPack.h | 6 ++ src/MaaCore/Config/ResourceLoader.cpp | 16 ++--- 3 files changed, 61 insertions(+), 28 deletions(-) diff --git a/src/MaaCore/Config/Miscellaneous/OcrPack.cpp b/src/MaaCore/Config/Miscellaneous/OcrPack.cpp index 26c7eacdcf..64b81595de 100644 --- a/src/MaaCore/Config/Miscellaneous/OcrPack.cpp +++ b/src/MaaCore/Config/Miscellaneous/OcrPack.cpp @@ -32,42 +32,42 @@ bool asst::OcrPack::load(const std::filesystem::path& path) LogTraceFunction; Log.info("load", path); - fastdeploy::RuntimeOption option; - option.UseOrtBackend(); - using namespace asst::utils::path_literals; const auto det_dir = path / "det"_p; const auto det_model_file = det_dir / "inference.onnx"_p; - if (std::filesystem::exists(det_model_file)) { - auto det_model = asst::utils::read_file(det_model_file); - option.SetModelBuffer(det_model.data(), det_model.size(), nullptr, 0, fastdeploy::ModelFormat::ONNX); - m_det = std::make_unique("dummy.onnx", std::string(), option, - fastdeploy::ModelFormat::ONNX); + if (std::filesystem::exists(det_model_file) && m_det_model_path != det_model_file) { + m_det_model_path = det_model_file; + m_det = nullptr; } const auto rec_dir = path / "rec"_p; const auto rec_model_file = rec_dir / "inference.onnx"_p; const auto rec_label_file = rec_dir / "keys.txt"_p; - if (std::filesystem::exists(rec_model_file) && std::filesystem::exists(rec_label_file)) { - auto rec_model = asst::utils::read_file(rec_model_file); - std::string rec_label = asst::utils::read_file(rec_label_file); - option.SetModelBuffer(rec_model.data(), rec_model.size(), nullptr, 0, fastdeploy::ModelFormat::ONNX); - m_rec = std::make_unique("dummy.onnx", std::string(), rec_label, option, - fastdeploy::ModelFormat::ONNX); + if (std::filesystem::exists(rec_model_file) && m_rec_model_path != rec_model_file) { + m_rec_model_path = rec_model_file; + m_rec = nullptr; + } + if (std::filesystem::exists(rec_label_file) && m_rec_model_path != rec_label_file) { + m_rec_label_path = rec_label_file; + m_rec = nullptr; } if (m_det && m_rec) { m_ocr = std::make_unique(m_det.get(), m_rec.get()); } - return m_det != nullptr && m_rec != nullptr && m_ocr != nullptr && m_det->Initialized() && m_rec->Initialized() && - m_ocr->Initialized(); + return !m_det_model_path.empty() && !m_rec_model_path.empty() && !m_rec_label_path.empty(); } asst::OcrPack::ResultsVec asst::OcrPack::recognize(const cv::Mat& image, bool without_det) { + if (!check_and_load()) { + Log.error(__FUNCTION__, "check_and_load failed"); + return {}; + } + fastdeploy::vision::OCRResult ocr_result; auto start_time = std::chrono::steady_clock::now(); @@ -121,3 +121,38 @@ asst::OcrPack::ResultsVec asst::OcrPack::recognize(const cv::Mat& image, bool wi Log.trace(class_type, raw_results, without_det ? "by Rec" : "by Pipeline", ", cost", costs, "ms"); return raw_results; } + +bool asst::OcrPack::check_and_load() +{ + if (m_det && m_rec) { + return true; + } + + LogTraceFunction; + + fastdeploy::RuntimeOption option; + option.UseOrtBackend(); + + auto det_model = asst::utils::read_file(m_det_model_path); + option.SetModelBuffer(det_model.data(), det_model.size(), nullptr, 0, fastdeploy::ModelFormat::ONNX); + m_det = std::make_unique("dummy.onnx", std::string(), option, + fastdeploy::ModelFormat::ONNX); + + auto rec_model = asst::utils::read_file(m_rec_model_path); + std::string rec_label = asst::utils::read_file(m_rec_label_path); + option.SetModelBuffer(rec_model.data(), rec_model.size(), nullptr, 0, fastdeploy::ModelFormat::ONNX); + m_rec = std::make_unique("dummy.onnx", std::string(), rec_label, option, + fastdeploy::ModelFormat::ONNX); + + if (m_det && m_rec) { + m_ocr = std::make_unique(m_det.get(), m_rec.get()); + } + + bool det_inited = m_det && m_det->Initialized(); + bool rec_inited = m_rec && m_rec->Initialized(); + bool ocr_inited = m_ocr && m_ocr->Initialized(); + + Log.info("det", det_inited, "rec", rec_inited, "ocr", ocr_inited); + + return det_inited && rec_inited && ocr_inited; +} diff --git a/src/MaaCore/Config/Miscellaneous/OcrPack.h b/src/MaaCore/Config/Miscellaneous/OcrPack.h index 9832110dc7..2984f86419 100644 --- a/src/MaaCore/Config/Miscellaneous/OcrPack.h +++ b/src/MaaCore/Config/Miscellaneous/OcrPack.h @@ -44,9 +44,15 @@ namespace asst protected: OcrPack(); + bool check_and_load(); + std::unique_ptr m_det; std::unique_ptr m_rec; std::unique_ptr m_ocr; + + std::filesystem::path m_det_model_path; + std::filesystem::path m_rec_model_path; + std::filesystem::path m_rec_label_path; }; class WordOcr final : public SingletonHolder, public OcrPack diff --git a/src/MaaCore/Config/ResourceLoader.cpp b/src/MaaCore/Config/ResourceLoader.cpp index 49d11aae28..3a70642f86 100644 --- a/src/MaaCore/Config/ResourceLoader.cpp +++ b/src/MaaCore/Config/ResourceLoader.cpp @@ -63,16 +63,6 @@ bool asst::ResourceLoader::load(const std::filesystem::path& path) LogTraceFunction; using namespace asst::utils::path_literals; - auto word_ocr_future = std::async(std::launch::async, [&]() -> bool { - LoadResourceAndCheckRet(WordOcr, "PaddleOCR"_p); - return true; - }); - - auto char_ocr_future = std::async(std::launch::async, [&]() -> bool { - LoadResourceAndCheckRet(CharOcr, "PaddleCharOCR"_p); - return true; - }); - /* load resource with json files*/ LoadResourceAndCheckRet(GeneralConfig, "config.json"_p); LoadResourceAndCheckRet(RecruitConfig, "recruitment.json"_p); @@ -103,13 +93,15 @@ bool asst::ResourceLoader::load(const std::filesystem::path& path) /* tiles info */ LoadResourceAndCheckRet(TilePack, "Arknights-Tile-Pos"_p); + /* ocr */ + LoadResourceAndCheckRet(WordOcr, "PaddleOCR"_p); + LoadResourceAndCheckRet(CharOcr, "PaddleCharOCR"_p); + #undef LoadTemplByConfigAndCheckRet #undef LoadResourceAndCheckRet #undef LoadCacheWithoutRet m_loaded = true; - m_loaded &= word_ocr_future.get(); - m_loaded &= char_ocr_future.get(); return m_loaded; }