Go 調用 C-- 動態庫實現車牌識別

1. 前言

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

主要內容如下：

• 模型部分：利用 yolov8-pose 對車牌數據集進行訓練，然後利用 OCR 模型對檢測矯正後的車牌字符識別，主要參考這個項目**「yolov8 車牌識別算法，支持 12 種中文車牌類型」**

• 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 之前說了注意偏移量的問題，另外就是對輸出轉置處理一下方便解析，當然這個操作也可以在模型導出前改一下源碼實現。偏移部分實現大致如下

 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 篩選，得到最終保留結果。保存目標的結構體定義如下：

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 模型的前後處理，由於大貨車存在雙行車牌的情況，因此需要對車牌上下部分切分然後橫向拼接再給模型推理，大致實現如下：

// 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++ 版本代碼，

// 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 並沒有生成頭文件的工具，這也讓我們有更多的靈活性去定義對應的接口。比如我的接口定義如下：

#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*。然後實現對應接口

#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的庫, 並且上面需要添加編譯註釋，「同時保證二者之間不能有空行」

/*
#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 的實現

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 也試試，看看到底哪個實現性能最高，再次挖坑。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/A3tvkkJ5RpHfRT6zjK8buQ