1、识别监测代码独立出来；

CMakeLists.txt

@@ -0,0 +1,106 @@
+set(SON_PROJECT_NAME yoloseg)
+
+message("========== Load son project ${SON_PROJECT_NAME} ==========" )
+
+unset(OPTION_ENABLE_TENSORRT_DETECT_CODE CACHE)
+option(OPTION_ENABLE_TENSORRT_DETECT_CODE "" OFF)
+message("<=${SON_PROJECT_NAME}=> OPTION_ENABLE_TENSORRT_DETECT_CODE: " ${OPTION_ENABLE_TENSORRT_DETECT_CODE})
+
+if (OPTION_ENABLE_TENSORRT_DETECT_CODE)
+    set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86)
+    set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
+    project(${SON_PROJECT_NAME} LANGUAGES CXX CUDA)
+
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++17 -O3")
+    set(CMAKE_CXX_STANDARD 17)
+    option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
+
+    # CUDA
+    find_package(CUDA REQUIRED)
+    message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n")
+    get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY)
+    message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n")
+
+    # OpenCV
+    message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n")
+    message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n")
+    message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n")
+
+    if(${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0)
+        message(STATUS "Build with -DBATCHED_NMS")
+        add_definitions(-DBATCHED_NMS)
+    endif()
+
+    # TensorRT
+    set(TensorRT_INCLUDE_DIRS /usr/include/aarch64-linux-gnu)
+    set(TensorRT_LIBRARIES /usr/lib/aarch64-linux-gnu)
+    message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n")
+    message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n")
+
+    list(APPEND INCLUDE_DIRS
+            ${CUDA_INCLUDE_DIRS}
+            ${TensorRT_INCLUDE_DIRS}
+    )
+
+    list(APPEND ALL_LIBS
+            nvinfer
+            nvinfer_plugin
+            ${CUDA_LIBRARIES}
+            ${CUDA_LIB_DIR}
+            ${TensorRT_LIBRARIES}
+    )
+
+    aux_source_directory(${CMAKE_CURRENT_LIST_DIR}/tensorrt_detect yolov8_detect)
+else ()
+    aux_source_directory(${CMAKE_CURRENT_LIST_DIR}/onnx_detect yolov8_detect)
+endif ()
+
+find_package(gRPC CONFIG REQUIRED)
+# 获取系统架构
+message("SYSTEM: ${CMAKE_SYSTEM_PROCESSOR}")
+
+aux_source_directory(${CMAKE_CURRENT_LIST_DIR}/proto proto)
+
+include_directories(
+        ${INCLUDE_DIRS}
+        ${CMAKE_CURRENT_LIST_DIR}
+        ${OpenCV_INCLUDE_DIRS}
+        ${PROTOBUF_INCLUDE_DIR}
+        ${NebulaSDK_PATH}/include
+)
+
+aux_source_directory(${CMAKE_CURRENT_LIST_DIR}/detect detect)
+
+set(SON_PROJECT_SOURCE_LIST ${SON_PROJECT_SOURCE_LIST}
+        ${proto}
+        ${yolov8_detect}
+        ${CMAKE_CURRENT_LIST_DIR}/segMain.cpp
+        ${CMAKE_CURRENT_LIST_DIR}/net_message_trans.h
+        ${CMAKE_CURRENT_LIST_DIR}/net_message_trans.cpp
+        ${CMAKE_CURRENT_LIST_DIR}/yolo_seg_mqtt_async.h
+        ${CMAKE_CURRENT_LIST_DIR}/yolo_seg_mqtt_async.cpp
+)
+
+set(SON_PROJECT_DEPEND_LIST ${SON_PROJECT_DEPEND_LIST}
+        zxpahoc
+        zxthread
+        protobuf::libprotobuf
+        gRPC::grpc++_reflection
+        ${OpenCV_LIBRARIES}
+        ${ALL_LIBS}
+)
+
+add_executable(${SON_PROJECT_NAME} ${SON_PROJECT_SOURCE_LIST})
+target_link_libraries(${SON_PROJECT_NAME} ${SON_PROJECT_DEPEND_LIST})
+target_compile_definitions(${SON_PROJECT_NAME} PRIVATE PROJECT_NAME="${SON_PROJECT_NAME}")
+if (OPTION_ENABLE_TENSORRT_DETECT_CODE)
+        target_compile_definitions(${SON_PROJECT_NAME} PRIVATE ENABLE_TENSORRT_DETECT)
+endif ()
+
+install(TARGETS ${SON_PROJECT_NAME}
+        LIBRARY DESTINATION lib  # 动态库安装路径
+        ARCHIVE DESTINATION lib  # 静态库安装路径
+        RUNTIME DESTINATION bin  # 可执行文件安装路径
+        PUBLIC_HEADER DESTINATION include  # 头文件安装路径
+)
+

net_message_trans.h

@@ -6,18 +6,33 @@
 
 class NetMessageTrans {
 public:
-    static void lableMat2Proto(int lable, cv::Mat &mat, JetStream::LabelImage &image) {
+    static void lableMat2Proto(int lable, cv::Mat &mat, JetStream::LabelImage &image, bool encode = false) {
         image.set_label(lable);
         image.mutable_ir()->set_width(mat.cols);
         image.mutable_ir()->set_height(mat.rows);
         image.mutable_ir()->set_channel(mat.type());
-        image.mutable_ir()->set_data(mat.data, mat.dataend - mat.datastart);
+        image.mutable_ir()->set_encode(encode);
+        if (encode) {
+            std::vector<uchar> merge_encode_data;
+//            cv::cvtColor(mat, mat, cv::COLOR_GRAY2RGB);
+            cv::imencode(".jpg", mat, merge_encode_data, {cv::IMWRITE_JPEG_QUALITY, 80});
+            image.mutable_ir()->set_data(merge_encode_data.data(), merge_encode_data.size());
+        } else {
+            image.mutable_ir()->set_data(mat.data, mat.dataend - mat.datastart);
+        }
         // LOG(INFO) << "mat.dataend - mat.datastart = " << mat.dataend - mat.datastart;
     }
 
     static void Proto2lableMat(JetStream::LabelImage &image, int &lable, cv::Mat &mat) {
         lable = image.label();
-        mat = cv::Mat(image.ir().height(), image.ir().width(), image.ir().channel(), (void*)image.ir().data().data());
+        if (image.ir().encode()) {
+            std::vector<uchar> data;
+            data.resize(image.ir().data().size());
+            memcpy((void *)data.data(), (void *)image.ir().data().data(), data.size());
+            mat = cv::imdecode(data, 0);
+        } else {
+            mat = cv::Mat(image.ir().height(), image.ir().width(), image.ir().channel(), (void*)image.ir().data().data());
+        }
     }
 
     static void getSegBoxData() {

onnx_detect/inference.cpp

@@ -0,0 +1,183 @@
+#include "inference.h"
+
+Inference::Inference(const std::string &onnxModelPath, const cv::Size &modelInputShape, const std::string &classesTxtFile, const bool &runWithCuda)
+{
+    modelPath = onnxModelPath;
+    modelShape = modelInputShape;
+    classesPath = classesTxtFile;
+    cudaEnabled = runWithCuda;
+
+    loadOnnxNetwork();
+    // loadClassesFromFile(); The classes are hard-coded for this example
+}
+
+std::vector<Detection> Inference::runInference(const cv::Mat &input)
+{
+    cv::Mat modelInput = input;
+    if (letterBoxForSquare && modelShape.width == modelShape.height)
+        modelInput = formatToSquare(modelInput);
+
+    cv::Mat blob;
+    cv::dnn::blobFromImage(modelInput, blob, 1.0/255.0, modelShape, cv::Scalar(), true, false);
+    net.setInput(blob);
+
+    std::vector<cv::Mat> outputs;
+    net.forward(outputs, net.getUnconnectedOutLayersNames());
+
+    int rows = outputs[0].size[1];
+    int dimensions = outputs[0].size[2];
+
+    bool yolov8 = false;
+    // yolov5 has an output of shape (batchSize, 25200, 85) (Num classes + box[x,y,w,h] + confidence[c])
+    // yolov8 has an output of shape (batchSize, 84,  8400) (Num classes + box[x,y,w,h])
+    if (dimensions > rows) // Check if the shape[2] is more than shape[1] (yolov8)
+    {
+        yolov8 = true;
+        rows = outputs[0].size[2];
+        dimensions = outputs[0].size[1];
+
+        outputs[0] = outputs[0].reshape(1, dimensions);
+        cv::transpose(outputs[0], outputs[0]);
+    }
+    float *data = (float *)outputs[0].data;
+
+    float x_factor = modelInput.cols / modelShape.width;
+    float y_factor = modelInput.rows / modelShape.height;
+
+    std::vector<int> class_ids;
+    std::vector<float> confidences;
+    std::vector<cv::Rect> boxes;
+
+    for (int i = 0; i < rows; ++i)
+    {
+        if (yolov8)
+        {
+            float *classes_scores = data+4;
+
+            cv::Mat scores(1, classes.size(), CV_32FC1, classes_scores);
+            cv::Point class_id;
+            double maxClassScore;
+
+            minMaxLoc(scores, 0, &maxClassScore, 0, &class_id);
+
+            if (maxClassScore > modelScoreThreshold)
+            {
+                confidences.push_back(maxClassScore);
+                class_ids.push_back(class_id.x);
+
+                float x = data[0];
+                float y = data[1];
+                float w = data[2];
+                float h = data[3];
+
+                int left = int((x - 0.5 * w) * x_factor);
+                int top = int((y - 0.5 * h) * y_factor);
+
+                int width = int(w * x_factor);
+                int height = int(h * y_factor);
+
+                boxes.push_back(cv::Rect(left, top, width, height));
+            }
+        }
+        else // yolov5
+        {
+            float confidence = data[4];
+
+            if (confidence >= modelConfidenceThreshold)
+            {
+                float *classes_scores = data+5;
+
+                cv::Mat scores(1, classes.size(), CV_32FC1, classes_scores);
+                cv::Point class_id;
+                double max_class_score;
+
+                minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
+
+                if (max_class_score > modelScoreThreshold)
+                {
+                    confidences.push_back(confidence);
+                    class_ids.push_back(class_id.x);
+
+                    float x = data[0];
+                    float y = data[1];
+                    float w = data[2];
+                    float h = data[3];
+
+                    int left = int((x - 0.5 * w) * x_factor);
+                    int top = int((y - 0.5 * h) * y_factor);
+
+                    int width = int(w * x_factor);
+                    int height = int(h * y_factor);
+
+                    boxes.push_back(cv::Rect(left, top, width, height));
+                }
+            }
+        }
+
+        data += dimensions;
+    }
+
+    std::vector<int> nms_result;
+    cv::dnn::NMSBoxes(boxes, confidences, modelScoreThreshold, modelNMSThreshold, nms_result);
+
+    std::vector<Detection> detections{};
+    for (unsigned long i = 0; i < nms_result.size(); ++i)
+    {
+        int idx = nms_result[i];
+
+        Detection result;
+        result.class_id = class_ids[idx];
+        result.confidence = confidences[idx];
+
+        std::random_device rd;
+        std::mt19937 gen(rd());
+        std::uniform_int_distribution<int> dis(100, 255);
+        result.color = cv::Scalar(dis(gen),
+                                  dis(gen),
+                                  dis(gen));
+
+        result.className = classes[result.class_id];
+        result.box = boxes[idx];
+
+        detections.push_back(result);
+    }
+
+    return detections;
+}
+
+void Inference::loadClassesFromFile()
+{
+    std::ifstream inputFile(classesPath);
+    if (inputFile.is_open())
+    {
+        std::string classLine;
+        while (std::getline(inputFile, classLine))
+            classes.push_back(classLine);
+        inputFile.close();
+    }
+}
+
+void Inference::loadOnnxNetwork()
+{
+    net = cv::dnn::readNetFromONNX(modelPath);
+    if (cudaEnabled)
+    {
+        net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
+        net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
+    }
+    else
+    {
+        net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
+        net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
+    }
+}
+
+cv::Mat Inference::formatToSquare(const cv::Mat &source)
+{
+    int col = source.cols;
+    int row = source.rows;
+    int _max = MAX(col, row);
+    cv::Mat result = cv::Mat::zeros(_max, _max, CV_8UC3);
+    source.copyTo(result(cv::Rect(0, 0, col, row)));
+    return result;
+}

onnx_detect/inference.h

@@ -0,0 +1,50 @@
+#pragma once
+
+// Cpp native
+#include <fstream>
+#include <vector>
+#include <string>
+#include <random>
+
+// OpenCV / DNN / Inference
+#include <opencv2/imgproc.hpp>
+#include <opencv2/opencv.hpp>
+#include <opencv2/dnn.hpp>
+
+struct Detection
+{
+    int class_id{0};
+    std::string className{};
+    float confidence{0.0};
+    cv::Scalar color{};
+    cv::Rect box{};
+};
+
+class Inference
+{
+public:
+    Inference(const std::string &onnxModelPath, const cv::Size &modelInputShape = {640, 640}, const std::string &classesTxtFile = "", const bool &runWithCuda = false);
+    std::vector<Detection> runInference(const cv::Mat &input);
+
+private:
+    void loadClassesFromFile();
+    void loadOnnxNetwork();
+    cv::Mat formatToSquare(const cv::Mat &source);
+
+    std::string modelPath{};
+    std::string classesPath{};
+    bool cudaEnabled{};
+
+    std::vector<std::string> classes{"car", "wheel"};
+
+    cv::Size2f modelShape{};
+
+    float modelConfidenceThreshold {0.25};
+    float modelScoreThreshold      {0.45};
+    float modelNMSThreshold        {0.50};
+
+    bool letterBoxForSquare = true;
+
+    cv::dnn::Net net;
+};
+

onnx_detect/wheel-detector.cpp

@@ -0,0 +1,66 @@
+#include "wheel-detector.h"
+
+TensorrtWheelDetector::TensorrtWheelDetector(const std::string &model_file, const std::string &class_file){
+    yolov8_ = new Inference(model_file, cv::Size{640, 640}, class_file, false);
+}
+
+TensorrtWheelDetector::~TensorrtWheelDetector(){
+}
+
+bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs){
+    std::vector<Detection> rets = yolov8_->runInference(img);
+
+    for (auto &ret: rets) {
+        Object obj;
+       obj.rect = ret.box;
+        obj.label = ret.class_id;
+        obj.prob = ret.confidence;
+        objs.push_back(obj);
+    }
+    return true;
+}
+
+bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res){
+    std::vector<Detection> rets = yolov8_->runInference(img);
+
+    for (auto &ret: rets) {
+        Object obj;
+        obj.rect = ret.box;
+        obj.label = ret.class_id;
+        obj.prob = ret.confidence;
+        obj.boxMask = img;
+        objs.push_back(obj);
+
+        cv::Rect box = ret.box;
+        cv::Scalar color = ret.color;
+
+        // Detection box
+        cv::rectangle(res, box, color, 2);
+
+        // Detection box text
+        std::string classString = ret.className + ' ' + std::to_string(ret.confidence).substr(0, 4);
+        cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
+        cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);
+
+        cv::rectangle(res, textBox, color, cv::FILLED);
+        cv::putText(res, classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
+    }
+    return true;
+}
+
+std::vector<cv::Point> TensorrtWheelDetector::getPointsFromObj(const Object &obj){
+    std::vector<cv::Point> ret;
+
+    int x=int(obj.rect.x+0.5);
+    int y=int(obj.rect.y+0.5);
+    int width=int(obj.rect.width);
+    int height=int(obj.rect.height);
+
+    for(int i=0;i<height;++i){
+        for(int j=0;j<width;++j){
+            ret.emplace_back(x+j,y+i);
+        }
+    }
+
+    return ret;
+}

onnx_detect/wheel-detector.h

@@ -0,0 +1,29 @@
+#pragma once
+
+#include "inference.h"
+
+struct Object {
+    cv::Rect_<float> rect;
+    int              label = 0;
+    float            prob  = 0.0;
+    cv::Mat          boxMask;
+};
+    
+class TensorrtWheelDetector{
+public:
+    TensorrtWheelDetector(const std::string &model_file, const std::string &class_file);
+    ~TensorrtWheelDetector();
+
+    bool detect(cv::Mat& img,std::vector<Object>& objs);
+    bool detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res);
+
+    static std::vector<cv::Point> getPointsFromObj(const Object &obj);
+
+private:
+    Inference*  yolov8_;
+    cv::Size     imgsz_;
+    int      seg_h_        = 120;
+    int      seg_w_        = 160;
+    int      seg_channels_ = 32;
+};
+

proto/def.grpc.proto

@@ -0,0 +1,93 @@
+syntax = "proto3";
+package JetStream;
+message RequestCmd
+{
+  int32 Id=1;  // 对于获取对应id的相机数据。1,2,3,4,其他表示所有
+  //对于打开视频流
+}
+message Response{
+  string info=2;
+}
+message Image{
+  int32 width=1;
+  int32 height=2;
+  int32 channel=3;
+  bool encode = 4;
+  bytes data= 5;
+}
+message PointCloud{
+  int32 size=1;
+  bytes data=2;
+}
+message LabelImage{
+  int32 label=1;
+  Image ir=2;
+}
+message Line{
+  int32 begin=1;
+  int32 end=2;
+}
+message SegBox{
+  float confidence = 1;
+  repeated Line lines= 2;
+}
+message LabelYolo{
+  int32 label=1;
+  repeated SegBox boxes=2;
+}
+message ResImage{
+  Image img1=1;
+  Image img2=2;
+  Image img3=3;
+  Image img4=4;
+}
+message ResCloud{
+  PointCloud cloud1=1;
+  PointCloud cloud2=2;
+  PointCloud cloud3=3;
+  PointCloud cloud4=4;
+}
+message MeasureInfo {
+  optional float x=1;
+  optional float y=2;
+  optional float trans_x=3;
+  optional float trans_y=4;
+  optional float theta=5;
+  optional float width=6;
+  optional float wheelbase=7;
+  optional float ftheta=8;
+  // 超界说明，此参数没有表示正常。
+  // 1:超宽，2：轴距超长，
+  // 3：前超界，4：后超界，5：左，6：右，7：左倾角过大，8：右倾角过大，9：前轮角过大
+  optional int32 border_plc=9;
+  optional int32 border_display=10;
+  optional string error=11;
+}
+message ResFrame{
+  ResImage images=1;
+  ResCloud clouds=2;
+  MeasureInfo measure_info=3;
+}
+service StreamServer{
+  rpc OpenImageStream(RequestCmd) returns(stream ResImage){}
+  rpc OpenMeasureDataStream(RequestCmd) returns(stream MeasureInfo){}
+  rpc CloseImageStream(RequestCmd) returns(Response){}
+  rpc CloseMeasureDataStream(RequestCmd) returns(Response){}
+  rpc GetCloud(RequestCmd) returns(ResCloud){}
+  rpc GetImage(RequestCmd) returns(ResImage){}
+}
+message Boundary{
+  float radius=1; //转台半径
+  float left=2;
+  float right=3;
+  float buttom=4;
+  float minAngle=5;
+  float maxAngle=6;
+  float fwheelAngle=7;
+}
+message Limit{
+  Boundary plc_limit=1;
+  Boundary display_limit=2;
+  float maxWidth=3;
+  float maxWheelBase=4;
+}

segMain.cpp

@@ -0,0 +1,22 @@
+//
+// Created by zx on 2023/10/7.
+//
+#include <iostream>
+#include "tool/log.hpp"
+#include "tool/load_protobuf.hpp"
+#include "yolo_seg_mqtt_async.h"
+
+int main(int argc, char * argv[]) {
+    // 初始化日志系统
+    ZX::InitGlog(PROJECT_NAME, ETC_PATH PROJECT_NAME "/LogInfo/");
+
+    YoloSegmentMqttAsyncClient::iter()->init(ETC_PATH PROJECT_NAME "/mqtt.json");
+
+    // 加载看门狗 TODO：添加看门狗，监视线程
+    LOG(INFO) << "---------- all init complate ----------";
+    while (true) {
+        usleep(100 * 1000);
+    }
+
+    return EXIT_SUCCESS;
+}

tensorrt_detect/common.hpp

@@ -0,0 +1,142 @@
+//
+// Created by ubuntu on 3/16/23.
+//
+
+#ifndef JETSON_SEGMENT_COMMON_HPP
+#define JETSON_SEGMENT_COMMON_HPP
+#include "NvInfer.h"
+#include "opencv2/opencv.hpp"
+#include <sys/stat.h>
+#include <unistd.h>
+
+#define CHECK(call)                                                                                                    \
+    do {                                                                                                               \
+        const cudaError_t error_code = call;                                                                           \
+        if (error_code != cudaSuccess) {                                                                               \
+            printf("CUDA Error:\n");                                                                                   \
+            printf("    File:       %s\n", __FILE__);                                                                  \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
+            printf("    Error code: %d\n", error_code);                                                                \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
+            exit(1);                                                                                                   \
+        }                                                                                                              \
+    } while (0)
+
+class Logger: public nvinfer1::ILogger {
+public:
+    nvinfer1::ILogger::Severity reportableSeverity;
+
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
+        reportableSeverity(severity)
+    {
+    }
+
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    {
+        if (severity > reportableSeverity) {
+            return;
+        }
+        switch (severity) {
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                std::cerr << "INTERNAL_ERROR: ";
+                break;
+            case nvinfer1::ILogger::Severity::kERROR:
+                std::cerr << "ERROR: ";
+                break;
+            case nvinfer1::ILogger::Severity::kWARNING:
+                std::cerr << "WARNING: ";
+                break;
+            case nvinfer1::ILogger::Severity::kINFO:
+                std::cerr << "INFO: ";
+                break;
+            default:
+                std::cerr << "VERBOSE: ";
+                break;
+        }
+        std::cerr << msg << std::endl;
+    }
+};
+
+inline int get_size_by_dims(const nvinfer1::Dims& dims)
+{
+    int size = 1;
+    for (int i = 0; i < dims.nbDims; i++) {
+        size *= dims.d[i];
+    }
+    return size;
+}
+
+inline int type_to_size(const nvinfer1::DataType& dataType)
+{
+    switch (dataType) {
+        case nvinfer1::DataType::kFLOAT:
+            return 4;
+        case nvinfer1::DataType::kHALF:
+            return 2;
+        case nvinfer1::DataType::kINT32:
+            return 4;
+        case nvinfer1::DataType::kINT8:
+            return 1;
+        case nvinfer1::DataType::kBOOL:
+            return 1;
+        default:
+            return 4;
+    }
+}
+
+inline static float clamp(float val, float min, float max)
+{
+    return val > min ? (val < max ? val : max) : min;
+}
+
+inline bool IsPathExist(const std::string& path)
+{
+    if (access(path.c_str(), 0) == F_OK) {
+        return true;
+    }
+    return false;
+}
+
+inline bool IsFile(const std::string& path)
+{
+    if (!IsPathExist(path)) {
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
+    }
+    struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
+}
+
+inline bool IsFolder(const std::string& path)
+{
+    if (!IsPathExist(path)) {
+        return false;
+    }
+    struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
+}
+
+namespace seg {
+struct Binding {
+    size_t         size  = 1;
+    size_t         dsize = 1;
+    nvinfer1::Dims dims;
+    std::string    name;
+};
+
+struct Object {
+    cv::Rect_<float> rect;
+    int              label = 0;
+    float            prob  = 0.0;
+    cv::Mat          boxMask;
+};
+
+struct PreParam {
+    float ratio  = 1.0f;
+    float dw     = 0.0f;
+    float dh     = 0.0f;
+    float height = 0;
+    float width  = 0;
+};
+}  // namespace seg
+#endif  // JETSON_SEGMENT_COMMON_HPP

tensorrt_detect/wheel-detector.cpp

@@ -0,0 +1,65 @@
+#include "wheel-detector.h"
+
+TensorrtWheelDetector::TensorrtWheelDetector(const std::string &model_file, const std::string &class_file){
+    cudaSetDevice(0);
+    
+    yolov8_ = new YOLOv8_seg(model_file);
+    yolov8_->make_pipe(false);
+    imgsz_=cv::Size{640, 480};
+    seg_h_        = 120;
+    seg_w_        = 160;
+    seg_channels_ = 32;
+}
+
+TensorrtWheelDetector::~TensorrtWheelDetector(){
+    if(yolov8_!=nullptr){
+        delete yolov8_;
+        yolov8_=nullptr;
+    }
+
+}
+
+bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs){
+    if(yolov8_==nullptr){
+        return false;
+    }
+    
+    yolov8_->copy_from_Mat(img, imgsz_);
+    yolov8_->infer();
+    float score_thres=0.9;
+    float iou_thres=0.65;
+    int topk=10;
+    yolov8_->postprocess(objs, score_thres, iou_thres, topk, seg_channels_, seg_h_, seg_w_);
+    return true;
+
+}
+
+bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res){
+    if(detect(img,objs))
+    {
+        const std::vector<std::string> classes={"none","wheel"};
+        const std::vector<std::vector<unsigned int>> colors = {{0, 114, 189},   {0, 255, 0}};
+        const std::vector<std::vector<unsigned int>> mask_colors = {{255, 56, 56},  {255, 0, 0}};
+        yolov8_->draw_objects(img, res, objs, classes, colors, mask_colors);
+        return true;
+    }else{
+        return false;
+    }
+}
+
+std::vector<cv::Point> TensorrtWheelDetector::getPointsFromObj(Object obj){
+    std::vector<cv::Point> ret;
+    int x=int(obj.rect.x+0.5);
+    int y=int(obj.rect.y+0.5);
+    int width=int(obj.rect.width);
+    int height=int(obj.rect.height);
+    
+    for(int i=0;i<height;++i){
+        for(int j=0;j<width;++j){
+            if(obj.boxMask.at<uchar>(i,j)!=0){
+                ret.push_back(cv::Point((x+j),y+i));
+            }
+        }
+    }
+    return ret;
+}

