You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
87 lines
2.3 KiB
C++
87 lines
2.3 KiB
C++
#include "pch.h"
|
|
#include "OCRToolBox.h"
|
|
#include "OCRConfig.h"
|
|
|
|
bool uns::OCRToolBox::CheckFile(const std::wstring& file)
|
|
{
|
|
namespace fs = std::filesystem;
|
|
try
|
|
{
|
|
if (file.empty())
|
|
return false;
|
|
if (!fs::exists(file))
|
|
return false;
|
|
if (!fs::is_regular_file(file))
|
|
return false;
|
|
auto perms = fs::status(file).permissions();
|
|
return ((perms & fs::perms::owner_read) != fs::perms::none && (perms & fs::perms::owner_write) != fs::perms::none);
|
|
}
|
|
catch (...)
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
void uns::OCRToolBox::InitOrtSessionOptions(Ort::SessionOptions& se_opt)
|
|
{
|
|
se_opt.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
|
se_opt.EnableCpuMemArena(); // 启用内存池
|
|
se_opt.EnableMemPattern(); // 启用内存模式优化
|
|
}
|
|
|
|
bool uns::OCRToolBox::AutoSelectEP(const OrtApi* ort, Ort::SessionOptions& se_opt, bool& fallback_to_cpu)
|
|
{
|
|
fallback_to_cpu = false;
|
|
if (G_OCRConfig.GetGPUUsage() == easyocr::GPUUsage::CPUOnly)
|
|
return true;
|
|
try
|
|
{
|
|
OrtStatusPtr status = OrtSessionOptionsAppendExecutionProvider_CUDA(se_opt, 0);
|
|
if (status)
|
|
{
|
|
ort->ReleaseStatus(status);
|
|
if (G_OCRConfig.GetGPUUsage() == easyocr::GPUUsage::ForceGPU)
|
|
return false;
|
|
se_opt = Ort::SessionOptions(); //Reset Session Options (Fallback to CPU)
|
|
fallback_to_cpu = true;
|
|
return true;
|
|
}
|
|
else
|
|
return true; //GPU Enable Successful
|
|
}
|
|
catch (const Ort::Exception&)
|
|
{
|
|
se_opt = Ort::SessionOptions(); //Reset Session Options (Fallback to CPU)
|
|
fallback_to_cpu = true;
|
|
return true;
|
|
}
|
|
catch (...)
|
|
{
|
|
return false;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
void uns::OCRToolBox::GetInputOutputNames(Ort::Session* ort_session, IONames& input_names, IONamesStorage& input_ns, IONames& output_names, IONamesStorage& output_ns)
|
|
{
|
|
input_ns.clear();
|
|
output_ns.clear();
|
|
input_names.clear();
|
|
output_names.clear();
|
|
|
|
if (ort_session == nullptr)
|
|
return;
|
|
|
|
auto input_name_alloc = ort_session->GetInputNameAllocated(0, Ort::AllocatorWithDefaultOptions());
|
|
std::string input_name = input_name_alloc.get();
|
|
input_ns = { input_name };
|
|
for (auto& name : input_ns)
|
|
input_names.push_back(name.c_str());
|
|
|
|
auto output_name_alloc = ort_session->GetOutputNameAllocated(0, Ort::AllocatorWithDefaultOptions());
|
|
std::string output_name = output_name_alloc.get();
|
|
output_ns = { output_name };
|
|
for (auto& name : output_ns)
|
|
output_names.push_back(name.c_str());
|
|
}
|