Go调用C++动态库实现车牌识别

1. 前言

很久没更新博客,这次正好趁着这次机会来更新一个稍微有点意思的内容,利用C++中Opencv、TensorRT等库编译出动态库供Go调用,再写个简单的api对上传的车辆图片进行车牌识别。究其原因,天下苦Java久矣,每次写JNI去给公司Java后端服务调用,而我不喜欢Java那我每次写好的模型动态库就到此为止了?白白浪费之前那么多计算资源于心不忍,因此打算收集一些已有模型,做一个自己的模型服务仓库。

主要内容如下:

  • C++部分:实现TensorRT推理以及对应模型的前后处理,最后写cgo对应接口以及实现
  • Go部分:调用C++编译后的动态库,加载模型,实现辅助功能函数以及完成接口

2 . 开始

2.1 模型部分

打开上面的链接,README中也提到了pytorch1.8以上的版本会有问题,实际尝试确实如此,总会在一个Conv的地方报错,魔改了一番代码还是无法解决,因此下载了车牌检测的数据集本地重新训练模型。需要注意的是,和官方yolov8-pose的输出结果中类别数目不同,因为按照该仓库的yaml文件设置会有两类,因此后处理阶段需要注意。

训练参数等不过多介绍,yolov8文档十分详细可以自己去查看。来看看最后导出的onnx模型。

最后输出为14*8400,其中14=4+2+8,含义分别是bbox的四个点,对应两个类别概率以及四个关键点的(x,y)坐标,后处理阶段就要注意对应的偏移量分别是4,2,8.

然后OCR模型直接用它提供的预训练权重导出就好,精度基本一致。得到onnx之后可以直接利用trtexec转为对应的engine文件。

2.2 C++部分

为推理引擎反序列化构建,host以及device的内存分配等共有操作实现基类,然后重载不同模型的构造函数和前后处理函数。这个部分可以去参考网上一些开源教程,大多模板一致。在这里有两个点需要注意:

  1. 如果希望两个模型运行在不同显卡上,记得在所有有关上下文操作前后加上cudaSetDevice()
  2. 对于不同模型,构造函数传参大多不一致,目前几种解决方法:工厂模式输入modelType对应不同实例化,读取json/yaml等配置文件参数实例化,最后一种恶心办法无脑统一实例化接口,大不了某些参数不用。最优方法当然是写配置文件,用yaml-cpp或者其他文件解析库实现对配置文件参数解析,然后入参就统一为配置文件路径以及一些共有参数(如deviceId可以在服务端或者前端设置因此保留)。可惜这个意见没被接受,不得已提交的那一版写的是最恶心的方式,后来改成了第一种通过传入模型类别去实例化。

稍微说说前后处理部分,对于yolov8-pose之前说了注意偏移量的问题,另外就是对输出转置处理一下方便解析,当然这个操作也可以在模型导出前改一下源码实现。偏移部分实现大致如下

ini 复制代码
 auto row_ptr    = output.row(i).ptr<float>();
 auto bboxes_ptr = row_ptr;
 auto scores_ptr = row_ptr + 4;
 auto  max_s_ptr  = std::max_element(scores_ptr, scores_ptr + this->class_nums);
 auto kps_ptr    = row_ptr + 6;

然后将所有结果经过nms筛选,得到最终保留结果。保存目标的结构体定义如下:

c 复制代码
struct Object {
    int              label = 0;
    float            prob  = 0.0;
    std::vector<cv::Point2f> kps;
    cv::Rect_<int> rect;
    std::string plateContent;
    std::string colorType;
};

对于OCR模型

模型输入大小为(48,168),输出为5和(21,78),其中5代表黑蓝绿白黄五种车牌颜色,78代表78个可识别的字符包括开头的#号占位符,0-9的数字,英文字母以及中文汉字,21为最大识别车牌字符长度。然后来看看OCR模型的前后处理,由于大货车存在双行车牌的情况,因此需要对车牌上下部分切分然后横向拼接再给模型推理,大致实现如下:

css 复制代码
// merge double plate
void mergePlate(const cv::Mat& src,cv::Mat& dst) {
    int width = src.cols;
    int height = src.rows;
    cv::Mat upper = src(cv::Rect(0,0,width,int(height*5.0/12)));
    cv::Mat lower = src(cv::Rect(0,int(height*1.0/3.),width,height-int(height*1.0/3.0)));
​
    cv::resize(upper,upper,lower.size());
    dst = cv::Mat(lower.rows,lower.cols+upper.cols,CV_8UC3,cv::Scalar(114,114,114));
    upper.copyTo(dst(cv::Rect(0,0,upper.cols,upper.rows)));
    lower.copyTo(dst(cv::Rect(upper.cols,0,lower.cols,lower.rows)));
}
​
​
/*
preprocess
​
0. Perspective
1. merge plate if label is double
2. resize to (48,168)
3. normalize to 0-1 and standard (mean = 0.588 , std = 0.193)
*/
if(obj.label == 1) {
        mergePlate(dst,dst);
}

仅仅对于label为1也就是双行车牌进行拼接操作,当然这个是透视变换后的车牌。关于透视变换可以根据仓库中Python代码翻译出对应的C++版本代码,

css 复制代码
// Perspective
// the kps means pose model's KeyPoints,which is (tl,tr,br,bl)
void Transform(const cv::Mat& src,cv::Mat& dst,const std::vector<cv::Point2f>& kps) {
    float widthA = sqrt(pow((kps[2].x-kps[3].x),2)+pow((kps[2].y-kps[3].y),2));
    float widthB = sqrt(pow((kps[1].x-kps[0].x),2)+pow((kps[1].y-kps[0].y),2));
    float maxWidth = std::max(int(widthA),int(widthB));
​
    float heightA = sqrt(powf((kps[1].x-kps[2].x),2)+powf((kps[1].y-kps[2].y),2));
    float heightB = sqrt(powf((kps[0].x-kps[3].x),2)+powf((kps[0].y-kps[3].y),2));
    float maxHeight = std::max(int(heightA),int(heightB));
​
    std::vector<cv::Point2f> dstTri {
        cv::Point2f(0,0),cv::Point2f(maxWidth,0),
        cv::Point2f(maxWidth,maxHeight),cv::Point2f(0,maxHeight)
    };
    cv::Mat M = cv::getPerspectiveTransform(kps,dstTri);                      cv::warpPerspective(src,dst,M,cv::Size(maxWidth,maxHeight),cv::INTER_LINEAR,cv::BORDER_REPLICATE);
}

Blob部分和Python一样,减去均值除以方差。然后后处理解析部分,0输出的是5维颜色,1输出的是(21,78),和分类任务后处理一致,找最大值下标即为对应类别。注意遍历识别字符时需要过滤操作,即对于下标0和已识别出的相邻同样字符进行过滤。找最大值下标可以利用std::distance()很方便的找到。

最后就是书写对应的cgo接口,相比起JNI直接根据类定义使用javah生成的头文件来写而言,cgo并没有生成头文件的工具,这也让我们有更多的灵活性去定义对应的接口。比如我的接口定义如下:

arduino 复制代码
#include<stdio.h>
#include<string.h>
#ifndef GOWRAP_H
#define GOWRAP_H
#ifdef __cplusplus
extern "C"
{
#endif
extern void* init(const char* modelType, const char* enginePath, int deviceId, int classNums, int kps);
extern char* detect(void* model1,void* model2,const char* base64Img,float score,float iou);
extern void release(void*);
​
#ifdef __cplusplus
}
#endif
​
#endif //GOWRAP_H
​

因为go不能调用c++的类,也不能使用c++的std::string等,所以这里全部是char*。然后实现对应接口

arduino 复制代码
#include "../include/gowrap.h"
#include "../include/plate.hpp"
#include "../include/pose.hpp"
#include "../include/factory.hpp"
#include "../include/base64.h"
void* init(const char* modelType, const char* enginePath, int deviceId, int classNums, int kps) {
    std::string type(modelType);
    std::string engine(enginePath);
    auto model = modelInit(type,engine,deviceId,classNums,kps);
    model->make_pipe(true);
    return (void*)model;
}
​
char* detect(void* m1, void* m2,const char* base64Img, float score, float iou) {
    std::string base64(base64Img);
    cv::Mat image = Base2Mat(base64);
    std::vector<Object> objs;
​
    // get model
    auto* model1 = (YOLOV8_Pose*)m1;
    auto* model2 = (Plate*)m2;
​
    model1->predict(image, objs, score,iou,100);
    model2->predict(image,objs);
​
    // obj trans to json
    Json::Value root;
    Json::Value resObjs;
    Json::Value resObj;
    Json::Value objRec;
    Json::FastWriter writer;
​
    for(const auto&obj : objs){
        Json::Value attrObj;
        attrObj["color"] = obj.colorType;
        attrObj["lineType"] = obj.label;
        attrObj["plate"] = obj.plateContent;
        resObj["attr"] = Json::Value(attrObj);
        resObj["class_id"] = (int)obj.label;
        resObj["conf"] = (float)obj.prob;
        int x = (int)obj.rect.x;
        int y = (int)obj.rect.y;
        int width = (int)obj.rect.width;
        int height = (int)obj.rect.height;
​
        objRec["x"] = x;
        objRec["y"] = y;
        objRec["width"] = width;
        objRec["height"] = height;
​
        resObj["position"]=Json::Value(objRec);
        resObjs.append(resObj);
    }
​
    root["result"] = Json::Value(resObjs);
    std::string resObjs_str = writer.write(root);
    return strdup(resObjs_str.c_str());
}
​
​
void release(void* modelHandle) {
    auto model = (TRTInfer*) modelHandle;
    delete model;
}

这里分别实现了模型实例化,推理以及模型销毁,最后推理结果返回的是json格式的字符串,这部分大多还是沿用之前JNI的写法。最后就是写个CMakeLists然后编译,现在来看看C++上的推理结果图

对于这种角度的车牌人眼都需要细看才能识别正确,模型居然也能正确识别,看来模型还是可以的,而且在家里这个服务器上推理耗时也仅仅1.3ms左右,速度与精度都完全可以接受。

2.3 Go部分

经过一系列操作,我们终于编译得到了.so动态库文件,现在就是加载这个动态库然后写个服务今天的任务就算完成啦。来看看go调用动态库的部分,首先需要调用C的库,并且上面需要添加编译注释,同时保证二者之间不能有空行

shell 复制代码
/*
#cgo LDFLAGS: -L./ -lshelgi_plate -lstdc++
#cgo CPPFLAGS: -I ../include -I /usr/include -I /usr/local/include
#cgo CFLAGS: -std=gnu11
#include<stdio.h>
#include<stdlib.h>
#include "gowrap.h"
*/
import "C"

其实最主要就是第一行LDFLAGS去加载对应的动态库。剩下的步骤就是根据刚才C++定义的函数来对应写Go的实现

go 复制代码
type Object struct {
    p unsafe.Pointer
}
​
func NewModel(modelType, enginePath string, deviceId int, classNums int, kps int) *Object {
    obj := &Object{p: C.init(C.CString(modelType), C.CString(enginePath), C.int(deviceId), C.int(classNums), C.int(kps))}
    return obj
}
​
func detect(m1, m2 *Object, img string, score, iou float32) string {
    res := C.detect(m1.p, m2.p, C.CString(img), C.float(score), C.float(iou))
    result := C.GoString(res)
    return result
}
​
func release(m *Object) {
    C.release(m.p)
}

剩余一些函数,比如base64,unicode与string的转换,对于推理后json字符串的解析等等略过,最后用gin写个简单的POST推理路由以及上传路由。下面来看看效果:

传入图片:

推理结果:

成功识别出两辆车的车牌,响应延时为153ms,经过多次测试,平均在100ms左右,对于单个车辆的图片延时在50ms左右,基本满足需求。

3. 最后

其实这部分内容也是临时想到的,后期打算用Rust也试试,看看到底哪个实现性能最高,再次挖坑。

相关推荐
axxy200025 分钟前
leetcode之hot100---24两两交换链表中的节点(C++)
c++·leetcode·链表
车载诊断技术1 小时前
电子电气架构 --- 什么是EPS?
网络·人工智能·安全·架构·汽车·需求分析
若亦_Royi1 小时前
C++ 的大括号的用法合集
开发语言·c++
KevinRay_1 小时前
Python超能力:高级技巧让你的代码飞起来
网络·人工智能·python·lambda表达式·列表推导式·python高级技巧
跃跃欲试-迪之2 小时前
animatediff 模型网盘分享
人工智能·stable diffusion
慕城南风2 小时前
Go语言中的defer,panic,recover 与错误处理
golang·go
Captain823Jack2 小时前
nlp新词发现——浅析 TF·IDF
人工智能·python·深度学习·神经网络·算法·自然语言处理
被制作时长两年半的个人练习生2 小时前
【AscendC】ReduceSum中指定workLocal大小时如何计算
人工智能·算子开发·ascendc
Captain823Jack3 小时前
w04_nlp大模型训练·中文分词
人工智能·python·深度学习·神经网络·算法·自然语言处理·中文分词
Black_mario3 小时前
链原生 Web3 AI 网络 Chainbase 推出 AVS 主网, 拓展 EigenLayer AVS 应用场景
网络·人工智能·web3