detect_topic

最新推荐文章于 2024-07-06 21:01:07 发布

涂鸦c

最新推荐文章于 2024-07-06 21:01:07 发布

阅读量190

点赞数

文章标签： linux ubuntu YOLO c++

本文链接：https://blog.csdn.net/qq_16377055/article/details/133787595

版权

1.接收ros图像

2.yolo识别

3.sort跟踪

yolo.h

#ifndef YOLO_H
#define YOLO_H

#include <cuda_runtime_api.h>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "NvInferPlugin.h"
#include "NvInferRuntimeCommon.h"
#include "NvOnnxParser.h"

#include <math.h>
#include <array>
#include <fstream>
#include <iostream>
#include <random>
#include <sstream>
#include <string>
#include <vector>

using nvinfer1::Dims2;
using nvinfer1::Dims3;
using nvinfer1::IBuilder;
using nvinfer1::IBuilderConfig;
using nvinfer1::ICudaEngine;
using nvinfer1::IExecutionContext;
using nvinfer1::IHostMemory;
using nvinfer1::ILogger;
using nvinfer1::INetworkDefinition;
using Severity = nvinfer1::ILogger::Severity;

using cv::Mat;
using std::array;
using std::cout;
using std::endl;
using std::ifstream;
using std::ios;
using std::ofstream;
using std::string;
using std::vector;


class Logger : public ILogger {
 public:
  void log(Severity severity, const char* msg) noexcept override {
    if (severity != Severity::kINFO) {
      std::cout << msg << std::endl;
    }
  }
};



class Yolo {
 public:
  Yolo(char* model_path);
  float letterbox(
      const cv::Mat& image,
      cv::Mat& out_image,
      const cv::Size& new_shape,
      int stride,
      const cv::Scalar& color,
      bool fixed_shape,
      bool scale_up);
  float* blobFromImage(cv::Mat& img);
  void draw_objects(const cv::Mat& img, float* Boxes, int* ClassIndexs, int* BboxNum, float* Scores);
  void Init(char* model_path);
  void Infer(
      int aWidth,
      int aHeight,
      int aChannel,
      unsigned char* aBytes,
      float* Boxes,
      int* ClassIndexs,
      int* BboxNum,
      float* Scores);
  ~Yolo();

 private:
  nvinfer1::ICudaEngine* engine = nullptr;
  nvinfer1::IRuntime* runtime = nullptr;
  nvinfer1::IExecutionContext* context = nullptr;
  cudaStream_t stream = nullptr;
  void* buffs[5];
  int iH, iW, in_size, out_size1, out_size2, out_size3, out_size4;
  Logger gLogger;
};

#endif // YOLO_H

yolo.cpp

#include "yolo.h"

float Yolo::letterbox(
    const cv::Mat& image,
    cv::Mat& out_image,
    const cv::Size& new_shape = cv::Size(640, 640),
    int stride = 32,
    const cv::Scalar& color = cv::Scalar(114, 114, 114),
    bool fixed_shape = false,
    bool scale_up = true) {
  cv::Size shape = image.size();
  float r = std::min(
      (float)new_shape.height / (float)shape.height, (float)new_shape.width / (float)shape.width);
  if (!scale_up) {
    r = std::min(r, 1.0f);
  }

  int newUnpad[2]{
      (int)std::round((float)shape.width * r), (int)std::round((float)shape.height * r)};

  cv::Mat tmp;
  if (shape.width != newUnpad[0] || shape.height != newUnpad[1]) {
    cv::resize(image, tmp, cv::Size(newUnpad[0], newUnpad[1]));
  } else {
    tmp = image.clone();
  }

  float dw = new_shape.width - newUnpad[0];
  float dh = new_shape.height - newUnpad[1];

  if (!fixed_shape) {
    dw = (float)((int)dw % stride);
    dh = (float)((int)dh % stride);
  }

  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, out_image, top, bottom, left, right, cv::BORDER_CONSTANT, color);

  return 1.0f / r;
}