tensorrt_detect/wheel-detector.h

@@ -0,0 +1,25 @@
+#ifndef WHELL_DETECTOR__HH___
+#define WWHELL_DETECTOR__HH___
+#include "yolov8-seg.h"
+
+class TensorrtWheelDetector{
+public:
+    TensorrtWheelDetector(const std::string &model_file, const std::string &class_file);
+    ~TensorrtWheelDetector();
+
+
+    bool detect(cv::Mat& img,std::vector<Object>& objs);
+    bool detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res);
+
+    static std::vector<cv::Point> getPointsFromObj(Object obj);
+
+private:
+    YOLOv8_seg*  yolov8_;
+    cv::Size     imgsz_;
+    int      seg_h_        = 120;
+    int      seg_w_        = 160;
+    int      seg_channels_ = 32;
+
+};
+
+#endif

tensorrt_detect/yolov8-seg.cpp

@@ -0,0 +1,319 @@
+//
+// Created by ubuntu on 3/16/23.
+//
+
+#include "yolov8-seg.h"
+
+
+using namespace seg;
+
+YOLOv8_seg::YOLOv8_seg(const std::string& engine_file_path)
+{
+    std::ifstream file(engine_file_path, std::ios::binary);
+    assert(file.good());
+    file.seekg(0, std::ios::end);
+    auto size = file.tellg();
+    file.seekg(0, std::ios::beg);
+    char* trtModelStream = new char[size];
+    assert(trtModelStream);
+    file.read(trtModelStream, size);
+    file.close();
+    initLibNvInferPlugins(&this->gLogger, "");
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    assert(this->runtime != nullptr);
+
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    assert(this->engine != nullptr);
+    delete[] trtModelStream;
+    this->context = this->engine->createExecutionContext();
+
+    assert(this->context != nullptr);
+    cudaStreamCreate(&this->stream);
+    this->num_bindings = this->engine->getNbBindings();
+
+    for (int i = 0; i < this->num_bindings; ++i) {
+        Binding            binding;
+        nvinfer1::Dims     dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
+        binding.name             = name;
+        binding.dsize            = type_to_size(dtype);
+
+        bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
+            this->num_inputs += 1;
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
+            binding.size = get_size_by_dims(dims);
+            binding.dims = dims;
+            this->input_bindings.push_back(binding);
+            // set max opt shape
+            this->context->setBindingDimensions(i, dims);
+        }
+        else {
+            dims         = this->context->getBindingDimensions(i);
+            binding.size = get_size_by_dims(dims);
+            binding.dims = dims;
+            this->output_bindings.push_back(binding);
+            this->num_outputs += 1;
+        }
+        // printf("name: %s, size: %ld, dims: %d %d %d %d %d\n", 
+        //         name.c_str(), binding.dsize, dims.nbDims, dims.d[0], dims.d[1], dims.d[2], dims.d[3]);
+    }
+}
+
+YOLOv8_seg::~YOLOv8_seg()
+{
+    this->context->destroy();
+    this->engine->destroy();
+    this->runtime->destroy();
+    cudaStreamDestroy(this->stream);
+    for (auto& ptr : this->device_ptrs) {
+        CHECK(cudaFree(ptr));
+    }
+
+    for (auto& ptr : this->host_ptrs) {
+        CHECK(cudaFreeHost(ptr));
+    }
+}
+
+void YOLOv8_seg::make_pipe(bool warmup)
+{
+
+    for (auto& bindings : this->input_bindings) {
+        void* d_ptr;
+        CHECK(cudaMalloc(&d_ptr, bindings.size * bindings.dsize));
+        this->device_ptrs.push_back(d_ptr);
+    }
+
+    for (auto& bindings : this->output_bindings) {
+        void * d_ptr, *h_ptr;
+        size_t size = bindings.size * bindings.dsize;
+        CHECK(cudaMalloc(&d_ptr, size));
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
+        this->device_ptrs.push_back(d_ptr);
+        this->host_ptrs.push_back(h_ptr);
+    }
+
+    if (warmup) {
+        for (int i = 0; i < 10; i++) {
+            for (auto& bindings : this->input_bindings) {
+                size_t size  = bindings.size * bindings.dsize;
+                void*  h_ptr = malloc(size);
+                memset(h_ptr, 0, size);
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
+                free(h_ptr);
+            }
+            this->infer();
+        }
+        printf("model warmup 10 times\n");
+    }
+}
+
+void YOLOv8_seg::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
+{
+    const float inp_h  = size.height;
+    const float inp_w  = size.width;
+    float       height = image.rows;
+    float       width  = image.cols;
+
+    float r    = std::min(inp_h / height, inp_w / width);
+    int   padw = std::round(width * r);
+    int   padh = std::round(height * r);
+
+    cv::Mat tmp;
+    if ((int)width != padw || (int)height != padh) {
+        cv::resize(image, tmp, cv::Size(padw, padh));
+    }
+    else {
+        tmp = image.clone();
+    }
+
+    float dw = inp_w - padw;
+    float dh = inp_h - padh;
+
+    dw /= 2.0f;
+    dh /= 2.0f;
+    int top    = int(std::round(dh - 0.1f));
+    int bottom = int(std::round(dh + 0.1f));
+    int left   = int(std::round(dw - 0.1f));
+    int right  = int(std::round(dw + 0.1f));
+
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
+
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
+    this->pparam.ratio  = 1 / r;
+    this->pparam.dw     = dw;
+    this->pparam.dh     = dh;
+    this->pparam.height = height;
+    this->pparam.width  = width;
+}
+
+void YOLOv8_seg::copy_from_Mat(const cv::Mat& image)
+{
+    cv::Mat  nchw;
+    auto&    in_binding = this->input_bindings[0];
+    auto     width      = in_binding.dims.d[3];
+    auto     height     = in_binding.dims.d[2];
+    cv::Size size{width, height};
+    this->letterbox(image, nchw, size);
+
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
+
+    CHECK(cudaMemcpyAsync(
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
+}
+
+void YOLOv8_seg::copy_from_Mat(const cv::Mat& image, cv::Size& size)
+{
+    cv::Mat nchw;
+    this->letterbox(image, nchw, size);
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
+    CHECK(cudaMemcpyAsync(
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
+}
+
+void YOLOv8_seg::infer()
+{
+
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
+    for (int i = 0; i < this->num_outputs; i++) {
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
+        CHECK(cudaMemcpyAsync(
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
+    }
+    cudaStreamSynchronize(this->stream);
+}
+
+void YOLOv8_seg::postprocess(
+    std::vector<Object>& objs, float score_thres, float iou_thres, int topk, int seg_channels, int seg_h, int seg_w)
+{
+    objs.clear();
+    auto input_h      = this->input_bindings[0].dims.d[2];
+    auto input_w      = this->input_bindings[0].dims.d[3];
+    auto num_anchors  = this->output_bindings[0].dims.d[1];
+    auto num_channels = this->output_bindings[0].dims.d[2];
+
+    auto& dw     = this->pparam.dw;
+    auto& dh     = this->pparam.dh;
+    auto& width  = this->pparam.width;
+    auto& height = this->pparam.height;
+    auto& ratio  = this->pparam.ratio;
+
+    auto*   output = static_cast<float*>(this->host_ptrs[0]);
+    cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F, static_cast<float*>(this->host_ptrs[1]));
+
+    std::vector<int>      labels;
+    std::vector<float>    scores;
+    std::vector<cv::Rect> bboxes;
+    std::vector<cv::Mat>  mask_confs;
+    std::vector<int>      indices;
+
+    for (int i = 0; i < num_anchors; i++) {
+        float* ptr   = output + i * num_channels;
+        float  score = *(ptr + 4);
+		
+         /*if (score > score_thres) {
+             printf("num_channels: %d, score: %f\n", num_channels, score);
+         }*/
+	
+        if (score > score_thres) {
+            float x0 = *ptr++ - dw;
+            float y0 = *ptr++ - dh;
+            float x1 = *ptr++ - dw;
+            float y1 = *ptr++ - dh;
+
+            x0 = clamp(x0 * ratio, 0.f, width);
+            y0 = clamp(y0 * ratio, 0.f, height);
+            x1 = clamp(x1 * ratio, 0.f, width);
+            y1 = clamp(y1 * ratio, 0.f, height);
+
+            int     label     = *(++ptr);
+            cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, ++ptr);
+            mask_confs.push_back(mask_conf);
+            labels.push_back(label);
+            scores.push_back(score);
+            bboxes.push_back(cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0));
+        }
+    }
+
+#if defined(BATCHED_NMS)
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
+#else
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
+#endif
+
+    cv::Mat masks;
+    int     cnt = 0;
+    for (auto& i : indices) {
+        if (cnt >= topk) {
+            break;
+        }
+        cv::Rect tmp = bboxes[i];
+        Object   obj;
+        obj.label = labels[i];
+        obj.rect  = tmp;
+        obj.prob  = scores[i];
+        masks.push_back(mask_confs[i]);
+        objs.push_back(obj);
+        cnt += 1;
+    }
+    if (masks.empty()) {
+        // masks is empty
+    }
+    else {
+        cv::Mat matmulRes = (masks * protos).t();
+        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_h, seg_w});
+
+        std::vector<cv::Mat> maskChannels;
+        cv::split(maskMat, maskChannels);
+        int scale_dw = dw / input_w * seg_w;
+        int scale_dh = dh / input_h * seg_h;
+
+        cv::Rect roi(scale_dw, scale_dh, seg_w - 2 * scale_dw, seg_h - 2 * scale_dh);
+
+        for (int i = 0; i < indices.size(); i++) {
+            cv::Mat dest, mask;
+            cv::exp(-maskChannels[i], dest);
+            dest = 1.0 / (1.0 + dest);
+            dest = dest(roi);
+            cv::resize(dest, mask, cv::Size((int)width, (int)height), cv::INTER_LINEAR);
+            objs[i].boxMask = mask(objs[i].rect) > 0.5f;
+        }
+    }
+}
+
+void YOLOv8_seg::draw_objects(const cv::Mat&                                image,
+                              cv::Mat&                                      res,
+                              const std::vector<Object>&                    objs,
+                              const std::vector<std::string>&               CLASS_NAMES,
+                              const std::vector<std::vector<unsigned int>>& COLORS,
+                              const std::vector<std::vector<unsigned int>>& MASK_COLORS)
+{
+    res          = image.clone();
+    cv::Mat mask = image.clone();
+    for (auto& obj : objs) {
+        int        idx   = obj.label;
+        cv::Scalar color = cv::Scalar(COLORS[idx][0], COLORS[idx][1], COLORS[idx][2]);
+        cv::Scalar mask_color =
+            cv::Scalar(MASK_COLORS[idx % 20][0], MASK_COLORS[idx % 20][1], MASK_COLORS[idx % 20][2]);
+        cv::rectangle(res, obj.rect, color, 2);
+
+        char text[256];
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[idx].c_str(), obj.prob * 100);
+        mask(obj.rect).setTo(mask_color, obj.boxMask);
+
+        int      baseLine   = 0;
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
+
+        int x = (int)obj.rect.x;
+        int y = (int)obj.rect.y + 1;
+
+        if (y > res.rows)
+            y = res.rows;
+
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
+
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
+    }
+    cv::addWeighted(res, 0.5, mask, 0.8, 1, res);
+}

