Go調(diào)用C++動態(tài)庫實現(xiàn)車牌識別的示例代碼

更新時間：2023年12月12日 08:52:00 作者：shelgi

本文主要介紹了如何利用C++中Opencv、TensorRT等庫編譯出動態(tài)庫供Go調(diào)用,再寫個簡單的api對上傳的車輛圖片進(jìn)行車牌識別,文中通過代碼示例給大家介紹的非常詳細(xì),需要的朋友可以參考下

1. 前言

很久沒更新博客，這次正好趁著這次機(jī)會來更新一個稍微有點意思的內(nèi)容，利用C++中Opencv、TensorRT等庫編譯出動態(tài)庫供Go調(diào)用，再寫個簡單的api對上傳的車輛圖片進(jìn)行車牌識別。究其原因，天下苦Java久矣，每次寫JNI去給公司Java后端服務(wù)調(diào)用，而我不喜歡Java那我每次寫好的模型動態(tài)庫就到此為止了？白白浪費(fèi)之前那么多計算資源于心不忍，因此打算收集一些已有模型，做一個自己的模型服務(wù)倉庫。

主要內(nèi)容如下：

模型部分：利用yolov8-pose對車牌數(shù)據(jù)集進(jìn)行訓(xùn)練，然后利用OCR模型對檢測矯正后的車牌字符識別，主要參考這個項目yolov8車牌識別算法，支持12種中文車牌類型

C++部分：實現(xiàn)TensorRT推理以及對應(yīng)模型的前后處理，最后寫cgo對應(yīng)接口以及實現(xiàn)
Go部分：調(diào)用C++編譯后的動態(tài)庫，加載模型，實現(xiàn)輔助功能函數(shù)以及完成接口

2 . 開始

2.1 模型部分

打開上面的鏈接，README中也提到了pytorch1.8以上的版本會有問題，實際嘗試確實如此，總會在一個Conv的地方報錯，魔改了一番代碼還是無法解決，因此下載了車牌檢測的數(shù)據(jù)集本地重新訓(xùn)練模型。需要注意的是，和官方y(tǒng)olov8-pose的輸出結(jié)果中類別數(shù)目不同，因為按照該倉庫的yaml文件設(shè)置會有兩類，因此后處理階段需要注意。

訓(xùn)練參數(shù)等不過多介紹，yolov8文檔十分詳細(xì)可以自己去查看。來看看最后導(dǎo)出的onnx模型。

最后輸出為14*8400，其中14=4+2+8，含義分別是bbox的四個點，對應(yīng)兩個類別概率以及四個關(guān)鍵點的(x,y)坐標(biāo)，后處理階段就要注意對應(yīng)的偏移量分別是4，2，8.

然后OCR模型直接用它提供的預(yù)訓(xùn)練權(quán)重導(dǎo)出就好，精度基本一致。得到onnx之后可以直接利用trtexec轉(zhuǎn)為對應(yīng)的engine文件。

2.2 C++部分

為推理引擎反序列化構(gòu)建，host以及device的內(nèi)存分配等共有操作實現(xiàn)基類，然后重載不同模型的構(gòu)造函數(shù)和前后處理函數(shù)。這個部分可以去參考網(wǎng)上一些開源教程，大多模板一致。在這里有兩個點需要注意：

如果希望兩個模型運(yùn)行在不同顯卡上，記得在所有有關(guān)上下文操作前后加上cudaSetDevice()。
對于不同模型，構(gòu)造函數(shù)傳參大多不一致，目前幾種解決方法：工廠模式輸入modelType對應(yīng)不同實例化，讀取json/yaml等配置文件參數(shù)實例化，最后一種惡心辦法無腦統(tǒng)一實例化接口，大不了某些參數(shù)不用。最優(yōu)方法當(dāng)然是寫配置文件，用yaml-cpp或者其他文件解析庫實現(xiàn)對配置文件參數(shù)解析，然后入?yún)⒕徒y(tǒng)一為配置文件路徑以及一些共有參數(shù)(如deviceId可以在服務(wù)端或者前端設(shè)置因此保留)?？上н@個意見沒被接受，不得已提交的那一版寫的是最惡心的方式，后來改成了第一種通過傳入模型類別去實例化。

稍微說說前后處理部分，對于yolov8-pose之前說了注意偏移量的問題，另外就是對輸出轉(zhuǎn)置處理一下方便解析，當(dāng)然這個操作也可以在模型導(dǎo)出前改一下源碼實現(xiàn)。偏移部分實現(xiàn)大致如下

 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;

然后將所有結(jié)果經(jīng)過nms篩選，得到最終保留結(jié)果。保存目標(biāo)的結(jié)構(gòu)體定義如下：

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代表黑藍(lán)綠白黃五種車牌顏色，78代表78個可識別的字符包括開頭的#號占位符，0-9的數(shù)字，英文字母以及中文漢字，21為最大識別車牌字符長度。然后來看看OCR模型的前后處理，由于大貨車存在雙行車牌的情況，因此需要對車牌上下部分切分然后橫向拼接再給模型推理，大致實現(xiàn)如下：

// 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也就是雙行車牌進(jìn)行拼接操作，當(dāng)然這個是透視變換后的車牌。關(guān)于透視變換可以根據(jù)倉庫中Python代碼翻譯出對應(yīng)的C++版本代碼，

// 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)，和分類任務(wù)后處理一致，找最大值下標(biāo)即為對應(yīng)類別。注意遍歷識別字符時需要過濾操作，即對于下標(biāo)0和已識別出的相鄰?fù)瑯幼址M(jìn)行過濾。找最大值下標(biāo)可以利用std::distance()很方便的找到。

最后就是書寫對應(yīng)的cgo接口，相比起JNI直接根據(jù)類定義使用javah生成的頭文件來寫而言，cgo并沒有生成頭文件的工具，這也讓我們有更多的靈活性去定義對應(yīng)的接口。比如我的接口定義如下：

#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不能調(diào)用c++的類，也不能使用c++的std::string等，所以這里全部是char*。然后實現(xiàn)對應(yīng)接口

#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;
}

這里分別實現(xiàn)了模型實例化，推理以及模型銷毀，最后推理結(jié)果返回的是json格式的字符串，這部分大多還是沿用之前JNI的寫法。最后就是寫個CMakeLists然后編譯，現(xiàn)在來看看C++上的推理結(jié)果圖

對于這種角度的車牌人眼都需要細(xì)看才能識別正確，模型居然也能正確識別，看來模型還是可以的，而且在家里這個服務(wù)器上推理耗時也僅僅1.3ms左右，速度與精度都完全可以接受。

2.3 Go部分

經(jīng)過一系列操作，我們終于編譯得到了.so動態(tài)庫文件，現(xiàn)在就是加載這個動態(tài)庫然后寫個服務(wù)今天的任務(wù)就算完成啦。來看看go調(diào)用動態(tài)庫的部分，首先需要調(diào)用C的庫,并且上面需要添加編譯注釋，同時保證二者之間不能有空行

/*
#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去加載對應(yīng)的動態(tài)庫。剩下的步驟就是根據(jù)剛才C++定義的函數(shù)來對應(yīng)寫Go的實現(xiàn)

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)
}

剩余一些函數(shù)，比如base64，unicode與string的轉(zhuǎn)換，對于推理后json字符串的解析等等略過，最后用gin寫個簡單的POST推理路由以及上傳路由。下面來看看效果：