float* Yolo::blobFromImage(cv::Mat& img) {
  float* blob = new float[img.total() * 3];
  int channels = 3;
  int img_h = img.rows;
  int img_w = img.cols;
  for (size_t c = 0; c < channels; c++) {
    for (size_t h = 0; h < img_h; h++) {
      for (size_t w = 0; w < img_w; w++) {
        blob[c * img_w * img_h + h * img_w + w] = (float)img.at<cv::Vec3b>(h, w)[c] / 255.0;
      }
    }
  }
  return blob;
}

void Yolo::draw_objects(const cv::Mat& img, float* Boxes, int* ClassIndexs, int* BboxNum, float* Scores) {
  for (int j = 0; j < BboxNum[0]; j++) {
    cv::Rect rect(Boxes[j * 4], Boxes[j * 4 + 1], Boxes[j * 4 + 2], Boxes[j * 4 + 3]);
    cv::rectangle(img, rect, cv::Scalar(0x27, 0xC1, 0x36), 2);
    cv::putText(
        img,
        std::to_string(ClassIndexs[j])+":"+std::to_string(Scores[j]),
        cv::Point(rect.x, rect.y - 1),
        cv::FONT_HERSHEY_PLAIN,
        1.2,
        cv::Scalar(0xFF, 0xFF, 0xFF),
        2);
    //cv::imwrite("result.jpg", img);
  }
}

Yolo::Yolo(char* model_path) {
  ifstream ifile(model_path, ios::in | ios::binary);
  if (!ifile) {
    cout << "read serialized file failed\n";
    std::abort();
  }

  ifile.seekg(0, ios::end);
  const int mdsize = ifile.tellg();
  ifile.clear();
  ifile.seekg(0, ios::beg);
  vector<char> buf(mdsize);
  ifile.read(&buf[0], mdsize);
  ifile.close();
  cout << "model size: " << mdsize << endl;

  runtime = nvinfer1::createInferRuntime(gLogger);
  initLibNvInferPlugins(&gLogger, "");
  engine = runtime->deserializeCudaEngine((void*)&buf[0], mdsize, nullptr);
  auto in_dims = engine->getBindingDimensions(engine->getBindingIndex("images"));
  iH = in_dims.d[2];
  iW = in_dims.d[3];
  in_size = 1;
  for (int j = 0; j < in_dims.nbDims; j++) {
    in_size *= in_dims.d[j];
  }
  auto out_dims1 = engine->getBindingDimensions(engine->getBindingIndex("num_dets"));
  out_size1 = 1;
  for (int j = 0; j < out_dims1.nbDims; j++) {
    out_size1 *= out_dims1.d[j];
  }
  auto out_dims2 = engine->getBindingDimensions(engine->getBindingIndex("det_boxes"));
  out_size2 = 1;
  for (int j = 0; j < out_dims2.nbDims; j++) {
    out_size2 *= out_dims2.d[j];
  }
  auto out_dims3 = engine->getBindingDimensions(engine->getBindingIndex("det_scores"));
  out_size3 = 1;
  for (int j = 0; j < out_dims3.nbDims; j++) {
    out_size3 *= out_dims3.d[j];
  }
  auto out_dims4 = engine->getBindingDimensions(engine->getBindingIndex("det_classes"));
  out_size4 = 1;
  for (int j = 0; j < out_dims4.nbDims; j++) {
    out_size4 *= out_dims4.d[j];
  }
  context = engine->createExecutionContext();
  if (!context) {
    cout << "create execution context failed\n";
    std::abort();
  }

  cudaError_t state;
  state = cudaMalloc(&buffs[0], in_size * sizeof(float));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }
  state = cudaMalloc(&buffs[1], out_size1 * sizeof(int));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaMalloc(&buffs[2], out_size2 * sizeof(float));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaMalloc(&buffs[3], out_size3 * sizeof(float));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaMalloc(&buffs[4], out_size4 * sizeof(int));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaStreamCreate(&stream);
  if (state) {
    cout << "create stream failed\n";
    std::abort();
  }
}