tensorrt_detect/yolov8-seg.h

@@ -0,0 +1,53 @@
+//
+// Created by ubuntu on 3/16/23.
+//
+#ifndef JETSON_SEGMENT_YOLOV8_SEG_HPP
+#define JETSON_SEGMENT_YOLOV8_SEG_HPP
+#include "NvInferPlugin.h"
+#include "common.hpp"
+#include <fstream>
+
+using namespace seg;
+
+class YOLOv8_seg {
+public:
+    explicit YOLOv8_seg(const std::string& engine_file_path);
+    ~YOLOv8_seg();
+
+    void                 make_pipe(bool warmup = true);
+    void                 copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
+    void                 infer();
+    void                 postprocess(std::vector<Object>& objs,
+                                     float                score_thres  = 0.25f,
+                                     float                iou_thres    = 0.65f,
+                                     int                  topk         = 100,
+                                     int                  seg_channels = 32,
+                                     int                  seg_h        = 160,
+                                     int                  seg_w        = 160);
+    static void          draw_objects(const cv::Mat&                                image,
+                                      cv::Mat&                                      res,
+                                      const std::vector<Object>&                    objs,
+                                      const std::vector<std::string>&               CLASS_NAMES,
+                                      const std::vector<std::vector<unsigned int>>& COLORS,
+                                      const std::vector<std::vector<unsigned int>>& MASK_COLORS);
+    int                  num_bindings;
+    int                  num_inputs  = 0;
+    int                  num_outputs = 0;
+    std::vector<Binding> input_bindings;
+    std::vector<Binding> output_bindings;
+    std::vector<void*>   host_ptrs;
+    std::vector<void*>   device_ptrs;
+
+    PreParam pparam;
+
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
+};
+
+#endif  // JETSON_SEGMENT_YOLOV8_SEG_HPP

