用C++写一个高性能OCR推理服务-TrWebOCR.cpp
SamuraiBUPT 2024-09-16 14:33:11 阅读 99
github: https://github.com/SamuraiBUPT/TrWebOCR.cpp
前因
公司有比较要紧的任务，要在短时间内跑完20万条数据的OCR识别，一方面是生产环境要限制成本，不可能全部使用大模型来跑数据，OCR识别的可信率在60%左右（人工认为的满意度），目前在生产前线，OCR仍然是一个有用的、高效的文字提取方案；另一方面是原本Tornado backend使用的是协程支持，在推理方面的速度对于当时公司的任务以及领导给定的DDL来说并不能完全满足，于是自己写了一套新的serving backend，提高部署Web服务的请求吞吐量。
依赖
Tr: 这份代码被公司其他同事认为是生产环境效果比较好的OCR工具
市面上的Chinese OCR工具很多，但是使用tr是因为稳定性。毕竟在生产环境里面稳定性就是最重要的。httplib：C++网络库，对于开发Pythonic backend的人比较友好，因为设计理念相近，基本上开箱即用。nlohmann/json：json库，处理json数据的，非常方便。opencv：主要用来读取图像文件，尝试了不同的轻量级读取图片库例如stb_image等，发现在读取图片的结果上与pillow、opencv的读取结果又出入。因为底层的tr_run处于黑盒状态，不知道它期待接收的图片数据是什么，所以采用了与原仓库中同样的读取手段。其实对于opencv库，我们用到的仅仅只有读取功能、转化灰度功能。
其实在架构选型之初是准备使用stb_image的，毕竟体量小。在我们这里的功能需求里面，也只有一个读取图像、转化灰度功能，剩下的都不需要了，但是后来发现stb_image读取的图片数据再输入tr的接口中运行，并不能识别出正确的结果，所以最后还是选择使用了opencv。
参考
TrWebOCR
在原版TrWebOCR中，使用的是tornado框架做了推理部署，其实这个代码已经能够支持正常的OCR服务了，并且也相当完善，但是本人非要用C++重写一套的原因，还是推理性能：
tornado使用协程做推理支持，推理throughput并不高，对于我们在这里的紧急任务来说，确实不是一个很好的选择tr既然支持多线程运行，那就说明一定有一个方法是可以提高推理速度、短时间之内满足公司需求的。在GPU运行模式下，有CUDA OOM风险。支持Tr老版本的接口（Tr 2.3），后面没有继续维护了。
综上，我自己做了一个TrWebOCR.cpp，前人栽树，后人乘凉。
TrWebOCR.cpp的特性
高性能OCR Web推理Update with Tr 2.8支持上传图片文件、图片base64编码支持图像旋转CUDA内存管理，规避Out of Memory问题。
说来话长，因为tr这个项目本质上是一个黑盒状态——他只开源了动态链接库<code>.so文件，对于它底层的代码是如何实现的，我们未可知。所以只能基于这个接口在顶层做一系列的优化操作。
设计思路
把tr_run这个函数接口看做是一个task设计多线程-任务池的软件架构。一个ThreadPool持有一个task_queue，利用锁来管理并行程序数据竞争。外部，通过future和promise来获取结果。内部，只管enqueue，并且内部的各个线程会从task_queue这个类内成员变量里面拿取task进行处理在USE_GPU的时候，会再单独开启一个线程，监视显存占用情况，如果有OOM风险，会block住所有请求的运行，清理显存、重启服务，之后再resume各个请求的执行。
为什么会OOM？
这是tr本身的问题。
要是单独使用tr，应该理论上是不会有OOM错误的。在web service部署的场景下，可能是因为tr本身没有做好垃圾回收 （或者显存释放操作），导致随着到来的请求越来越多（图片越来越多），废弃的显存越来越多而没有释放，最终导致OOM的情况。
这对于web service，或者说高吞吐量的推理引擎来说是比较致命的。因为总不可能让一个service在运行过程中直接挂掉了吧？
所以对于这样的情况，有两种解决方案：
准备好容灾措施，多开两个服务，哪个服务挂了重启哪个手动管理CUDA Memory，做好清理、重启操作
在这里我们是使用的CPP进行编写service，追求high-throughput，那我们就选择第二个方案，手动管理内存。
性能比较
1. GPU利用率比较
GPU：cpp后端：
+-----------------------------------------------------------------------------------------+code>
| NVIDIA-SMI 550.90.07 Driver Version: 550.90.07 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 2080 Ti On | 00000000:B2:00.0 Off | N/A |
| 33% 45C P2 135W / 250W | 2673MiB / 11264MiB | 73% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+
GPU：tornado后端
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.90.07 Driver Version: 550.90.07 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 2080 Ti On | 00000000:B2:00.0 Off | N/A |
| 33% 45C P2 135W / 250W | 875MiB / 11264MiB | 39% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+
在这里可以明显看出，在并发20个请求的时候，Tornado在GPU模式里面，无论是显存占用还是GPU利用率都没有CPP 后端高，个人认为可能是Tornado使用的协程，没法真正做到并行推理。这也能解释为什么cpp做了多线程处理之后要快一点。
2. 吞吐量
这就是cpp后端带来的性能提升。
更多内容，可以移步至github，目前代码已全部开源：
github: https://github.com/SamuraiBUPT/TrWebOCR.cpp
后续-web service开发
使用httplib做了基本的serving开发，后面可能会选择更利于高性能情况下的网络库，目前是为了图敏捷开发而选择的。
1. 核心接口定义
<code>svr.Post("/api/trocr", [&tr_task_pool, ctpn_id, crnn_id, &rotations](const Request& req, Response& res) { -- -->code>
try {
#ifdef USE_GPU
while (isBlock.load()) {
std::this_thread::sleep_for(std::chrono::milliseconds(10)); // 自旋等待block解除
}
#endif
int width, height, channels;
cv::Mat img;
if (req.has_file("file")) {
const auto& file = req.get_file_value("file");
// 从内存中读取图像
std::vector<unsigned char> img_data(file.content.begin(), file.content.end());
img = cv::imdecode(img_data, cv::IMREAD_COLOR); // 读取为彩色图像
} else if (req.has_param("img_base64")) {
std::string img_base64 = req.get_param_value("img_base64");
std::string decoded_data = base64_decode(img_base64); // 解码Base64数据
// 将Base64解码的数据转为vector并加载为彩色图像
std::vector<unsigned char> img_data(decoded_data.begin(), decoded_data.end());
img = cv::imdecode(img_data, cv::IMREAD_COLOR); // 读取为彩色图像
} else {
res.set_content("Missing file or img_base64 parameter", "text/plain");
return;
}
if (img.empty()) {
res.set_content("Failed to load image", "text/plain");
return;
}
// 转为灰度图像
cv::Mat gray_img;
cv::cvtColor(img, gray_img, cv::COLOR_BGR2GRAY);
// 现在开始核心的推理进程，旋转图像来看看到底是不是有用的数据
auto start = std::chrono::high_resolution_clock::now();
std::string plain_text;
plain_text = "";
width = gray_img.cols;
height = gray_img.rows;
channels = gray_img.channels();
unsigned char* image_data = gray_img.data;
auto future = tr_task_pool.enqueue(image_data, height, width, channels, ctpn_id, crnn_id);
std::vector<TrResult> results = future.get(); // inference
// 处理 OCR 结果
for (const auto& result : results) {
std::string txt = std::get<1>(result);
plain_text += txt + "|";
}
res.set_content(plain_text, "text/plain; charset=UTF-8");
// log
auto end = std::chrono::high_resolution_clock::now();
auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
std::string msg = std::string("Latency: ") + std::to_string(duration.count() / 1000.0) + " ms";
std::string level("INFO");
tr_log(msg, level);
} catch (const std::exception& e) {
res.set_content("Invalid JSON data", "text/plain");
}
});
这个就是最主要的推理接口了，
接收图片数据，支持图片文件、base64编码两种形式的数据上传图片读取、灰度图转化启动推理，通过future获得结果。
2. 中文识别支持（C++方面处理）
为了支持中文内容的识别、输出，做中文支持也是必不可少的内容：
std::string unichr_utf8(int unicode) {
std::string result;
if (unicode <= 0x7F) {
result += static_cast<char>(unicode);
} else if (unicode <= 0x7FF) {
result += static_cast<char>(0xC0 | ((unicode >> 6) & 0x1F));
result += static_cast<char>(0x80 | (unicode & 0x3F));
} else if (unicode <= 0xFFFF) {
result += static_cast<char>(0xE0 | ((unicode >> 12) & 0x0F));
result += static_cast<char>(0x80 | ((unicode >> 6) & 0x3F));
result += static_cast<char>(0x80 | (unicode & 0x3F));
} else if (unicode <= 0x10FFFF) {
result += static_cast<char>(0xF0 | ((unicode >> 18) & 0x07));
result += static_cast<char>(0x80 | ((unicode >> 12) & 0x3F));
result += static_cast<char>(0x80 | ((unicode >> 6) & 0x3F));
result += static_cast<char>(0x80 | (unicode & 0x3F));
}
return result;
}
std::pair<std::string, float> parse(const int* unicode_arr, const float* prob_arr, int num) {
std::string txt;
float prob = 0.0;
int unicode_pre = -1;
int count = 0;
for (int pos = 0; pos < num; ++pos) {
int unicode = unicode_arr[pos];
if (unicode >= 0) {
if (unicode != unicode_pre) {
txt += unichr_utf8(unicode); // 将 Unicode 转换为 UTF-8 编码
}
count += 1;
prob += prob_arr[pos];
}
unicode_pre = unicode;
}
float confidence = prob / std::max(count, 1); // 计算平均置信度
return { txt, confidence};
}
我们的网络库是支持中英文OCR识别的，效果都还挺不错。
代码已全部开源至github：
github: https://github.com/SamuraiBUPT/TrWebOCR.cpp
效果图展示
1. 中英文效果展示
2. GPU模式下、规避OOM操作服务器日志
github: https://github.com/SamuraiBUPT/TrWebOCR.cpp
代码已全部开源，欢迎使用。
声明
本文内容仅代表作者观点，或转载于其他网站，本站不以此文作为商业用途
如有涉及侵权，请联系本站进行删除
转载本站原创文章，请注明来源及作者。