diff --git a/src/paddleocr-ncnn/paddleocr.cpp b/src/paddleocr-ncnn/paddleocr.cpp index 2cd72dc..c66d7e2 100644 --- a/src/paddleocr-ncnn/paddleocr.cpp +++ b/src/paddleocr-ncnn/paddleocr.cpp @@ -111,6 +111,14 @@ void PaddleOCRApp::resetNet() needReset = false; } +void PaddleOCRApp::clearInferenceAllocators() +{ + detBlobAllocator.clear(); + detWorkspaceAllocator.clear(); + recBlobAllocator.clear(); + recWorkspaceAllocator.clear(); +} + void PaddleOCRApp::initNet() { if(currentPath.empty()) { @@ -162,6 +170,8 @@ void PaddleOCRApp::initNet() detNet = new ncnn::Net; detNet->opt = option; + detNet->opt.blob_allocator = &detBlobAllocator; + detNet->opt.workspace_allocator = &detWorkspaceAllocator; detNet->opt.num_threads = static_cast(maxThreadsUsed); detNet->load_param_bin((detModel + paramSuffix).c_str()); detNet->load_model((detModel + binSuffix).c_str()); @@ -173,6 +183,8 @@ void PaddleOCRApp::initNet() recNet = new ncnn::Net; recNet->opt = option; + recNet->opt.blob_allocator = &recBlobAllocator; + recNet->opt.workspace_allocator = &recWorkspaceAllocator; //由于检测网络的速度足够快,因此GPU设备仅给识别网络使用以节省GPU初始化时间 if (!gpuCanUse.empty()) { @@ -605,6 +617,7 @@ bool PaddleOCRApp::analyze() allResult.clear(); boxesResult.clear(); needBreak = false; + clearInferenceAllocators(); return false; } else { //对识别结果进行最后清理,将未识别到文字的检测框排除掉 @@ -617,7 +630,9 @@ bool PaddleOCRApp::analyze() } } - return !textBoxes.empty(); + const bool result = !textBoxes.empty(); + clearInferenceAllocators(); + return result; } } diff --git a/src/paddleocr-ncnn/paddleocr.h b/src/paddleocr-ncnn/paddleocr.h index ac08a4e..4b711c5 100644 --- a/src/paddleocr-ncnn/paddleocr.h +++ b/src/paddleocr-ncnn/paddleocr.h @@ -30,6 +30,7 @@ #include #include +#include namespace ncnn { class Net; @@ -64,11 +65,16 @@ class PaddleOCRApp : public DeepinOCRPlugin::Plugin std::atomic_bool needBreak = false; ncnn::Net *detNet = nullptr; ncnn::Net *recNet = nullptr; + ncnn::PoolAllocator detBlobAllocator; + ncnn::PoolAllocator detWorkspaceAllocator; + ncnn::PoolAllocator recBlobAllocator; + ncnn::PoolAllocator recWorkspaceAllocator; std::vector keys; PaddleOCR::PostProcessor postProcessor; PaddleOCR::Utility utilityTool; void resetNet(); //重置网络 void initNet(); //初始化网络 + void clearInferenceAllocators(); std::vector>> detect(const cv::Mat &src, float thresh, float boxThresh, float unclipRatio); //检测 std::pair> ctcDecode(const std::vector &recNetOutputData, int h, int w); //CTC解码 void rec(const std::vector &detectImg); //识别