yolo_seg_mqtt_async.cpp

@@ -3,13 +3,14 @@
 void YoloSegmentMqttAsyncClient::init(const std::string &file) {
     m_message_arrived_callback_ = CloudDataArrived;
 
-    MqttAsyncClient::init_from_proto(file);
-
 #ifdef ENABLE_TENSORRT_DETECT
     detector = new TensorrtWheelDetector(ETC_PATH PROJECT_NAME "/model.engine", ETC_PATH PROJECT_NAME "/class.txt");
 #else
     detector = new TensorrtWheelDetector(ETC_PATH PROJECT_NAME "/model.onnx", ETC_PATH PROJECT_NAME "/class.txt");
 #endif
+
+    MqttAsyncClient::init_from_proto(file);
+
 }
 
 // 识别结果数据包含识别结果、时间序号、具体识别信息
@@ -23,22 +24,24 @@ int YoloSegmentMqttAsyncClient::CloudDataArrived(void *client, char *topicName,
         int lable;
         cv::Mat merge_mat;
         JetStream::LabelImage testimage2;
+        JetStream::LabelYolo seg_results;
         testimage2.ParseFromArray(message->payload, message->payloadlen);
         NetMessageTrans::Proto2lableMat(testimage2, lable, merge_mat);
+        seg_results.set_label(lable);
+        seg_results.add_boxes();
+        seg_results.add_boxes();
+        seg_results.add_boxes();
+        seg_results.add_boxes();
 
         // 模型识别
         std::vector<Object> objs;
         cv::cvtColor(merge_mat, merge_mat, cv::COLOR_GRAY2RGB);
         tof_client->detector->detect(merge_mat, objs, merge_mat);
+//        cv::imshow("merge_mat", merge_mat);
+//        cv::waitKey(1);
 
         // 分离识别结果
         auto t = std::chrono::steady_clock::now();
-        JetStream::LabelYolo seg_results;
-        seg_results.set_label(lable);
-        seg_results.add_boxes();
-        seg_results.add_boxes();
-        seg_results.add_boxes();
-        seg_results.add_boxes();
         for (auto iter = objs.begin(); iter != objs.end(); iter++) {
             auto seg_points = tof_client->detector->getPointsFromObj(*iter);
             int device_index = (int(iter->rect.x / 640) * 0x01) | (int(iter->rect.y / 480) << 1);
@@ -51,7 +54,8 @@ int YoloSegmentMqttAsyncClient::CloudDataArrived(void *client, char *topicName,
                 continue;
             }
 
-            if (iter->prob > 0.9) {
+            seg_results.mutable_boxes(device_index)->set_confidence(iter->prob);
+            if (iter->prob > 0.6) {
                 int box_contour[480][2] = {0};
                 int x_alpha = (device_index & 0x1);
                 int y_alpha = ((device_index & 0x2) >> 1);

yolo_seg_mqtt_async.h

@@ -1,20 +1,15 @@
 #pragma once
 
-#include <pcl/point_types.h>
-#include <pcl/point_cloud.h>
-
 #include "pahoc/mqtt_async.h"
 #include "tool/log.hpp"
 #include "net_message_trans.h"
 
 #ifdef ENABLE_TENSORRT_DETECT
-#include "detect/tensorrt_detect/wheel-detector.h"
+#include "tensorrt_detect/wheel-detector.h"
 #else
-#include "detect/onnx_detect/wheel-detector.h"
+#include "onnx_detect/wheel-detector.h"
 #endif
 
-#include "communication/data_buf_lable.hpp"
-
 class YoloSegmentMqttAsyncClient : public MqttAsyncClient {
     friend class MqttAsyncClient;
 public: