加入bytetrack处理

CMakeLists.txt

@@ -2,16 +2,17 @@ cmake_minimum_required(VERSION 3.10)
 project(rm_detector)
 
 if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
-        set(CMAKE_CUDA_ARCHITECTURES 70 75 80)
+  set(CMAKE_CUDA_ARCHITECTURES 70 75 80)
 endif(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
 
 add_definitions(-std=c++11)
 add_definitions(-DAPI_EXPORTS)
-option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
 set(CMAKE_CXX_STANDARD 11)
 set(CMAKE_BUILD_TYPE Debug)
 
-# TODO(Call for PR): make cmake compatible with Windows
+# CUDA
+set(CUDA_VERSION "11.6")
+option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
 set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
 enable_language(CUDA)
 
@@ -23,37 +24,55 @@ find_package(catkin REQUIRED COMPONENTS
         cv_bridge
         nodelet
         pluginlib
+        rm_msgs
         )
 #find_package(OpenCV REQUIRED)
 find_package(OpenCV 4 REQUIRED)
 
-
 include_directories(/usr/local/cuda/include)
+
 link_directories(/usr/local/cuda/lib64)
 
-include_directories(/opt/TensorRT-8.4.0.6/include/)
-link_directories(/opt/TensorRT-8.4.0.6/targets/x86_64-linux-gnu/lib/)
+# TensorRT
+set(TensorRT_VERSION "8.4.0.6")
+
+if (NOT EXISTS "/home/ywj/opt/TensorRT-${TensorRT_VERSION}")
+  MESSAGE(WARNING "/home/ywj/opt/TensorRT-${TensorRT_VERSION} not found.")
+endif()
+
+include_directories("/home/ywj/opt/TensorRT-${TensorRT_VERSION}/include/")
+
+link_directories("/home/ywj/opt/TensorRT-${TensorRT_VERSION}/targets/x86_64-linux-gnu/lib/")
+
+file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/utils/*.cpp ${PROJECT_SOURCE_DIR}/utils/*.cu)
 
-include_directories(./utils/)
-include_directories(./plugin/)
-file(GLOB_RECURSE SRCS ./utils/*.cpp ${PROJECT_SOURCE_DIR}/utils/*.cu)
 file(GLOB_RECURSE PLUGIN_SRCS ./plugin/*.cu)
 
-add_library(myplugins SHARED ${PLUGIN_SRCS})
-target_link_libraries(myplugins nvinfer cudart)
+add_library(network SHARED ${PLUGIN_SRCS})
 
-generate_dynamic_reconfigure_options(
-        cfg/dynamic.cfg
-)
+add_library(TRTx SHARED ${SRCS})
+
+target_link_libraries(network
+        nvinfer
+        cudart
+        )
+
+target_link_libraries(TRTx network)
 
 include_directories(
         include
+        ./utils/
+        ./plugin/
         ${catkin_INCLUDE_DIRS}
         ${OpenCV_INCLUDE_DIRS}
 )
 
 link_directories(${OpenCV_LIBRARY_DIRS})
 
+generate_dynamic_reconfigure_options(
+        cfg/dynamic.cfg
+)
+
 catkin_package(
         INCLUDE_DIRS
         include
@@ -67,18 +86,36 @@ catkin_package(
         LIBRARIES ${PROJECT_NAME}
 )
 
+if (NOT EXISTS "/usr/local/cuda-${CUDA_VERSION}/include")
+  MESSAGE(WARNING "/usr/local/cuda-${CUDA_VERSION}/include not found.")
+endif()
+
 include_directories(
         include
         ${catkin_INCLUDE_DIRS}
-        /usr/local/cuda-11.6/include
+        "/usr/local/cuda-${CUDA_VERSION}/include"
+        /usr/include/eigen3
 )
-## Declare a C++ library
-add_library(${PROJECT_NAME} src/detector.cpp
+
+add_library(${PROJECT_NAME}
+        src/detector.cpp
         src/inferencer.cpp
-        ${SRCS}
-        src/get_engine.cpp
+        src/rm_bytetrack/BYTETracker.cpp
+        src/rm_bytetrack/kalmanFilter.cpp
+        src/rm_bytetrack/lapjv.cpp
+        src/rm_bytetrack/STrack.cpp
 )
 
+add_executable(get_engine src/get_engine.cpp)
+
+target_link_libraries(get_engine
+        nvinfer
+        cudart
+        network
+        TRTx
+        ${OpenCV_LIBS}
+        )
+
 target_link_libraries(${PROJECT_NAME}
         ${catkin_LIBRARIES}
         ${InferenceEngine_LIBRARIES}
@@ -87,19 +124,12 @@ target_link_libraries(${PROJECT_NAME}
         ${TENSORRT_LIB}
         ${CUDA_LIB}
         ${glog_LIB}
+        nvinfer
+        cudart
+        network
+        TRTx
         )
 
-
-target_link_libraries(rm_detector nvinfer)
-target_link_libraries(rm_detector cudart)
-target_link_libraries(rm_detector myplugins)
-target_link_libraries(rm_detector ${OpenCV_LIBS})
-
-add_executable(get_engine src/get_engine.cpp)
-target_link_libraries(get_engine nvinfer)
-target_link_libraries(get_engine cudart)
-target_link_libraries(get_engine myplugins)
-target_link_libraries(get_engine ${OpenCV_LIBS})
-
 add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS})
+
 add_dependencies(${PROJECT_NAME} ${PROJECT_NAME}_gencfg)

cfg/dynamic.cfg

@@ -7,12 +7,14 @@ from dynamic_reconfigure.parameter_generator_catkin import *
 
 gen = ParameterGenerator()
 
-gen.add("g_conf_thresh", double_t, 1, "g_conf_thresh", 0.5, 0.1, 1.0)
-gen.add("g_nms_thresh",double_t, 2, "g_nms_thresh", 0.45, 0.1, 1.0)
+gen.add("g_car_conf_thresh", double_t, 1, "g_car_conf_thresh", 0.1, 0.01, 1.0)
+gen.add("g_car_nms_thresh",double_t, 2, "g_car_nms_thresh", 0.7, 0.1, 1.0)
+gen.add("g_armor_conf_thresh", double_t, 1, "g_armor_conf_thresh", 0.5, 0.1, 1.0)
+gen.add("g_armor_nms_thresh",double_t, 2, "g_armor_nms_thresh", 0.45, 0.1, 1.0)
 gen.add("g_conf_thresh2", double_t, 3, "g_conf_thresh2", 0.5, 0.1, 1.0)
 gen.add("g_nms_thresh2",double_t, 4, "g_nms_thresh2", 0.45, 0.1, 1.0)
 gen.add("g_turn_on_image",bool_t,5,"g_turn_on_image",True)
 gen.add("target_is_blue",bool_t,6,"target_is_blue",True)
 
 exit(gen.generate(PACKAGE, PACKAGE, "dynamic"))
-u
+u

include/rm_bytetrack/BYTETracker.h

@@ -0,0 +1,60 @@
+//
+// Created by ywj on 24-1-14.
+//
+
+#ifndef RM_RADAR_BYTETRACK_BYTETRACKER_H
+#define RM_RADAR_BYTETRACK_BYTETRACKER_H
+
+#include "rm_bytetrack/STrack.h"
+#include "rm_bytetrack/dataType.h"
+#include "rm_msgs/RadarTargetObject.h"
+#include "rm_bytetrack/lapjv.h"
+#include "types.h"
+
+namespace rm_bytetrack
+{
+class BYTETracker
+{
+public:
+  explicit BYTETracker(const int& frame_rate = 30, const int& track_buffer = 30, const float& track_thresh = 0.5,
+                       const float& high_thresh = 0.6, const float& match_thresh = 0.8);
+  ~BYTETracker();
+
+  std::vector<STrack> update(const std::vector<Object>& objects);
+  static cv::Scalar get_color(int idx);
+
+private:
+  static std::vector<STrack*> joint_stracks(std::vector<STrack*>& tlista, std::vector<STrack>& tlistb);
+  static std::vector<STrack> joint_stracks(std::vector<STrack>& tlista, std::vector<STrack>& tlistb);
+
+  static std::vector<STrack> sub_stracks(std::vector<STrack>& tlista, std::vector<STrack>& tlistb);
+  void remove_duplicate_stracks(std::vector<STrack>& resa, std::vector<STrack>& resb, std::vector<STrack>& stracksa,
+                                std::vector<STrack>& stracksb);
+
+  void linear_assignment(std::vector<std::vector<float> >& cost_matrix, int cost_matrix_size, int cost_matrix_size_size,
+                         float thresh, std::vector<std::vector<int> >& matches, std::vector<int>& unmatched_a,
+                         std::vector<int>& unmatched_b);
+  static std::vector<std::vector<float> > iou_distance(std::vector<STrack*>& atracks, std::vector<STrack>& btracks,
+                                                       int& dist_size, int& dist_size_size);
+  static std::vector<std::vector<float> > iou_distance(std::vector<STrack>& atracks, std::vector<STrack>& btracks);
+  static std::vector<std::vector<float> > ious(std::vector<std::vector<float> >& atlbrs,
+                                               std::vector<std::vector<float> >& btlbrs);
+
+  static double lapjv(const std::vector<std::vector<float> >& cost, std::vector<int>& rowsol, std::vector<int>& colsol,
+                      bool extend_cost = false, float cost_limit = LONG_MAX, bool return_cost = true);
+
+private:
+  float track_thresh_;
+  float high_thresh_;
+  float match_thresh_;
+  int frame_id_;
+  int max_time_lost_;
+
+  std::vector<STrack> tracked_stracks_;
+  std::vector<STrack> lost_stracks_;
+  std::vector<STrack> removed_stracks_;
+  KalmanFilter kalman_filter_;
+};
+}  // namespace rm_bytetrack
+
+#endif  // RM_RADAR_BYTETRACK_BYTETRACKER_H

include/rm_bytetrack/STrack.h

@@ -0,0 +1,64 @@
+//
+// Created by ywj on 24-1-14.
+//
+
+#ifndef RM_RADAR_BYTETRACK_STRACK_H
+#define RM_RADAR_BYTETRACK_STRACK_H
+
+#include <opencv2/opencv.hpp>
+#include "rm_bytetrack/kalmanFilter.h"
+
+enum TrackState
+{
+  New = 0,
+  Tracked,
+  Lost,
+  Removed
+};
+
+namespace rm_bytetrack
+{
+class STrack
+{
+public:
+  STrack(std::vector<float> tlwh_, float score, int class_id);
+  ~STrack();
+
+  std::vector<float> static tlbr_to_tlwh(std::vector<float>& tlbr);
+  void static multi_predict(std::vector<STrack*>& stracks, KalmanFilter& kalman_filter);
+  void static_tlwh();
+  void static_tlbr();
+  static std::vector<float> tlwh_to_xyah(std::vector<float> tlwh_tmp);
+  std::vector<float> to_xyah() const;
+  void mark_lost();
+  void mark_removed();
+  static int next_id();
+  int end_frame() const;
+
+  void activate(KalmanFilter& kalman_filter, int frame_id);
+  void re_activate(STrack& new_track, int frame_id, bool new_id = false);
+  void update(STrack& new_track, int frame_id);
+
+public:
+  bool is_activated_;
+  int track_class_id_;
+  int state_;
+  //        int class_id;
+
+  std::vector<float> _tlwh_;
+  std::vector<float> tlwh_;
+  std::vector<float> tlbr_;
+  int frame_id_;
+  int tracklet_len_;
+  int start_frame_;
+
+  KAL_MEAN mean_;
+  KAL_COVA covariance_;
+  float score_;
+
+private:
+  KalmanFilter kalman_filter_;
+};
+}  // namespace rm_bytetrack
+
+#endif  // RM_RADAR_BYTETRACK_STRACK_H

include/rm_bytetrack/dataType.h

@@ -0,0 +1,52 @@
+//
+// Created by ywj on 24-1-14.
+//
+
+#ifndef RM_RADAR_BYTETRACK_DATATYPE_H
+#define RM_RADAR_BYTETRACK_DATATYPE_H
+
+#include <Eigen/Core>
+#include <Eigen/Dense>
+#include <cstddef>
+#include <vector>
+#include <opencv2/opencv.hpp>
+
+typedef Eigen::Matrix<float, 1, 4, Eigen::RowMajor> DETECTBOX;
+typedef Eigen::Matrix<float, -1, 4, Eigen::RowMajor> DETECTBOXSS;
+typedef Eigen::Matrix<float, 1, 128, Eigen::RowMajor> FEATURE;
+typedef Eigen::Matrix<float, Eigen::Dynamic, 128, Eigen::RowMajor> FEATURESS;
+// typedef std::vector<FEATURE> FEATURESS;
+
+// Kalmanfilter
+// typedef Eigen::Matrix<float, 8, 8, Eigen::RowMajor> KAL_FILTER;
+typedef Eigen::Matrix<float, 1, 8, Eigen::RowMajor> KAL_MEAN;
+typedef Eigen::Matrix<float, 8, 8, Eigen::RowMajor> KAL_COVA;
+typedef Eigen::Matrix<float, 1, 4, Eigen::RowMajor> KAL_HMEAN;
+typedef Eigen::Matrix<float, 4, 4, Eigen::RowMajor> KAL_HCOVA;
+using KAL_DATA = std::pair<KAL_MEAN, KAL_COVA>;
+using KAL_HDATA = std::pair<KAL_HMEAN, KAL_HCOVA>;
+
+// main
+using RESULT_DATA = std::pair<int, DETECTBOX>;
+
+// tracker:
+using TRACKER_DATA = std::pair<int, FEATURESS>;
+using MATCH_DATA = std::pair<int, int>;
+typedef struct t
+{
+  std::vector<MATCH_DATA> matches;
+  std::vector<int> unmatched_tracks;
+  std::vector<int> unmatched_detections;
+} TRACHER_MATCHD;
+
+// linear_assignment:
+typedef Eigen::Matrix<float, -1, -1, Eigen::RowMajor> DYNAMICM;
+
+struct Object
+{
+  cv::Rect_<int> rect;
+  int label;
+  float prob;
+};
+
+#endif  // RM_RADAR_BYTETRACK_DATATYPE_H

include/rm_bytetrack/kalmanFilter.h

@@ -0,0 +1,33 @@
+//
+// Created by ywj on 24-1-14.
+//
+
+#ifndef RM_RADAR_BYTETRACK_KALMANFILTER_H
+#define RM_RADAR_BYTETRACK_KALMANFILTER_H
+
+#include "rm_bytetrack/dataType.h"
+
+namespace rm_bytetrack
+{
+class KalmanFilter
+{
+public:
+  static const double chi2inv95[10];
+  KalmanFilter();
+  KAL_DATA initiate(const DETECTBOX& measurement);
+  void predict(KAL_MEAN& mean, KAL_COVA& covariance);
+  KAL_HDATA project(const KAL_MEAN& mean, const KAL_COVA& covariance);
+  KAL_DATA update(const KAL_MEAN& mean, const KAL_COVA& covariance, const DETECTBOX& measurement);
+
+  Eigen::Matrix<float, 1, -1> gating_distance(const KAL_MEAN& mean, const KAL_COVA& covariance,
+                                              const std::vector<DETECTBOX>& measurements, bool only_position = false);
+
+private:
+  Eigen::Matrix<float, 8, 8, Eigen::RowMajor> _motion_mat;
+  Eigen::Matrix<float, 4, 8, Eigen::RowMajor> _update_mat;
+  float _std_weight_position;
+  float _std_weight_velocity;
+};
+}  // namespace rm_bytetrack
+
+#endif  // RM_RADAR_BYTETRACK_KALMANFILTER_H