void Yolo::Infer(
    int aWidth,
    int aHeight,
    int aChannel,
    unsigned char* aBytes,
    float* Boxes,
    int* ClassIndexs,
    int* BboxNum,
    float* Scores) {
  cv::Mat img(aHeight, aWidth, CV_MAKETYPE(CV_8U, aChannel), aBytes);
  cv::Mat pr_img;
  float scale = letterbox(img, pr_img, {iW, iH}, 32, {114, 114, 114}, true);
  cv::cvtColor(pr_img, pr_img, cv::COLOR_BGR2RGB);
  float* blob = blobFromImage(pr_img);

  static int* num_dets = new int[out_size1];
  static float* det_boxes = new float[out_size2];
  static float* det_scores = new float[out_size3];
  static int* det_classes = new int[out_size4];

  cudaError_t state =
      cudaMemcpyAsync(buffs[0], &blob[0], in_size * sizeof(float), cudaMemcpyHostToDevice, stream);
  if (state) {
    cout << "transmit to device failed\n";
    std::abort();
  }
  context->enqueueV2(&buffs[0], stream, nullptr);
  state =
      cudaMemcpyAsync(num_dets, buffs[1], out_size1 * sizeof(int), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  state = cudaMemcpyAsync(
      det_boxes, buffs[2], out_size2 * sizeof(float), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  state = cudaMemcpyAsync(
      det_scores, buffs[3], out_size3 * sizeof(float), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }

  state = cudaMemcpyAsync(
      det_classes, buffs[4], out_size4 * sizeof(int), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  BboxNum[0] = num_dets[0];
  int img_w = img.cols;
  int img_h = img.rows;
  int x_offset = (iW * scale - img_w) / 2;
  int y_offset = (iH * scale - img_h) / 2;
  for (size_t i = 0; i < num_dets[0]; i++) {
    float x0 = (det_boxes[i * 4]) * scale - x_offset;
    float y0 = (det_boxes[i * 4 + 1]) * scale - y_offset;
    float x1 = (det_boxes[i * 4 + 2]) * scale - x_offset;
    float y1 = (det_boxes[i * 4 + 3]) * scale - y_offset;
    x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f);
    y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f);
    x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f);
    y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f);
    Boxes[i * 4] = x0;
    Boxes[i * 4 + 1] = y0;
    Boxes[i * 4 + 2] = x1 - x0;
    Boxes[i * 4 + 3] = y1 - y0;
    ClassIndexs[i] = det_classes[i];
    Scores[i] = det_scores[i];
  }
  delete blob;
}

Yolo::~Yolo() {
  cudaStreamSynchronize(stream);
  cudaFree(buffs[0]);
  cudaFree(buffs[1]);
  cudaFree(buffs[2]);
  cudaFree(buffs[3]);
  cudaFree(buffs[4]);
  cudaStreamDestroy(stream);
  context->destroy();
  engine->destroy();
  runtime->destroy();
}

listenerMat.cpp

// 1.包含头文件；
#include "rclcpp/rclcpp.hpp"
#include "std_msgs/msg/string.hpp"
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/msg/image.hpp>
#include <opencv4/opencv2/opencv.hpp>
#include "utils.h"
#include "kalmanboxtracker.h"
#include "sort.h"
#include <algorithm>
#include <vector>
#include <iostream>

#include "std_msgs/msg/string.hpp"
 
using namespace std::chrono_literals;
using namespace std::chrono_literals;
using std::placeholders::_1;
//using std::placeholders::_1;
using std::placeholders::_2;

#include "yolo.h"


Yolo *myyolo;
sort *mot_tracker;
 
std::vector<cv::Vec3b> colors;

std::string msgstr = "";

cv::Mat detectResultMat = cv::Mat::zeros(1080,720,CV_8UC3);


void filter(float* Boxes,int* ClassIndexs,int* BboxNum,float* Scores,
            float conf_thres, std::vector<int> classes,float *newBoxes,
                    int *newClassIndexs,int *newBboxNum,float *newScores)
{
    int count = 0;
    for (int i =0;i<BboxNum[0];i++) {
      if (Scores[i]>conf_thres){
        std::vector<int>::iterator iter=std::find(classes.begin(),classes.end(),ClassIndexs[i]);
        if ( iter!=classes.end())
        {
          //cout << "Not found" << endl;
          for (int j=0;j<4;j++) {
            newBoxes[count*4+j]=Boxes[i*4+j];
            }

            newClassIndexs[count]=ClassIndexs[i];

            newScores[count]=Scores[i];

            count++;
        }
      }

    }

    newBboxNum[0]=count;
}

// 3.定义节点类；
class MinimalSubscriber : public rclcpp::Node
{
  public:
    MinimalSubscriber(std::string name)
    : Node(name)
    {
      srand(time(0));
      for (size_t i = 0; i < 100; i++)
        colors.push_back(cv::Vec3b(rand() % 255, rand() % 255, rand() % 255));

      mot_tracker = new sort(5, 4, 0.25);

      // 3-1.创建订阅方；
      //subscription_ = this->create_subscription<sensor_msgs::msg::CompressedImage>("pubImageTopic", 10, std::bind(&MinimalSubscriber::topic_callback, this, _1));
      subscription_ = this->create_subscription<sensor_msgs::msg::CompressedImage>("pubImageTopic", 10, std::bind(&MinimalSubscriber::topic_callback, this, _1));


       // 3-1.创建发布方；
      msgPublisher_ = this->create_publisher<std_msgs::msg::String>("detectResultMsgtopic", 10);
      // 3-2.创建定时器；
      msgTimer = this->create_wall_timer(50ms, std::bind(&MinimalSubscriber::msg_timer_callback, this));

      //mat topic
      matPublisher_ = this->create_publisher<sensor_msgs::msg::CompressedImage>("detectResultMatTopic", 10);
      matTimer = this->create_wall_timer(50ms, std::bind(&MinimalSubscriber::mat_timer_callback, this));
    }
 
  private:
    // 3-2.处理订阅到的消息；

    //rclcpp::Subscription<sensor_msgs::msg::CompressedImage> subscription_;
    rclcpp::Subscription<sensor_msgs::msg::CompressedImage>::SharedPtr subscription_;

    //void topic_callback(const sensor_msgs::msg::CompressedImage & msg) const

    rclcpp::TimerBase::SharedPtr msgTimer;
    rclcpp::Publisher<std_msgs::msg::String>::SharedPtr msgPublisher_;

    //std_msgs::msg::String detectResultMessage;// = std_msgs::msg::String();
    //detectResultMessage.data  = "null";


    rclcpp::TimerBase::SharedPtr matTimer;
    rclcpp::Publisher<sensor_msgs::msg::CompressedImage>::SharedPtr matPublisher_;

    //sensor_msgs::msg::CompressedImage::SharedPtr detect_result_img_compressed_;// = cv_bridge::CvImage(std_msgs::msg::Header(), "bgr8", des1080).toCompressedImageMsg();



    void msg_timer_callback(){
      // 3-3.组织消息并发布。
      auto detectResultMessage = std_msgs::msg::String();
      detectResultMessage.data = msgstr;
      RCLCPP_INFO(this->get_logger(), "发布的消息：'%s'", detectResultMessage.data.c_str());
      msgPublisher_ ->publish(detectResultMessage);
    }


    void mat_timer_callback(){
      sensor_msgs::msg::CompressedImage::SharedPtr detect_result_img_compressed_ = cv_bridge::CvImage(std_msgs::msg::Header(), "bgr8", detectResultMat).toCompressedImageMsg();
      //RCLCPP_INFO(this->get_logger(), "发布的消息：检测图片");
      this->matPublisher_ ->publish(*detect_result_img_compressed_);
    }


    
    void topic_callback(const sensor_msgs::msg::CompressedImage::ConstPtr  msg) const
    {

    try
    {
      cv_bridge::CvImagePtr cv_ptr_compressed = cv_bridge::toCvCopy(msg,sensor_msgs::image_encodings::BGR8);
      cv::Mat imgCallback = cv_ptr_compressed->image;
      //cv::imshow("imgCallback",imgCallback);
      //cv::waitKey(1);

      cv::Mat des1080;
 
      //cv::resize(imgCallback, des1080, cv::Size(1920, 1080), 0, 0, cv::INTER_NEAREST);

      cv::resize(imgCallback, des1080, cv::Size(1080, 720), 0, 0, cv::INTER_NEAREST);


      //std::cout<<"cv_ptr_compressed: "<<cv_ptr_compressed->image.cols<<" h: "<<cv_ptr_compressed->image.rows<<std::endl;
      //std::cout<<des1080.size()<<std::endl;

      float* Boxes = new float[400];
      int* BboxNum = new int[1];
      int* ClassIndexs = new int[100];
      float* Scores = new float[100];

      //auto start = std::chrono::system_clock::now();
      myyolo->Infer(des1080.cols, des1080.rows, des1080.channels(), des1080.data, Boxes, ClassIndexs, BboxNum,Scores);
      //auto end = std::chrono::system_clock::now();

      //std::cout<<std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count()<<"ms"<<std::endl;

      //myyolo->draw_objects(des1080, Boxes, ClassIndexs, BboxNum, Scores);
      float newBoxes[400];
      int newBboxNum[1]; 
      int newClassIndexs[100]; 
      float newScores[100];
      float conf_thres = 0.4;
      std::vector<int> classes;
      //classes.push_back(1);//识别车
      classes.push_back(0);//识别人

      filter(Boxes, ClassIndexs, BboxNum,Scores,conf_thres,classes,newBoxes,newClassIndexs,newBboxNum,newScores);
     
      std::vector<cv::Rect> detections;

      for (int i =0; i< newBboxNum[0]; i++){

           cv::Rect rect(newBoxes[i * 4], newBoxes[i * 4 + 1], newBoxes[i * 4 + 2], newBoxes[i * 4 + 3]);
           detections.push_back(rect);
       }

       std::vector<std::vector<float>> trackers = mot_tracker->update(detections);
       //qDebug()<<trackers;
       std::string sendstr = "";
       if(!trackers.empty()){

            for (auto tracker : trackers)
            {
                cv::rectangle(des1080, cv::Rect(tracker[0], tracker[1],
                            tracker[2]- tracker[0], tracker[3]- tracker[1]),
                            colors[int(tracker[4])%100], 2);
                cv::putText(des1080,std::to_string(int(tracker[4])),
                cv::Point(tracker[0], tracker[1]),
                cv::FONT_HERSHEY_PLAIN,1.5,
                colors[int(tracker[4])%100],
                2);
              
          std::string objnum = std::to_string(int(tracker[4]));
          std::string x0 = std::to_string(int(tracker[0]));
          std::string y0 = std::to_string(int(tracker[1]));

          std::string x1 = std::to_string(int(tracker[2]));
          std::string y1 = std::to_string(int(tracker[3]));
          sendstr = sendstr+objnum+","+x0+","+y0+","+x1+","+y1+";";
          }
        
          //this->detectResultMessage.data  = std::to_string(84);
          //msgstr = "sddh15151515";
          //std::string str
          msgstr = sendstr;

       }
       else{
          //msgstr = "null";
          msgstr = sendstr;
       }



      cv::imshow("des1080",des1080);
      cv::waitKey(1);


      detectResultMat = des1080;

      //auto end = std::chrono::system_clock::now();

      //std::cout<<std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count()<<"ms"<<std::endl;

    }
    catch (cv_bridge::Exception& e)
    {
      std::cout<<"Could not convert !!!"<<std::endl;
      //cv::Mat des1080;
      //detectResultMat = des1080;
    }  
    }
    
};
 
int main(int argc, char * argv[])
{

  std::string str = "./src/detect_topic/src/yolov7-tiny-nms.trt";
  char* chr = const_cast<char*>(str.c_str());
  myyolo = new Yolo(chr);
  // 2.初始化 ROS2 客户端；
  rclcpp::init(argc, argv);
  // 4.调用spin函数，并传入节点对象指针。
  auto node = std::make_shared<MinimalSubscriber>("detect");
  rclcpp::spin(node);
  //rclcpp::spin(std::make_shared<MinimalSubscriber>());
  // 5.释放资源；
  rclcpp::shutdown();
  return 0;
}

package.xml

<?xml version="1.0"?>
<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
<package format="3">
  <name>detect_topic</name>
  <version>0.0.0</version>
  <description>TODO: Package description</description>
  <maintainer email="1187506013@qq.com">tuya</maintainer>
  <license>TODO: License declaration</license>

  <buildtool_depend>ament_cmake</buildtool_depend>

  <depend>rclcpp</depend>
  <depend>std_msgs</depend>

  <test_depend>ament_lint_auto</test_depend>
  <test_depend>ament_lint_common</test_depend>

  <export>
    <build_type>ament_cmake</build_type>
  </export>
</package>

CMakeLists.txt

cmake_minimum_required(VERSION 3.8)
project(detect_topic)


add_definitions(-std=c++14)

# Default to C++14
#if(NOT CMAKE_CXX_STANDARD)
set(CMAKE_CXX_STANDARD 14)
#endif()


if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
  add_compile_options(-Wall -Wextra -Wpedantic)
endif()

# find dependencies
find_package(ament_cmake REQUIRED)
find_package(rclcpp REQUIRED)
find_package(std_msgs REQUIRED)

find_package(cv_bridge REQUIRED)
find_package(image_transport)
find_package(sensor_msgs REQUIRED)

find_package(CUDA REQUIRED)

include_directories(${CUDA_INCLUDE_DIRS})

include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/aarch64-linux-gnu/)
link_directories(/usr/lib/aarch64-linux-gnu/)

# tensorrt
#include_directories(/home/tuya/TensorRT-8.6.1.6/targets/x86_64-linux-gnu/include/)
#link_directories(/home/tuya/TensorRT-8.6.1.6/targets/x86_64-linux-gnu/lib/)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")

add_executable(detect
  src/listenerMat.cpp
  src/yolo.h 
  src/yolo.cpp
  src/hungarian.h
  src/hungarian.cpp
  src/kalmanboxtracker.h
  src/kalmanboxtracker.cpp
  src/sort.h
  src/sort.cpp
  src/utils.h
  src/utils.cpp
)

find_package(OpenCV REQUIRED)
include_directories( ${OpenCV_INCLUDE_DIRS})
target_link_libraries(detect ${OpenCV_LIBS})

# ros
target_link_libraries(detect
  ${rclcpp_LIBRARIES} 
)

target_link_libraries(detect nvinfer)
target_link_libraries(detect nvinfer_plugin)
target_link_libraries(detect cudart)

ament_target_dependencies(detect
  rclcpp
  std_msgs
  sensor_msgs 
  cv_bridge 
  OpenCV 
  image_transport
)
 


# 安装可执行文件
install(TARGETS detect
  DESTINATION lib/${PROJECT_NAME}
)

ament_package()

SORT代码

sort demo_涂鸦c的博客-CSDN博客