From 8d805042c4c4407a7a23a432fc1c008c64bf9a66 Mon Sep 17 00:00:00 2001 From: asitaku <92627302+asitaku@users.noreply.github.com> Date: Thu, 18 Jan 2024 22:19:21 +0800 Subject: [PATCH 1/5] Add files via upload --- CMakeLists.txt | 92 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 61 insertions(+), 31 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9ab6d5f..8615b0a 100644 --- a/CMakeLists.txt +++ b/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) From e5bc8bcdc74f336be3598e9b76c478d6a9e26b68 Mon Sep 17 00:00:00 2001 From: asitaku <92627302+asitaku@users.noreply.github.com> Date: Thu, 18 Jan 2024 22:22:08 +0800 Subject: [PATCH 2/5] Add files via upload --- include/rm_bytetrack/BYTETracker.h | 60 +++ include/rm_bytetrack/STrack.h | 64 +++ include/rm_bytetrack/dataType.h | 52 +++ include/rm_bytetrack/kalmanFilter.h | 33 ++ include/rm_bytetrack/lapjv.h | 89 ++++ include/rm_detector/detector.h | 101 ++-- include/rm_detector/inferencer.h | 51 ++- launch/armor_detector.launch | 14 - src/detector.cpp | 241 +++++----- src/get_engine.cpp | 92 ++-- src/inferencer.cpp | 189 ++++---- src/rm_bytetrack/BYTETracker.cpp | 682 ++++++++++++++++++++++++++++ src/rm_bytetrack/STrack.cpp | 206 +++++++++ src/rm_bytetrack/kalmanFilter.cpp | 139 ++++++ src/rm_bytetrack/lapjv.cpp | 372 +++++++++++++++ 15 files changed, 2038 insertions(+), 347 deletions(-) create mode 100644 include/rm_bytetrack/BYTETracker.h create mode 100644 include/rm_bytetrack/STrack.h create mode 100644 include/rm_bytetrack/dataType.h create mode 100644 include/rm_bytetrack/kalmanFilter.h create mode 100644 include/rm_bytetrack/lapjv.h create mode 100644 src/rm_bytetrack/BYTETracker.cpp create mode 100644 src/rm_bytetrack/STrack.cpp create mode 100644 src/rm_bytetrack/kalmanFilter.cpp create mode 100644 src/rm_bytetrack/lapjv.cpp diff --git a/include/rm_bytetrack/BYTETracker.h b/include/rm_bytetrack/BYTETracker.h new file mode 100644 index 0000000..71fae7c --- /dev/null +++ b/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 update(const std::vector& objects); + static cv::Scalar get_color(int idx); + +private: + static std::vector joint_stracks(std::vector& tlista, std::vector& tlistb); + static std::vector joint_stracks(std::vector& tlista, std::vector& tlistb); + + static std::vector sub_stracks(std::vector& tlista, std::vector& tlistb); + void remove_duplicate_stracks(std::vector& resa, std::vector& resb, std::vector& stracksa, + std::vector& stracksb); + + void linear_assignment(std::vector >& cost_matrix, int cost_matrix_size, int cost_matrix_size_size, + float thresh, std::vector >& matches, std::vector& unmatched_a, + std::vector& unmatched_b); + static std::vector > iou_distance(std::vector& atracks, std::vector& btracks, + int& dist_size, int& dist_size_size); + static std::vector > iou_distance(std::vector& atracks, std::vector& btracks); + static std::vector > ious(std::vector >& atlbrs, + std::vector >& btlbrs); + + static double lapjv(const std::vector >& cost, std::vector& rowsol, std::vector& 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 tracked_stracks_; + std::vector lost_stracks_; + std::vector removed_stracks_; + KalmanFilter kalman_filter_; +}; +} // namespace rm_bytetrack + +#endif // RM_RADAR_BYTETRACK_BYTETRACKER_H diff --git a/include/rm_bytetrack/STrack.h b/include/rm_bytetrack/STrack.h new file mode 100644 index 0000000..70308cb --- /dev/null +++ b/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 +#include "rm_bytetrack/kalmanFilter.h" + +enum TrackState +{ + New = 0, + Tracked, + Lost, + Removed +}; + +namespace rm_bytetrack +{ +class STrack +{ +public: + STrack(std::vector tlwh_, float score, int class_id); + ~STrack(); + + std::vector static tlbr_to_tlwh(std::vector& tlbr); + void static multi_predict(std::vector& stracks, KalmanFilter& kalman_filter); + void static_tlwh(); + void static_tlbr(); + static std::vector tlwh_to_xyah(std::vector tlwh_tmp); + std::vector 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 _tlwh_; + std::vector tlwh_; + std::vector 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 diff --git a/include/rm_bytetrack/dataType.h b/include/rm_bytetrack/dataType.h new file mode 100644 index 0000000..889db69 --- /dev/null +++ b/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 +#include +#include +#include +#include + +typedef Eigen::Matrix DETECTBOX; +typedef Eigen::Matrix DETECTBOXSS; +typedef Eigen::Matrix FEATURE; +typedef Eigen::Matrix FEATURESS; +// typedef std::vector FEATURESS; + +// Kalmanfilter +// typedef Eigen::Matrix KAL_FILTER; +typedef Eigen::Matrix KAL_MEAN; +typedef Eigen::Matrix KAL_COVA; +typedef Eigen::Matrix KAL_HMEAN; +typedef Eigen::Matrix KAL_HCOVA; +using KAL_DATA = std::pair; +using KAL_HDATA = std::pair; + +// main +using RESULT_DATA = std::pair; + +// tracker: +using TRACKER_DATA = std::pair; +using MATCH_DATA = std::pair; +typedef struct t +{ + std::vector matches; + std::vector unmatched_tracks; + std::vector unmatched_detections; +} TRACHER_MATCHD; + +// linear_assignment: +typedef Eigen::Matrix DYNAMICM; + +struct Object +{ + cv::Rect_ rect; + int label; + float prob; +}; + +#endif // RM_RADAR_BYTETRACK_DATATYPE_H diff --git a/include/rm_bytetrack/kalmanFilter.h b/include/rm_bytetrack/kalmanFilter.h new file mode 100644 index 0000000..b7d807f --- /dev/null +++ b/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 gating_distance(const KAL_MEAN& mean, const KAL_COVA& covariance, + const std::vector& measurements, bool only_position = false); + +private: + Eigen::Matrix _motion_mat; + Eigen::Matrix _update_mat; + float _std_weight_position; + float _std_weight_velocity; +}; +} // namespace rm_bytetrack + +#endif // RM_RADAR_BYTETRACK_KALMANFILTER_H diff --git a/include/rm_bytetrack/lapjv.h b/include/rm_bytetrack/lapjv.h new file mode 100644 index 0000000..5b43f3f --- /dev/null +++ b/include/rm_bytetrack/lapjv.h @@ -0,0 +1,89 @@ +// +// Created by ywj on 24-1-14. +// + +#ifndef RM_RADAR_BYTETRACK_LAPJV_H +#define RM_RADAR_BYTETRACK_LAPJV_H + +#define LARGE 1000000 + +#if !defined TRUE +#define TRUE 1 +#endif +#if !defined FALSE +#define FALSE 0 +#endif + +#define NEW(x, t, n) \ + if ((x = (t*)malloc(sizeof(t) * (n))) == 0) \ + { \ + return -1; \ + } +#define FREE(x) \ + if (x != 0) \ + { \ + free(x); \ + x = 0; \ + } +#define SWAP_INDICES(a, b) \ + { \ + int_t _temp_index = a; \ + a = b; \ + b = _temp_index; \ + } + +#if 0 +#include +#define ASSERT(cond) assert(cond) +#define PRINTF(fmt, ...) printf(fmt, ##__VA_ARGS__) +#define PRINT_COST_ARRAY(a, n) \ + while (1) \ + { \ + printf(#a " = ["); \ + if ((n) > 0) \ + { \ + printf("%f", (a)[0]); \ + for (uint_t j = 1; j < n; j++) \ + { \ + printf(", %f", (a)[j]); \ + } \ + } \ + printf("]\n"); \ + break; \ + } +#define PRINT_INDEX_ARRAY(a, n) \ + while (1) \ + { \ + printf(#a " = ["); \ + if ((n) > 0) \ + { \ + printf("%d", (a)[0]); \ + for (uint_t j = 1; j < n; j++) \ + { \ + printf(", %d", (a)[j]); \ + } \ + } \ + printf("]\n"); \ + break; \ + } +#else +#define ASSERT(cond) +#define PRINTF(fmt, ...) +#define PRINT_COST_ARRAY(a, n) +#define PRINT_INDEX_ARRAY(a, n) +#endif + +typedef signed int int_t; +typedef unsigned int uint_t; +typedef double cost_t; +typedef char boolean; +typedef enum fp_t +{ + FP_1 = 1, + FP_2 = 2, + FP_DYNAMIC = 3 +} fp_t; + +extern int_t lapjv_internal(const uint_t n, cost_t* cost[], int_t* x, int_t* y); + +#endif // RM_RADAR_BYTETRACK_LAPJV_H diff --git a/include/rm_detector/detector.h b/include/rm_detector/detector.h index 11e5995..0df0b2e 100644 --- a/include/rm_detector/detector.h +++ b/include/rm_detector/detector.h @@ -9,87 +9,78 @@ #include #include #include +// #include +// #include +#include "rm_msgs/RadarTargetDetection.h" +#include "rm_msgs/RadarTargetDetectionArray.h" #include #include "rm_detector/dynamicConfig.h" #include #include #include #include "rm_detector/inferencer.h" +#include "rm_bytetrack/BYTETracker.h" - -//static Logger gLogger_; +// static Logger gLogger_; namespace rm_detector { class Detector : public nodelet::Nodelet { public: - Detector(); - - ~Detector() override; - - void onInit() override; - - void receiveFromCam(const sensor_msgs::CompressedImageConstPtr & image); - - void publicMsg(); - - void initalizeInfer(); + Detector(); + ~Detector() override; - void dynamicCallback(rm_detector::dynamicConfig& config); + void onInit() override; + void receiveFromCam(const sensor_msgs::CompressedImageConstPtr& image); + void publicMsg(); + void initalizeInfer(); + void dynamicCallback(rm_detector::dynamicConfig& config); - cv_bridge::CvImagePtr cv_image_; + cv_bridge::CvImagePtr cv_image_; - std_msgs::Float32MultiArray roi_data_; - std::vector roi_point_vec_; - cv::Point2f roi_data_point_r_; - cv::Point2f roi_data_point_l_; + std::vector roi_point_vec_; + cv::Point2f roi_data_point_r_; + cv::Point2f roi_data_point_l_; - std::string car_model_path_; - std::string armor_model_path_; + std::string car_model_path_; + std::string armor_model_path_; + bool turn_on_image_; + dynamic_reconfigure::Server* server_; + dynamic_reconfigure::Server::CallbackType callback_; - bool turn_on_image_; - dynamic_reconfigure::Server* server_; - dynamic_reconfigure::Server::CallbackType callback_; + std::string camera_pub_name_; + std::string nodelet_name_; - std::string camera_pub_name_; - std::string nodelet_name_; + std::string roi_data1_name_; // hero + std::string roi_data2_name_; // engineer + std::string roi_data3_name_; // standard + std::string roi_data4_name_; // standard + std::string roi_data5_name_; // standard + std::string roi_data6_name_; // sentry + std::string roi_data7_name_; // our sentry - std::string roi_data1_name_; //敌方英雄 - std::string roi_data2_name_; //敌方工程 - std::string roi_data3_name_; //敌方步兵 - std::string roi_data4_name_; //敌方步兵 - std::string roi_data5_name_; //敌方步兵 - std::string roi_data6_name_; //敌方烧饼 - std::string roi_data7_name_; //我方烧饼 + bool target_is_blue_; + bool left_camera_; - bool target_is_blue_; - bool left_camera_; - - Inferencer carInferencer; - Inferencer armorInferencer; - - std::vector select_objects; - - ros::NodeHandle nh_; - Logger gLogger_; + Inferencer car_inferencer_; + Inferencer armor_inferencer_; + ros::NodeHandle nh_; + Logger gLogger_; private: + unsigned int num_frame_; + unsigned int total_ms_; + rm_bytetrack::BYTETracker* tracker_ = nullptr; + std::vector output_stracks_; - ros::Publisher camera_pub_; - - ros::Subscriber camera_sub_; - - std::vector roi_data_pub_vec; - ros::Publisher roi_data_pub1_; - ros::Publisher roi_data_pub2_; - ros::Publisher roi_data_pub3_; - ros::Publisher roi_data_pub4_; - ros::Publisher roi_data_pub5_; - ros::Publisher roi_data_pub6_; - ros::Publisher roi_data_pub7_; + ros::Publisher camera_pub_; + ros::Publisher camera_pub_track_; + ros::Subscriber camera_sub_; + ros::Publisher roi_datas_pub_; + rm_msgs::RadarTargetDetectionArray roi_array_{}; }; } // namespace rm_detector \ No newline at end of file diff --git a/include/rm_detector/inferencer.h b/include/rm_detector/inferencer.h index a7c7c9d..34d44c9 100644 --- a/include/rm_detector/inferencer.h +++ b/include/rm_detector/inferencer.h @@ -15,42 +15,43 @@ using namespace nvinfer1; -namespace rm_detector { - class Inferencer { - public: - Inferencer(); +namespace rm_detector +{ +class Inferencer +{ +public: + Inferencer(); - ~Inferencer(); + ~Inferencer(); - void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, - float** gpu_output_buffer, float** cpu_output_buffer); + void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer, + float** cpu_output_buffer); - void detect(cv::Mat &frame); + void detect(cv::Mat& frame); - void infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, int batchsize); + void infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, int batchsize); - void deserialize_engine(std::string& engine_name, IRuntime** runtime, - ICudaEngine** engine, IExecutionContext** context); + void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, + IExecutionContext** context); - void init(std::string& engine_path, Logger& logger); + void init(std::string& engine_path, Logger& logger); - IRuntime* runtime_{nullptr}; - ICudaEngine* engine_{nullptr}; - IExecutionContext* context_{nullptr}; - Logger* logger_{nullptr}; + IRuntime* runtime_{ nullptr }; + ICudaEngine* engine_{ nullptr }; + IExecutionContext* context_{ nullptr }; + Logger* logger_{ nullptr }; - float* gpu_buffers_[2]; - float* cpu_output_buffer_ = nullptr; + float* gpu_buffers_[2]; + float* cpu_output_buffer_ = nullptr; - const static int K_OUTPUT_SIZE = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; + const static int K_OUTPUT_SIZE = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; - std::vector target_objects; - CudaProprecess cudaProprecess_; + std::vector target_objects_; + CudaProprecess cudaProprecess_; - float conf_thresh_; - float nms_thresh_; - - }; + float conf_thresh_; + float nms_thresh_; +}; } // namespace rm_detector #endif \ No newline at end of file diff --git a/launch/armor_detector.launch b/launch/armor_detector.launch index b13c57e..648feef 100644 --- a/launch/armor_detector.launch +++ b/launch/armor_detector.launch @@ -17,13 +17,6 @@ - - - - - - - @@ -36,13 +29,6 @@ - - - - - - - diff --git a/src/detector.cpp b/src/detector.cpp index 1251e5e..6f02965 100644 --- a/src/detector.cpp +++ b/src/detector.cpp @@ -4,153 +4,154 @@ // #include +#include "rm_msgs/RadarTargetDetection.h" +#include "rm_msgs/RadarTargetDetectionArray.h" namespace rm_detector { Detector::Detector() { - + num_frame_ = 0; + total_ms_ = 0; + tracker_ = new rm_bytetrack::BYTETracker(50, 100); } void Detector::onInit() { - nh_ = getMTPrivateNodeHandle(); - nh_.getParam("g_car_model_path", car_model_path_); - nh_.getParam("g_armor_model_path", armor_model_path_); - - nh_.getParam("camera_pub_name", camera_pub_name_); - nh_.getParam("nodelet_name", nodelet_name_); - - nh_.getParam("roi_data1_name", roi_data1_name_); - nh_.getParam("roi_data2_name", roi_data2_name_); - nh_.getParam("roi_data3_name", roi_data3_name_); - nh_.getParam("roi_data4_name", roi_data4_name_); - nh_.getParam("roi_data5_name", roi_data5_name_); - nh_.getParam("roi_data6_name", roi_data6_name_); - nh_.getParam("roi_data7_name", roi_data7_name_); - nh_.getParam("left_camera", left_camera_); - - initalizeInfer(); - - ros::NodeHandle nh_reconfig(nh_, nodelet_name_ + "_reconfig"); - server_ = new dynamic_reconfigure::Server(nh_reconfig); - callback_ = boost::bind(&Detector::dynamicCallback, this, _1); - server_->setCallback(callback_); - - if (left_camera_) - camera_sub_ = nh_.subscribe("/galaxy_camera/galaxy_camera_left/image_raw/compressed", 1, &Detector::receiveFromCam, this); - else - camera_sub_ = nh_.subscribe("/galaxy_camera/galaxy_camera_right/image_raw/compressed", 1, &Detector::receiveFromCam, this); - - camera_pub_ = nh_.advertise(camera_pub_name_, 1); - - roi_data_pub1_ = nh_.advertise(roi_data1_name_, 1); - roi_data_pub2_ = nh_.advertise(roi_data2_name_, 1); - roi_data_pub3_ = nh_.advertise(roi_data3_name_, 1); - roi_data_pub4_ = nh_.advertise(roi_data4_name_, 1); - roi_data_pub5_ = nh_.advertise(roi_data5_name_, 1); - roi_data_pub6_ = nh_.advertise(roi_data6_name_, 1); - roi_data_pub7_ = nh_.advertise(roi_data7_name_, 1); - - roi_data_pub_vec.push_back(roi_data_pub1_); - roi_data_pub_vec.push_back(roi_data_pub2_); - roi_data_pub_vec.push_back(roi_data_pub3_); - roi_data_pub_vec.push_back(roi_data_pub4_); - roi_data_pub_vec.push_back(roi_data_pub5_); - roi_data_pub_vec.push_back(roi_data_pub6_); - roi_data_pub_vec.push_back(roi_data_pub7_); + nh_ = getMTPrivateNodeHandle(); + nh_.getParam("g_car_model_path", car_model_path_); + nh_.getParam("g_armor_model_path", armor_model_path_); + + nh_.getParam("camera_pub_name", camera_pub_name_); + nh_.getParam("nodelet_name", nodelet_name_); + + nh_.getParam("left_camera", left_camera_); + + initalizeInfer(); + + ros::NodeHandle nh_reconfig(nh_, nodelet_name_ + "_reconfig"); + server_ = new dynamic_reconfigure::Server(nh_reconfig); + callback_ = boost::bind(&Detector::dynamicCallback, this, _1); + server_->setCallback(callback_); + + if (left_camera_) // TODO: Should we use the subscribeCamera function to receive camera info? + camera_sub_ = + nh_.subscribe("/galaxy_camera/galaxy_camera_left/image_raw/compressed", 1, &Detector::receiveFromCam, this); + else + camera_sub_ = + nh_.subscribe("/galaxy_camera/galaxy_camera_right/image_raw/compressed", 1, &Detector::receiveFromCam, this); + + camera_pub_ = nh_.advertise(camera_pub_name_, 1); + + camera_pub_track_ = nh_.advertise(camera_pub_name_ + "_track_", 1); + roi_datas_pub_ = nh_.advertise("rm_radar/roi_datas", 10); } -void Detector::receiveFromCam(const sensor_msgs::CompressedImageConstPtr &image) { -// auto start = std::chrono::system_clock::now(); - cv_bridge::CvImagePtr cv_image_ = cv_bridge::toCvCopy(image, sensor_msgs::image_encodings::BGR8); - - carInferencer.detect(cv_image_->image); - if (!carInferencer.target_objects.empty()) { - for (auto &object: carInferencer.target_objects) { - cv::Mat armor_cls_image = cv_image_->image(get_rect(cv_image_->image, object.bbox)).clone(); - - armorInferencer.detect(armor_cls_image); - if (armorInferencer.target_objects.empty()) continue; - object.class_id = armorInferencer.target_objects[0].class_id; - - select_objects.push_back(object); - - } - if (turn_on_image_) { - draw_bbox(cv_image_->image, select_objects); - camera_pub_.publish(cv_bridge::CvImage(std_msgs::Header(), "bgr8", cv_image_->image).toImageMsg()); - } - if (!select_objects.empty()) - publicMsg(); - select_objects.clear(); -// auto end = std::chrono::system_clock::now(); -// ROS_INFO("inference time: %ld ms",std::chrono::duration_cast(end - start).count()); +void Detector::receiveFromCam(const sensor_msgs::CompressedImageConstPtr& image) +{ + if (num_frame_ > 1000) + { + num_frame_ = 0; + total_ms_ = 0; + } + num_frame_++; + auto start = std::chrono::system_clock::now(); + cv_image_ = cv_bridge::toCvCopy(image, sensor_msgs::image_encodings::BGR8); + + car_inferencer_.detect(cv_image_->image); + + if (!car_inferencer_.target_objects_.empty()) + { + for (auto& object : car_inferencer_.target_objects_) + { + cv::Mat armor_cls_image = cv_image_->image(get_rect(cv_image_->image, object.bbox)).clone(); + + armor_inferencer_.detect(armor_cls_image); + if (armor_inferencer_.target_objects_.empty()) + { + object.class_id = -1; + continue; + } + object.class_id = armor_inferencer_.target_objects_[0].class_id; + object.conf = 1.0f; + } + + std::vector objects; + for (auto& targetObject : car_inferencer_.target_objects_) + { + Object object; + object.rect = get_rect(cv_image_->image, targetObject.bbox); + object.label = targetObject.class_id; + object.prob = targetObject.conf; + objects.push_back(object); + } + output_stracks_.clear(); + output_stracks_ = tracker_->update(objects); + + if (turn_on_image_) + { + cv::Mat img_clone = cv_image_->image.clone(); + draw_bbox(img_clone, car_inferencer_.target_objects_); + camera_pub_.publish(cv_bridge::CvImage(std_msgs::Header(), "bgr8", img_clone).toImageMsg()); + auto end = std::chrono::system_clock::now(); + total_ms_ = total_ms_ + std::chrono::duration_cast(end - start).count(); + for (int i = 0; i < output_stracks_.size(); i++) + { + std::vector tlwh = output_stracks_[i].tlwh_; + putText(cv_image_->image, cv::format("%d", output_stracks_[i].track_class_id_), cv::Point(tlwh[0], tlwh[1] - 5), + 0, 0.6, cv::Scalar(0, 0, 255), 2, cv::LINE_AA); + rectangle(cv_image_->image, cv::Rect(tlwh[0], tlwh[1], tlwh[2], tlwh[3]), cv::Scalar(255, 0, 0), 2); + } + putText(cv_image_->image, + cv::format("frame: %d fps: %d num: %d", num_frame_, num_frame_ * 1000000 / total_ms_, + output_stracks_.size()), + cv::Point(0, 30), 0, 0.6, cv::Scalar(0, 0, 255), 2, cv::LINE_AA); + camera_pub_track_.publish(cv_bridge::CvImage(std_msgs::Header(), "bgr8", cv_image_->image).toImageMsg()); } + + if (!output_stracks_.empty()) + publicMsg(); + } } void Detector::dynamicCallback(rm_detector::dynamicConfig& config) { - carInferencer.conf_thresh_ = config.g_conf_thresh; - carInferencer.nms_thresh_ = config.g_nms_thresh; - armorInferencer.conf_thresh_ = config.g_conf_thresh2; - armorInferencer.nms_thresh_ = config.g_nms_thresh2; - turn_on_image_ = config.g_turn_on_image; - target_is_blue_ = config.target_is_blue; - ROS_INFO("Settings have been seted"); + car_inferencer_.conf_thresh_ = config.g_car_conf_thresh; + car_inferencer_.nms_thresh_ = config.g_car_nms_thresh; + armor_inferencer_.conf_thresh_ = config.g_armor_conf_thresh; + armor_inferencer_.nms_thresh_ = config.g_armor_nms_thresh; + turn_on_image_ = config.g_turn_on_image; + target_is_blue_ = config.target_is_blue; + ROS_INFO("Settings have been seted"); } void Detector::initalizeInfer() { - cudaSetDevice(kGpuId); - carInferencer.init(car_model_path_, gLogger_); - armorInferencer.init(armor_model_path_, gLogger_); + cudaSetDevice(kGpuId); + car_inferencer_.init(car_model_path_, gLogger_); + armor_inferencer_.init(armor_model_path_, gLogger_); + + // tracker_ = new rm_bytetrack::BYTETracker(50, 100); } Detector::~Detector() { + this->roi_array_.detections.clear(); } -void Detector::publicMsg() {//enemy is blue - std::vector target; - if (target_is_blue_){ - target = {0, 1, 2, 3, 4, 5, 11}; - } else {//red - target = {6, 7, 8, 9, 10, 11, 5}; - } - - int car_size = select_objects.size(); - - for (size_t i = 0; i < car_size; i++) - { - std::vector::iterator iter; - iter = std::find(target.begin(), target.end(), select_objects[i].class_id); - if (iter == target.end()) { - continue; - } - auto index = std::distance(target.begin(), iter); - - roi_point_vec_.clear(); - roi_data_.data.clear(); - - float* box = select_objects[i].bbox; - - roi_data_point_l_.x = box[0] - box[2] / 2; - roi_data_point_l_.y = box[1] - box[3] / 2; - roi_data_point_r_.x = box[0] + box[2] / 2; - roi_data_point_r_.y = box[1] + box[3] / 2; - - roi_point_vec_.push_back(roi_data_point_l_); - roi_point_vec_.push_back(roi_data_point_r_); - - roi_data_.data.push_back(roi_point_vec_[0].x); - roi_data_.data.push_back(roi_point_vec_[0].y); - roi_data_.data.push_back(roi_point_vec_[1].x); - roi_data_.data.push_back(roi_point_vec_[1].y); - - roi_data_pub_vec[index].publish(roi_data_); - } +void Detector::publicMsg() +{ + rm_msgs::RadarTargetDetectionArray array; + array.header.stamp = ros::Time::now(); + for (auto& output_strack : output_stracks_) + { + rm_msgs::RadarTargetDetection data; + data.id = output_strack.track_class_id_; + data.position.data.assign(output_strack.tlbr_.begin(), output_strack.tlbr_.end()); + array.detections.push_back(data); + } + roi_datas_pub_.publish(array); } } // namespace rm_detector PLUGINLIB_EXPORT_CLASS(rm_detector::Detector, nodelet::Nodelet) \ No newline at end of file diff --git a/src/get_engine.cpp b/src/get_engine.cpp index 15e6b79..b6b86d4 100644 --- a/src/get_engine.cpp +++ b/src/get_engine.cpp @@ -12,63 +12,71 @@ #include #include #include +#include using namespace nvinfer1; static Logger gLogger; const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; -void serialize_engine(unsigned int max_batchsize, bool& is_p6, float& gd, float& gw, std::string& wts_name, std::string& engine_name) { - // Create builder - IBuilder* builder = createInferBuilder(gLogger); - IBuilderConfig* config = builder->createBuilderConfig(); +void serialize_engine(unsigned int max_batchsize, bool& is_p6, float& gd, float& gw, std::string& wts_name, + std::string& engine_name) +{ + // Create builder + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); - // Create model to populate the network, then set the outputs and create an engine - ICudaEngine *engine = nullptr; - if (is_p6) { - engine = build_det_p6_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); - } else { - engine = build_det_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); - } - assert(engine != nullptr); + // Create model to populate the network, then set the outputs and create an engine + ICudaEngine* engine = nullptr; + if (is_p6) + { + engine = build_det_p6_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + } + else + { + engine = build_det_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + } + assert(engine != nullptr); - // Serialize the engine - IHostMemory* serialized_engine = engine->serialize(); - assert(serialized_engine != nullptr); + // Serialize the engine + IHostMemory* serialized_engine = engine->serialize(); + assert(serialized_engine != nullptr); - // Save engine to file - std::ofstream p(engine_name, std::ios::binary); - if (!p) { - std::cerr << "Could not open plan output file" << std::endl; - assert(false); - } - p.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); + // Save engine to file + std::ofstream p(engine_name, std::ios::binary); + if (!p) + { + std::cerr << "Could not open plan output file" << std::endl; + assert(false); + } + p.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); - // Close everything down - engine->destroy(); - config->destroy(); - serialized_engine->destroy(); - builder->destroy(); + // Close everything down + engine->destroy(); + config->destroy(); + serialized_engine->destroy(); + builder->destroy(); } -int main(int argc, char** argv) { - cudaSetDevice(kGpuId); +int main(int argc, char** argv) +{ + cudaSetDevice(kGpuId); - std::string wts_name = "/media/robotzero/Steins_Gate/rm/tensorrtx/yolov5/docs/armor_det.wts"; - std::string engine_name = "/media/robotzero/Steins_Gate/rm/tensorrtx/yolov5/docs/armor_det.engine"; - bool is_p6 = false; - float gd = 0.33f, gw = 0.5f; - std::string img_dir = "/media/robotzero/Steins_Gate/rm/tensorrtx/yolov5/image_armor"; + std::string wts_name = "/home/ywj/code_ws/src/rm_detector/docs/car.wts"; + std::string engine_name = "/home/ywj/code_ws/src/rm_detector/docs/car_det.engine"; + bool is_p6 = false; + float gd = 0.33f, gw = 0.5f; -// if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) { -// std::cerr << "arguments not right!" << std::endl; -// std::cerr << "./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file" << std::endl; -// std::cerr << "./yolov5_det -d [.engine] ../images // deserialize plan file and run inference" << std::endl; -// return -1; -// } + // if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) { + // std::cerr << "arguments not right!" << std::endl; + // std::cerr << "./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to + // plan file" << std::endl; std::cerr << "./yolov5_det -d [.engine] ../images // deserialize plan file and run + // inference" << std::endl; return -1; + // } - // Create a model using the API directly and serialize it to a file - if (!wts_name.empty()) { + // Create a model using the API directly and serialize it to a file + if (!wts_name.empty()) + { serialize_engine(kBatchSize, is_p6, gd, gw, wts_name, engine_name); return 0; } diff --git a/src/inferencer.cpp b/src/inferencer.cpp index 56dc4f8..a1cb9cd 100644 --- a/src/inferencer.cpp +++ b/src/inferencer.cpp @@ -7,96 +7,103 @@ using namespace cv; using namespace nvinfer1; - -namespace rm_detector { - - Inferencer::Inferencer() { - cudaSetDevice(kGpuId); - } - - Inferencer::~Inferencer() { - // Release stream and buffers -// cudaStreamDestroy(stream_); - CUDA_CHECK(cudaFree(gpu_buffers_[0])); - CUDA_CHECK(cudaFree(gpu_buffers_[1])); - delete[] cpu_output_buffer_; - cudaProprecess_.cuda_preprocess_destroy(); - // Destroy the engine - context_->destroy(); - engine_->destroy(); - runtime_->destroy(); - } - - void Inferencer::init(std::string& engine_path, Logger& logger) { - logger_ = &logger; - - deserialize_engine(engine_path, &runtime_, &engine_, &context_); - - cudaProprecess_.cuda_preprocess_init(kMaxInputImageSize); - - prepare_buffers(engine_, &gpu_buffers_[0], - &gpu_buffers_[1], &cpu_output_buffer_); - } - - void Inferencer::prepare_buffers(ICudaEngine *engine, float **gpu_input_buffer, float **gpu_output_buffer, - float **cpu_output_buffer) { - assert(engine->getNbBindings() == 2); - // In order to bind the buffers, we need to know the names of the input and output tensors. - // Note that indices are guaranteed to be less than IEngine::getNbBindings() - const int inputIndex = engine->getBindingIndex(kInputTensorName); - const int outputIndex = engine->getBindingIndex(kOutputTensorName); - assert(inputIndex == 0); - assert(outputIndex == 1); - // Create GPU buffers on device - CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float))); - CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer, kBatchSize * K_OUTPUT_SIZE * sizeof(float))); - - *cpu_output_buffer = new float[kBatchSize * K_OUTPUT_SIZE]; - } - - void Inferencer::detect(Mat &frame) { - target_objects.clear(); - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - cudaProprecess_.cuda_batch_preprocess(frame, gpu_buffers_[0], kInputW, - kInputH, stream); - - infer(*context_, stream, (void**)gpu_buffers_, cpu_output_buffer_, kBatchSize); - - batch_nms(target_objects, cpu_output_buffer_, conf_thresh_, nms_thresh_); - - cudaStreamDestroy(stream); - } - - void Inferencer::infer(IExecutionContext& context, cudaStream_t& stream, void **gpu_buffers, float *output, int batchsize) { - context.enqueue(batchsize, gpu_buffers, stream, nullptr); - CUDA_CHECK(cudaMemcpyAsync(output, gpu_buffers[1], batchsize * K_OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); - cudaStreamSynchronize(stream); - } - - void Inferencer::deserialize_engine(std::string &engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context) { - std::ifstream file(engine_name, std::ios::binary); - if (!file.good()) { - std::cerr << "read " << engine_name << " error!" << std::endl; - assert(false); - } - size_t size = 0; - file.seekg(0, file.end); - size = file.tellg(); - file.seekg(0, file.beg); - char* serialized_engine = new char[size]; - assert(serialized_engine); - file.read(serialized_engine, size); - file.close(); - - *runtime = createInferRuntime(*logger_); - assert(*runtime); - *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size); - assert(*engine); - *context = (*engine)->createExecutionContext(); - assert(*context); - delete[] serialized_engine; - } - +namespace rm_detector +{ +Inferencer::Inferencer() +{ + cudaSetDevice(kGpuId); +} + +Inferencer::~Inferencer() +{ + // Release stream and buffers + // cudaStreamDestroy(stream_); + CUDA_CHECK(cudaFree(gpu_buffers_[0])); + CUDA_CHECK(cudaFree(gpu_buffers_[1])); + delete[] cpu_output_buffer_; + cudaProprecess_.cuda_preprocess_destroy(); + // Destroy the engine + context_->destroy(); + engine_->destroy(); + runtime_->destroy(); +} + +void Inferencer::init(std::string& engine_path, Logger& logger) +{ + logger_ = &logger; + + deserialize_engine(engine_path, &runtime_, &engine_, &context_); + + cudaProprecess_.cuda_preprocess_init(kMaxInputImageSize); + + prepare_buffers(engine_, &gpu_buffers_[0], &gpu_buffers_[1], &cpu_output_buffer_); +} + +void Inferencer::prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer, + float** cpu_output_buffer) +{ + assert(engine->getNbBindings() == 2); + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine->getBindingIndex(kInputTensorName); + const int outputIndex = engine->getBindingIndex(kOutputTensorName); + assert(inputIndex == 0); + assert(outputIndex == 1); + // Create GPU buffers on device + CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float))); + CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer, kBatchSize * K_OUTPUT_SIZE * sizeof(float))); + + *cpu_output_buffer = new float[kBatchSize * K_OUTPUT_SIZE]; +} + +void Inferencer::detect(Mat& frame) +{ + target_objects_.clear(); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + cudaProprecess_.cuda_batch_preprocess(frame, gpu_buffers_[0], kInputW, kInputH, stream); + + infer(*context_, stream, (void**)gpu_buffers_, cpu_output_buffer_, kBatchSize); + + batch_nms(target_objects_, cpu_output_buffer_, conf_thresh_, nms_thresh_); + + cudaStreamDestroy(stream); +} + +void Inferencer::infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, + int batchsize) +{ + context.enqueue(batchsize, gpu_buffers, stream, nullptr); + CUDA_CHECK(cudaMemcpyAsync(output, gpu_buffers[1], batchsize * K_OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, + stream)); + cudaStreamSynchronize(stream); +} + +void Inferencer::deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, + IExecutionContext** context) +{ + std::ifstream file(engine_name, std::ios::binary); + if (!file.good()) + { + std::cerr << "read " << engine_name << " error!" << std::endl; + assert(false); + } + size_t size = 0; + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + char* serialized_engine = new char[size]; + assert(serialized_engine); + file.read(serialized_engine, size); + file.close(); + + *runtime = createInferRuntime(*logger_); + assert(*runtime); + *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size); + assert(*engine); + *context = (*engine)->createExecutionContext(); + assert(*context); + delete[] serialized_engine; +} } // namespace rm_detector \ No newline at end of file diff --git a/src/rm_bytetrack/BYTETracker.cpp b/src/rm_bytetrack/BYTETracker.cpp new file mode 100644 index 0000000..f8fa703 --- /dev/null +++ b/src/rm_bytetrack/BYTETracker.cpp @@ -0,0 +1,682 @@ +// +// Created by ywj on 24-1-14. +// + +#include "rm_bytetrack/BYTETracker.h" +#include "rm_bytetrack/STrack.h" + +#include + +namespace rm_bytetrack +{ +BYTETracker::BYTETracker(const int& frame_rate, const int& track_buffer, const float& track_thresh, + const float& high_thresh, const float& match_thresh) +{ + this->track_thresh_ = track_thresh; + this->high_thresh_ = high_thresh; + this->match_thresh_ = match_thresh; + + frame_id_ = 0; + max_time_lost_ = int(frame_rate / 30.0 * track_buffer); + std::cout << "Init ByteTrack!" << std::endl; +} + +BYTETracker::~BYTETracker() = default; + +std::vector BYTETracker::update(const std::vector& objects) +{ + // std::cout << "step 1" << std::endl; + ////////////////// Step 1: Get detections ////////////////// + this->frame_id_++; + std::vector activated_stracks; + std::vector refind_stracks; + std::vector removed_stracks; + std::vector lost_stracks; + std::vector detections; + std::vector detections_low; + + std::vector detections_cp; + std::vector tracked_stracks_swap; + std::vector resa, resb; + std::vector output_stracks; + + std::vector unconfirmed; + std::vector tracked_stracks; + std::vector strack_pool; + std::vector r_tracked_stracks; + + if (!objects.empty()) + { + for (auto& object : objects) + { + std::vector tlbr_; + tlbr_.resize(4); + tlbr_[0] = object.rect.x; + tlbr_[1] = object.rect.y; + tlbr_[2] = object.rect.x + object.rect.width; + tlbr_[3] = object.rect.y + object.rect.height; + + float score = object.prob; + + STrack strack(STrack::tlbr_to_tlwh(tlbr_), score, object.label); + + if (score >= track_thresh_) + { + detections.push_back(strack); + } + else + { + detections_low.push_back(strack); + } + } + } + + // Add newly detected tracklets to tracked_stracks + for (auto& tracked_strack : this->tracked_stracks_) + { + if (!tracked_strack.is_activated_) + unconfirmed.push_back(&tracked_strack); + else + tracked_stracks.push_back(&tracked_strack); + } + // std::cout << "step 2" << std::endl; + ////////////////// Step 2: First association, with IoU ////////////////// + strack_pool = joint_stracks(tracked_stracks, this->lost_stracks_); + STrack::multi_predict(strack_pool, this->kalman_filter_); + + std::vector> dists; + int dist_size = 0, dist_size_size = 0; + dists = iou_distance(strack_pool, detections, dist_size, dist_size_size); + + std::vector> matches; + std::vector u_track, u_detection; + linear_assignment(dists, dist_size, dist_size_size, match_thresh_, matches, u_track, u_detection); + + for (auto& matche : matches) + { + STrack* track = strack_pool[matche[0]]; + STrack* det = &detections[matche[1]]; + if (track->state_ == TrackState::Tracked) + { + track->update(*det, this->frame_id_); + activated_stracks.push_back(*track); + } + else + { + track->re_activate(*det, this->frame_id_, false); + refind_stracks.push_back(*track); + } + } + + // std::cout << "step 3" << std::endl; + ////////////////// Step 3: Second association, using low score dets ////////////////// + for (int i : u_detection) + { + detections_cp.push_back(detections[i]); + } + detections.clear(); + detections.assign(detections_low.begin(), detections_low.end()); + + for (int i : u_track) + { + if (strack_pool[i]->state_ == TrackState::Tracked) + { + r_tracked_stracks.push_back(strack_pool[i]); + } + } + + dists.clear(); + dists = iou_distance(r_tracked_stracks, detections, dist_size, dist_size_size); + + matches.clear(); + u_track.clear(); + u_detection.clear(); + linear_assignment(dists, dist_size, dist_size_size, 0.5, matches, u_track, u_detection); + + for (auto& matche : matches) + { + STrack* track = r_tracked_stracks[matche[0]]; + STrack* det = &detections[matche[1]]; + if (track->state_ == TrackState::Tracked) + { + track->update(*det, this->frame_id_); + activated_stracks.push_back(*track); + } + else + { + track->re_activate(*det, this->frame_id_, false); + refind_stracks.push_back(*track); + } + } + + for (int i : u_track) + { + STrack* track = r_tracked_stracks[i]; + if (track->state_ != TrackState::Lost) + { + track->mark_lost(); + lost_stracks.push_back(*track); + } + } + + // Deal with unconfirmed tracks, usually tracks with only one beginning frame + detections.clear(); + detections.assign(detections_cp.begin(), detections_cp.end()); + + dists.clear(); + dists = iou_distance(unconfirmed, detections, dist_size, dist_size_size); + + matches.clear(); + std::vector u_unconfirmed; + u_detection.clear(); + linear_assignment(dists, dist_size, dist_size_size, 0.7, matches, u_unconfirmed, u_detection); + + for (auto& matche : matches) + { + unconfirmed[matche[0]]->update(detections[matche[1]], this->frame_id_); + activated_stracks.push_back(*unconfirmed[matche[0]]); + } + + for (int i : u_unconfirmed) + { + STrack* track = unconfirmed[i]; + track->mark_removed(); + removed_stracks.push_back(*track); + } + + // std::cout << "step 4" << std::endl; + ////////////////// Step 4: Init new stracks ////////////////// + for (int i : u_detection) + { + STrack* track = &detections[i]; + if (track->score_ < this->high_thresh_) + continue; + track->activate(this->kalman_filter_, this->frame_id_); + activated_stracks.push_back(*track); + } + + // std::cout << "step 5" << std::endl; + ////////////////// Step 5: Update state ////////////////// + for (auto& lost_strack : this->lost_stracks_) + { + if (this->frame_id_ - lost_strack.end_frame() > this->max_time_lost_) + { + lost_strack.mark_removed(); + removed_stracks.push_back(lost_strack); + } + } + + for (auto& tracked_strack : this->tracked_stracks_) + { + if (tracked_strack.state_ == TrackState::Tracked) + { + tracked_stracks_swap.push_back(tracked_strack); + } + } + this->tracked_stracks_.clear(); + this->tracked_stracks_.assign(tracked_stracks_swap.begin(), tracked_stracks_swap.end()); + + this->tracked_stracks_ = joint_stracks(this->tracked_stracks_, activated_stracks); + this->tracked_stracks_ = joint_stracks(this->tracked_stracks_, refind_stracks); + + this->lost_stracks_ = sub_stracks(this->lost_stracks_, this->tracked_stracks_); + for (const auto& lost_strack : lost_stracks) + { + this->lost_stracks_.push_back(lost_strack); + } + + this->lost_stracks_ = sub_stracks(this->lost_stracks_, this->removed_stracks_); + for (const auto& removed_strack : removed_stracks) + { + this->removed_stracks_.push_back(removed_strack); + } + + remove_duplicate_stracks(resa, resb, this->tracked_stracks_, this->lost_stracks_); + + this->tracked_stracks_.clear(); + this->tracked_stracks_.assign(resa.begin(), resa.end()); + this->lost_stracks_.clear(); + this->lost_stracks_.assign(resb.begin(), resb.end()); + + for (auto& tracked_strack : this->tracked_stracks_) + { + if (tracked_strack.is_activated_) + { + output_stracks.push_back(tracked_strack); + } + } + return output_stracks; +} + +std::vector BYTETracker::joint_stracks(std::vector& tlista, std::vector& tlistb) +{ + std::map exists; + std::vector res; + for (auto& i : tlista) + { + exists.insert(std::pair(i->track_class_id_, 1)); + res.push_back(i); + } + for (auto& i : tlistb) + { + int tid = i.track_class_id_; + if (!exists[tid] || exists.count(tid) == 0) + { + exists[tid] = 1; + res.push_back(&i); + } + } + return res; +} + +std::vector BYTETracker::joint_stracks(std::vector& tlista, std::vector& tlistb) +{ + std::map exists; + std::vector res; + for (auto& i : tlista) + { + exists.insert(std::pair(i.track_class_id_, 1)); + res.push_back(i); + } + for (auto& i : tlistb) + { + int tid = i.track_class_id_; + if (!exists[tid] || exists.count(tid) == 0) + { + exists[tid] = 1; + res.push_back(i); + } + } + return res; +} + +std::vector BYTETracker::sub_stracks(std::vector& tlista, std::vector& tlistb) +{ + std::map stracks; + for (auto& i : tlista) + { + stracks.insert(std::pair(i.track_class_id_, i)); + } + for (auto& i : tlistb) + { + int tid = i.track_class_id_; + if (stracks.count(tid) != 0) + { + stracks.erase(tid); + } + } + + std::vector res; + std::map::iterator it; + for (it = stracks.begin(); it != stracks.end(); ++it) + { + res.push_back(it->second); + } + + return res; +} + +void BYTETracker::remove_duplicate_stracks(std::vector& resa, std::vector& resb, + std::vector& stracksa, std::vector& stracksb) +{ + std::vector> pdist = iou_distance(stracksa, stracksb); + std::vector> pairs; + for (int i = 0; i < pdist.size(); i++) + { + for (int j = 0; j < pdist[i].size(); j++) + { + if (pdist[i][j] < 0.15) + { + pairs.emplace_back(i, j); + } + } + } + + std::vector dupa, dupb; + for (auto& pair : pairs) + { + int timep = stracksa[pair.first].frame_id_ - stracksa[pair.first].start_frame_; + int timeq = stracksb[pair.second].frame_id_ - stracksb[pair.second].start_frame_; + if (timep > timeq) + dupb.push_back(pair.second); + else + dupa.push_back(pair.first); + } + + for (int i = 0; i < stracksa.size(); i++) + { + auto iter = find(dupa.begin(), dupa.end(), i); + if (iter == dupa.end()) + { + resa.push_back(stracksa[i]); + } + } + + for (int i = 0; i < stracksb.size(); i++) + { + auto iter = find(dupb.begin(), dupb.end(), i); + if (iter == dupb.end()) + { + resb.push_back(stracksb[i]); + } + } +} + +void BYTETracker::linear_assignment(std::vector>& cost_matrix, int cost_matrix_size, + int cost_matrix_size_size, float thresh, std::vector>& matches, + std::vector& unmatched_a, std::vector& unmatched_b) +{ + if (cost_matrix.empty()) + { + for (int i = 0; i < cost_matrix_size; i++) + { + unmatched_a.push_back(i); + } + for (int i = 0; i < cost_matrix_size_size; i++) + { + unmatched_b.push_back(i); + } + return; + } + + std::vector rowsol; + std::vector colsol; + float c = lapjv(cost_matrix, rowsol, colsol, true, thresh); + for (int i = 0; i < rowsol.size(); i++) + { + if (rowsol[i] >= 0) + { + std::vector match; + match.push_back(i); + match.push_back(rowsol[i]); + matches.push_back(match); + } + else + { + unmatched_a.push_back(i); + } + } + + for (int i = 0; i < colsol.size(); i++) + { + if (colsol[i] < 0) + { + unmatched_b.push_back(i); + } + } +} + +std::vector> BYTETracker::ious(std::vector>& atlbrs, + std::vector>& btlbrs) +{ + std::vector> ious; + if (atlbrs.size() * btlbrs.size() == 0) + return ious; + + ious.resize(atlbrs.size()); + for (auto& iou : ious) + { + iou.resize(btlbrs.size()); + } + + // bbox_ious + for (int k = 0; k < btlbrs.size(); k++) + { + std::vector ious_tmp; + float box_area = (btlbrs[k][2] - btlbrs[k][0] + 1) * (btlbrs[k][3] - btlbrs[k][1] + 1); + for (int n = 0; n < atlbrs.size(); n++) + { + float iw = std::min(atlbrs[n][2], btlbrs[k][2]) - std::max(atlbrs[n][0], btlbrs[k][0]) + 1; + if (iw > 0) + { + float ih = std::min(atlbrs[n][3], btlbrs[k][3]) - std::max(atlbrs[n][1], btlbrs[k][1]) + 1; + if (ih > 0) + { + float ua = (atlbrs[n][2] - atlbrs[n][0] + 1) * (atlbrs[n][3] - atlbrs[n][1] + 1) + box_area - iw * ih; + ious[n][k] = iw * ih / ua; + } + else + { + ious[n][k] = 0.0; + } + } + else + { + ious[n][k] = 0.0; + } + } + } + + return ious; +} + +std::vector> BYTETracker::iou_distance(std::vector& atracks, std::vector& btracks, + int& dist_size, int& dist_size_size) +{ + std::vector> cost_matrix; + if (atracks.size() * btracks.size() == 0) + { + dist_size = atracks.size(); + dist_size_size = btracks.size(); + return cost_matrix; + } + std::vector> atlbrs, btlbrs; + for (auto& atrack : atracks) + { + atlbrs.push_back(atrack->tlbr_); + } + for (auto& btrack : btracks) + { + btlbrs.push_back(btrack.tlbr_); + } + + dist_size = atracks.size(); + dist_size_size = btracks.size(); + + std::vector> _ious = ious(atlbrs, btlbrs); + + for (auto& i : _ious) + { + std::vector _iou; + for (float j : i) + { + _iou.push_back(1 - j); + } + cost_matrix.push_back(_iou); + } + + return cost_matrix; +} + +std::vector> BYTETracker::iou_distance(std::vector& atracks, std::vector& btracks) +{ + std::vector> atlbrs, btlbrs; + for (auto& atrack : atracks) + { + atlbrs.push_back(atrack.tlbr_); + } + for (auto& btrack : btracks) + { + btlbrs.push_back(btrack.tlbr_); + } + + std::vector> _ious = ious(atlbrs, btlbrs); + std::vector> cost_matrix; + for (auto& i : _ious) + { + std::vector _iou; + for (float j : i) + { + _iou.push_back(1 - j); + } + cost_matrix.push_back(_iou); + } + + return cost_matrix; +} + +double BYTETracker::lapjv(const std::vector>& cost, std::vector& rowsol, + std::vector& colsol, bool extend_cost, float cost_limit, bool return_cost) +{ + std::vector> cost_c; + cost_c.assign(cost.begin(), cost.end()); + + std::vector> cost_c_extended; + + int n_rows = cost.size(); + int n_cols = cost[0].size(); + rowsol.resize(n_rows); + colsol.resize(n_cols); + + int n = 0; + if (n_rows == n_cols) + { + n = n_rows; + } + else + { + if (!extend_cost) + { + std::cout << "set extend_cost=True" << std::endl; + system("pause"); + exit(0); + } + } + + if (extend_cost || cost_limit < LONG_MAX) + { + n = n_rows + n_cols; + cost_c_extended.resize(n); + for (auto& i : cost_c_extended) + i.resize(n); + + if (cost_limit < LONG_MAX) + { + for (auto& i : cost_c_extended) + { + for (float& j : i) + { + j = cost_limit / 2.0; + } + } + } + else + { + float cost_max = -1; + for (auto& i : cost_c) + { + for (float j : i) + { + if (j > cost_max) + cost_max = j; + } + } + for (auto& i : cost_c_extended) + { + for (float& j : i) + { + j = cost_max + 1; + } + } + } + + for (int i = n_rows; i < cost_c_extended.size(); i++) + { + for (int j = n_cols; j < cost_c_extended[i].size(); j++) + { + cost_c_extended[i][j] = 0; + } + } + for (int i = 0; i < n_rows; i++) + { + for (int j = 0; j < n_cols; j++) + { + cost_c_extended[i][j] = cost_c[i][j]; + } + } + + cost_c.clear(); + cost_c.assign(cost_c_extended.begin(), cost_c_extended.end()); + } + + double** cost_ptr; + cost_ptr = new double*[sizeof(double*) * n]; + for (int i = 0; i < n; i++) + cost_ptr[i] = new double[sizeof(double) * n]; + + for (int i = 0; i < n; i++) + { + for (int j = 0; j < n; j++) + { + cost_ptr[i][j] = cost_c[i][j]; + } + } + + int* x_c = new int[sizeof(int) * n]; + int* y_c = new int[sizeof(int) * n]; + + int ret = lapjv_internal(n, cost_ptr, x_c, y_c); + if (ret != 0) + { + std::cout << "Calculate Wrong!" << std::endl; + system("pause"); + exit(0); + } + + double opt = 0.0; + + if (n != n_rows) + { + for (int i = 0; i < n; i++) + { + if (x_c[i] >= n_cols) + x_c[i] = -1; + if (y_c[i] >= n_rows) + y_c[i] = -1; + } + for (int i = 0; i < n_rows; i++) + { + rowsol[i] = x_c[i]; + } + for (int i = 0; i < n_cols; i++) + { + colsol[i] = y_c[i]; + } + + if (return_cost) + { + for (int i = 0; i < rowsol.size(); i++) + { + if (rowsol[i] != -1) + { + // cout << i << "\t" << rowsol[i] << "\t" << cost_ptr[i][rowsol[i]] << endl; + opt += cost_ptr[i][rowsol[i]]; + } + } + } + } + else if (return_cost) + { + for (int i = 0; i < rowsol.size(); i++) + { + opt += cost_ptr[i][rowsol[i]]; + } + } + + for (int i = 0; i < n; i++) + { + delete[] cost_ptr[i]; + } + delete[] cost_ptr; + delete[] x_c; + delete[] y_c; + + return opt; +} + +cv::Scalar BYTETracker::get_color(int idx) +{ + idx += 3; + return cv::Scalar(37 * idx % 255, 17 * idx % 255, 29 * idx % 255); +} +} // namespace rm_bytetrack \ No newline at end of file diff --git a/src/rm_bytetrack/STrack.cpp b/src/rm_bytetrack/STrack.cpp new file mode 100644 index 0000000..fd78648 --- /dev/null +++ b/src/rm_bytetrack/STrack.cpp @@ -0,0 +1,206 @@ +// +// Created by ywj on 24-1-14. +// + +#include + +#include "rm_bytetrack/STrack.h" + +namespace rm_bytetrack +{ +STrack::STrack(std::vector tlwh, float score, int class_id) +{ + _tlwh_.resize(4); + _tlwh_.assign(tlwh.begin(), tlwh.end()); + + is_activated_ = false; + track_class_id_ = class_id; + state_ = TrackState::New; + + tlwh_.resize(4); + tlbr_.resize(4); + + static_tlwh(); + static_tlbr(); + frame_id_ = 0; + tracklet_len_ = 0; + this->score_ = score; + start_frame_ = 0; +} + +STrack::~STrack() = default; + +void STrack::activate(KalmanFilter& kalman_filter, int frame_id) +{ + this->kalman_filter_ = kalman_filter; + // this->track_class_id_ = STrack::next_id(); + + std::vector _tlwh_tmp(4); + _tlwh_tmp[0] = this->_tlwh_[0]; + _tlwh_tmp[1] = this->_tlwh_[1]; + _tlwh_tmp[2] = this->_tlwh_[2]; + _tlwh_tmp[3] = this->_tlwh_[3]; + std::vector xyah = tlwh_to_xyah(_tlwh_tmp); + DETECTBOX xyah_box; + xyah_box[0] = xyah[0]; + xyah_box[1] = xyah[1]; + xyah_box[2] = xyah[2]; + xyah_box[3] = xyah[3]; + auto mc = this->kalman_filter_.initiate(xyah_box); + this->mean_ = mc.first; + this->covariance_ = mc.second; + + static_tlwh(); + static_tlbr(); + + this->tracklet_len_ = 0; + this->state_ = TrackState::Tracked; + if (frame_id == 1) + { + this->is_activated_ = true; + } + // this->is_activated = true; + this->frame_id_ = frame_id; + this->start_frame_ = frame_id; +} + +void STrack::re_activate(STrack& new_track, int frame_id, bool new_id) +{ + std::vector xyah = tlwh_to_xyah(new_track.tlwh_); + DETECTBOX xyah_box; + xyah_box[0] = xyah[0]; + xyah_box[1] = xyah[1]; + xyah_box[2] = xyah[2]; + xyah_box[3] = xyah[3]; + auto mc = this->kalman_filter_.update(this->mean_, this->covariance_, xyah_box); + this->mean_ = mc.first; + this->covariance_ = mc.second; + + static_tlwh(); + static_tlbr(); + + this->tracklet_len_ = 0; + this->state_ = TrackState::Tracked; + this->is_activated_ = true; + this->frame_id_ = frame_id; + this->score_ = new_track.score_; + // if (new_id) + // this->track_class_id_ = next_id(); +} + +void STrack::update(STrack& new_track, int frame_id) +{ + this->frame_id_ = frame_id; + this->tracklet_len_++; + + std::vector xyah = tlwh_to_xyah(new_track.tlwh_); + DETECTBOX xyah_box; + xyah_box[0] = xyah[0]; + xyah_box[1] = xyah[1]; + xyah_box[2] = xyah[2]; + xyah_box[3] = xyah[3]; + + auto mc = this->kalman_filter_.update(this->mean_, this->covariance_, xyah_box); + this->mean_ = mc.first; + this->covariance_ = mc.second; + + static_tlwh(); + static_tlbr(); + + this->state_ = TrackState::Tracked; + this->is_activated_ = true; + + this->score_ = new_track.score_; + if (this->track_class_id_ == -1) + { + this->track_class_id_ = new_track.track_class_id_; + } +} + +void STrack::static_tlwh() +{ + if (this->state_ == TrackState::New) + { + tlwh_[0] = _tlwh_[0]; + tlwh_[1] = _tlwh_[1]; + tlwh_[2] = _tlwh_[2]; + tlwh_[3] = _tlwh_[3]; + + return; + } + + tlwh_[0] = mean_[0]; + tlwh_[1] = mean_[1]; + tlwh_[2] = mean_[2]; + tlwh_[3] = mean_[3]; + + tlwh_[2] *= tlwh_[3]; + tlwh_[0] -= tlwh_[2] / 2; + tlwh_[1] -= tlwh_[3] / 2; +} + +void STrack::static_tlbr() +{ + tlbr_.clear(); + tlbr_.assign(tlwh_.begin(), tlwh_.end()); + tlbr_[2] += tlbr_[0]; + tlbr_[3] += tlbr_[1]; +} + +std::vector STrack::tlwh_to_xyah(std::vector tlwh_tmp) +{ + std::vector tlwh_output = std::move(tlwh_tmp); + tlwh_output[0] += tlwh_output[2] / 2; + tlwh_output[1] += tlwh_output[3] / 2; + tlwh_output[2] /= tlwh_output[3]; + return tlwh_output; +} + +std::vector STrack::to_xyah() const +{ + return tlwh_to_xyah(tlwh_); +} + +std::vector STrack::tlbr_to_tlwh(std::vector& tlbr) +{ + tlbr[2] -= tlbr[0]; + tlbr[3] -= tlbr[1]; + return tlbr; +} + +void STrack::mark_lost() +{ + state_ = TrackState::Lost; +} + +void STrack::mark_removed() +{ + state_ = TrackState::Removed; +} + +int STrack::next_id() +{ + static int _count = 0; + _count++; + return _count; +} + +int STrack::end_frame() const +{ + return this->frame_id_; +} + +void STrack::multi_predict(std::vector& stracks, KalmanFilter& kalman_filter) +{ + for (auto& strack : stracks) + { + if (strack->state_ != TrackState::Tracked) + { + strack->mean_[7] = 0; + } + kalman_filter.predict(strack->mean_, strack->covariance_); + strack->static_tlwh(); + strack->static_tlbr(); + } +} +} // namespace rm_bytetrack \ No newline at end of file diff --git a/src/rm_bytetrack/kalmanFilter.cpp b/src/rm_bytetrack/kalmanFilter.cpp new file mode 100644 index 0000000..dda90bd --- /dev/null +++ b/src/rm_bytetrack/kalmanFilter.cpp @@ -0,0 +1,139 @@ +// +// Created by ywj on 24-1-14. +// + +#include "rm_bytetrack/kalmanFilter.h" +#include + +namespace rm_bytetrack +{ +const double KalmanFilter::chi2inv95[10] = { 0, 3.8415, 5.9915, 7.8147, 9.4877, 11.070, 12.592, 14.067, 15.507, 16.919 }; +KalmanFilter::KalmanFilter() +{ + int ndim = 4; + double dt = 1.; + + _motion_mat = Eigen::MatrixXf::Identity(8, 8); + for (int i = 0; i < ndim; i++) + { + _motion_mat(i, ndim + i) = dt; + } + _update_mat = Eigen::MatrixXf::Identity(4, 8); + + this->_std_weight_position = 1. / 20; + this->_std_weight_velocity = 1. / 160; +} + +KAL_DATA KalmanFilter::initiate(const DETECTBOX& measurement) +{ + DETECTBOX mean_pos = measurement; + DETECTBOX mean_vel; + for (int i = 0; i < 4; i++) + mean_vel(i) = 0; + + KAL_MEAN mean; + for (int i = 0; i < 8; i++) + { + if (i < 4) + mean(i) = mean_pos(i); + else + mean(i) = mean_vel(i - 4); + } + + KAL_MEAN std; + std(0) = 2 * _std_weight_position * measurement[3]; + std(1) = 2 * _std_weight_position * measurement[3]; + std(2) = 1e-2; + std(3) = 2 * _std_weight_position * measurement[3]; + std(4) = 10 * _std_weight_velocity * measurement[3]; + std(5) = 10 * _std_weight_velocity * measurement[3]; + std(6) = 1e-5; + std(7) = 10 * _std_weight_velocity * measurement[3]; + + KAL_MEAN tmp = std.array().square(); + KAL_COVA var = tmp.asDiagonal(); + return std::make_pair(mean, var); +} + +void KalmanFilter::predict(KAL_MEAN& mean, KAL_COVA& covariance) +{ + // revise the data; + DETECTBOX std_pos; + std_pos << _std_weight_position * mean(3), _std_weight_position * mean(3), 1e-2, _std_weight_position * mean(3); + DETECTBOX std_vel; + std_vel << _std_weight_velocity * mean(3), _std_weight_velocity * mean(3), 1e-5, _std_weight_velocity * mean(3); + KAL_MEAN tmp; + tmp.block<1, 4>(0, 0) = std_pos; + tmp.block<1, 4>(0, 4) = std_vel; + tmp = tmp.array().square(); + KAL_COVA motion_cov = tmp.asDiagonal(); + KAL_MEAN mean1 = this->_motion_mat * mean.transpose(); + KAL_COVA covariance1 = this->_motion_mat * covariance * (_motion_mat.transpose()); + covariance1 += motion_cov; + + mean = mean1; + covariance = covariance1; +} + +KAL_HDATA KalmanFilter::project(const KAL_MEAN& mean, const KAL_COVA& covariance) +{ + DETECTBOX std; + std << _std_weight_position * mean(3), _std_weight_position * mean(3), 1e-1, _std_weight_position * mean(3); + KAL_HMEAN mean1 = _update_mat * mean.transpose(); + KAL_HCOVA covariance1 = _update_mat * covariance * (_update_mat.transpose()); + Eigen::Matrix diag = std.asDiagonal(); + diag = diag.array().square().matrix(); + covariance1 += diag; + // covariance1.diagonal() << diag; + return std::make_pair(mean1, covariance1); +} + +KAL_DATA +KalmanFilter::update(const KAL_MEAN& mean, const KAL_COVA& covariance, const DETECTBOX& measurement) +{ + KAL_HDATA pa = project(mean, covariance); + KAL_HMEAN projected_mean = pa.first; + KAL_HCOVA projected_cov = pa.second; + + // chol_factor, lower = + // scipy.linalg.cho_factor(projected_cov, lower=True, check_finite=False) + // kalmain_gain = + // scipy.linalg.cho_solve((cho_factor, lower), + // np.dot(covariance, self._upadte_mat.T).T, + // check_finite=False).T + Eigen::Matrix B = (covariance * (_update_mat.transpose())).transpose(); + Eigen::Matrix kalman_gain = (projected_cov.llt().solve(B)).transpose(); // eg.8x4 + Eigen::Matrix innovation = measurement - projected_mean; // eg.1x4 + auto tmp = innovation * (kalman_gain.transpose()); + KAL_MEAN new_mean = (mean.array() + tmp.array()).matrix(); + KAL_COVA new_covariance = covariance - kalman_gain * projected_cov * (kalman_gain.transpose()); + return std::make_pair(new_mean, new_covariance); +} + +Eigen::Matrix KalmanFilter::gating_distance(const KAL_MEAN& mean, const KAL_COVA& covariance, + const std::vector& measurements, + bool only_position) +{ + KAL_HDATA pa = this->project(mean, covariance); + if (only_position) + { + printf("not implement!"); + exit(0); + } + KAL_HMEAN mean1 = pa.first; + KAL_HCOVA covariance1 = pa.second; + + // Eigen::Matrix d(size, 4); + DETECTBOXSS d(measurements.size(), 4); + int pos = 0; + for (DETECTBOX box : measurements) + { + d.row(pos++) = box - mean1; + } + Eigen::Matrix factor = covariance1.llt().matrixL(); + Eigen::Matrix z = factor.triangularView().solve(d).transpose(); + auto zz = ((z.array()) * (z.array())).matrix(); + auto square_maha = zz.colwise().sum(); + return square_maha; +} +} // namespace rm_bytetrack \ No newline at end of file diff --git a/src/rm_bytetrack/lapjv.cpp b/src/rm_bytetrack/lapjv.cpp new file mode 100644 index 0000000..ba0d305 --- /dev/null +++ b/src/rm_bytetrack/lapjv.cpp @@ -0,0 +1,372 @@ +// +// Created by ywj on 24-1-14. +// + +#include +#include + +#include "rm_bytetrack/lapjv.h" + +/** Column-reduction and reduction transfer for a dense cost matrix. + */ +int_t _ccrrt_dense(const uint_t n, cost_t* cost[], int_t* free_rows, int_t* x, int_t* y, cost_t* v) +{ + int_t n_free_rows; + boolean* unique; + + for (uint_t i = 0; i < n; i++) + { + x[i] = -1; + v[i] = LARGE; + y[i] = 0; + } + for (uint_t i = 0; i < n; i++) + { + for (uint_t j = 0; j < n; j++) + { + const cost_t c = cost[i][j]; + if (c < v[j]) + { + v[j] = c; + y[j] = i; + } + PRINTF("i=%d, j=%d, c[i,j]=%f, v[j]=%f y[j]=%d\n", i, j, c, v[j], y[j]); + } + } + PRINT_COST_ARRAY(v, n); + PRINT_INDEX_ARRAY(y, n); + NEW(unique, boolean, n); + memset(unique, TRUE, n); + { + int_t j = n; + do + { + j--; + const int_t i = y[j]; + if (x[i] < 0) + { + x[i] = j; + } + else + { + unique[i] = FALSE; + y[j] = -1; + } + } while (j > 0); + } + n_free_rows = 0; + for (uint_t i = 0; i < n; i++) + { + if (x[i] < 0) + { + free_rows[n_free_rows++] = i; + } + else if (unique[i]) + { + const int_t j = x[i]; + cost_t min = LARGE; + for (uint_t j2 = 0; j2 < n; j2++) + { + if (j2 == (uint_t)j) + { + continue; + } + const cost_t c = cost[i][j2] - v[j2]; + if (c < min) + { + min = c; + } + } + PRINTF("v[%d] = %f - %f\n", j, v[j], min); + v[j] -= min; + } + } + FREE(unique); + return n_free_rows; +} + +/** Augmenting row reduction for a dense cost matrix. + */ +int_t _carr_dense(const uint_t n, cost_t* cost[], const uint_t n_free_rows, int_t* free_rows, int_t* x, int_t* y, + cost_t* v) +{ + uint_t current = 0; + int_t new_free_rows = 0; + uint_t rr_cnt = 0; + PRINT_INDEX_ARRAY(x, n); + PRINT_INDEX_ARRAY(y, n); + PRINT_COST_ARRAY(v, n); + PRINT_INDEX_ARRAY(free_rows, n_free_rows); + while (current < n_free_rows) + { + int_t i0; + int_t j1, j2; + cost_t v1, v2, v1_new; + boolean v1_lowers; + + rr_cnt++; + PRINTF("current = %d rr_cnt = %d\n", current, rr_cnt); + const int_t free_i = free_rows[current++]; + j1 = 0; + v1 = cost[free_i][0] - v[0]; + j2 = -1; + v2 = LARGE; + for (uint_t j = 1; j < n; j++) + { + PRINTF("%d = %f %d = %f\n", j1, v1, j2, v2); + const cost_t c = cost[free_i][j] - v[j]; + if (c < v2) + { + if (c >= v1) + { + v2 = c; + j2 = j; + } + else + { + v2 = v1; + v1 = c; + j2 = j1; + j1 = j; + } + } + } + i0 = y[j1]; + v1_new = v[j1] - (v2 - v1); + v1_lowers = v1_new < v[j1]; + PRINTF("%d %d 1=%d,%f 2=%d,%f v1'=%f(%d,%g) \n", free_i, i0, j1, v1, j2, v2, v1_new, v1_lowers, v[j1] - v1_new); + if (rr_cnt < current * n) + { + if (v1_lowers) + { + v[j1] = v1_new; + } + else if (i0 >= 0 && j2 >= 0) + { + j1 = j2; + i0 = y[j2]; + } + if (i0 >= 0) + { + if (v1_lowers) + { + free_rows[--current] = i0; + } + else + { + free_rows[new_free_rows++] = i0; + } + } + } + else + { + PRINTF("rr_cnt=%d >= %d (current=%d * n=%d)\n", rr_cnt, current * n, current, n); + if (i0 >= 0) + { + free_rows[new_free_rows++] = i0; + } + } + x[free_i] = j1; + y[j1] = free_i; + } + return new_free_rows; +} + +/** Find columns with minimum d[j] and put them on the SCAN list. + */ +uint_t _find_dense(const uint_t n, uint_t lo, cost_t* d, int_t* cols, int_t* y) +{ + uint_t hi = lo + 1; + cost_t mind = d[cols[lo]]; + for (uint_t k = hi; k < n; k++) + { + int_t j = cols[k]; + if (d[j] <= mind) + { + if (d[j] < mind) + { + hi = lo; + mind = d[j]; + } + cols[k] = cols[hi]; + cols[hi++] = j; + } + } + return hi; +} + +// Scan all columns in TODO starting from arbitrary column in SCAN +// and try to decrease d of the TODO columns using the SCAN column. +int_t _scan_dense(const uint_t n, cost_t* cost[], uint_t* plo, uint_t* phi, cost_t* d, int_t* cols, int_t* pred, + int_t* y, cost_t* v) +{ + uint_t lo = *plo; + uint_t hi = *phi; + cost_t h, cred_ij; + + while (lo != hi) + { + int_t j = cols[lo++]; + const int_t i = y[j]; + const cost_t mind = d[j]; + h = cost[i][j] - v[j] - mind; + PRINTF("i=%d j=%d h=%f\n", i, j, h); + // For all columns in TODO + for (uint_t k = hi; k < n; k++) + { + j = cols[k]; + cred_ij = cost[i][j] - v[j] - h; + if (cred_ij < d[j]) + { + d[j] = cred_ij; + pred[j] = i; + if (cred_ij == mind) + { + if (y[j] < 0) + { + return j; + } + cols[k] = cols[hi]; + cols[hi++] = j; + } + } + } + } + *plo = lo; + *phi = hi; + return -1; +} + +/** Single iteration of modified Dijkstra shortest path algorithm as explained in the JV paper. + * + * This is a dense matrix version. + * + * \return The closest free column index. + */ +int_t find_path_dense(const uint_t n, cost_t* cost[], const int_t start_i, int_t* y, cost_t* v, int_t* pred) +{ + uint_t lo = 0, hi = 0; + int_t final_j = -1; + uint_t n_ready = 0; + int_t* cols; + cost_t* d; + + NEW(cols, int_t, n); + NEW(d, cost_t, n); + + for (uint_t i = 0; i < n; i++) + { + cols[i] = i; + pred[i] = start_i; + d[i] = cost[start_i][i] - v[i]; + } + PRINT_COST_ARRAY(d, n); + while (final_j == -1) + { + // No columns left on the SCAN list. + if (lo == hi) + { + PRINTF("%d..%d -> find\n", lo, hi); + n_ready = lo; + hi = _find_dense(n, lo, d, cols, y); + PRINTF("check %d..%d\n", lo, hi); + PRINT_INDEX_ARRAY(cols, n); + for (uint_t k = lo; k < hi; k++) + { + const int_t j = cols[k]; + if (y[j] < 0) + { + final_j = j; + } + } + } + if (final_j == -1) + { + PRINTF("%d..%d -> scan\n", lo, hi); + final_j = _scan_dense(n, cost, &lo, &hi, d, cols, pred, y, v); + PRINT_COST_ARRAY(d, n); + PRINT_INDEX_ARRAY(cols, n); + PRINT_INDEX_ARRAY(pred, n); + } + } + + PRINTF("found final_j=%d\n", final_j); + PRINT_INDEX_ARRAY(cols, n); + { + const cost_t mind = d[cols[lo]]; + for (uint_t k = 0; k < n_ready; k++) + { + const int_t j = cols[k]; + v[j] += d[j] - mind; + } + } + + FREE(cols); + FREE(d); + + return final_j; +} + +/** Augment for a dense cost matrix. + */ +int_t _ca_dense(const uint_t n, cost_t* cost[], const uint_t n_free_rows, int_t* free_rows, int_t* x, int_t* y, + cost_t* v) +{ + int_t* pred; + + NEW(pred, int_t, n); + + for (int_t* pfree_i = free_rows; pfree_i < free_rows + n_free_rows; pfree_i++) + { + int_t i = -1, j; + uint_t k = 0; + + PRINTF("looking at free_i=%d\n", *pfree_i); + j = find_path_dense(n, cost, *pfree_i, y, v, pred); + ASSERT(j >= 0); + ASSERT(j < n); + while (i != *pfree_i) + { + PRINTF("augment %d\n", j); + PRINT_INDEX_ARRAY(pred, n); + i = pred[j]; + PRINTF("y[%d]=%d -> %d\n", j, y[j], i); + y[j] = i; + PRINT_INDEX_ARRAY(x, n); + SWAP_INDICES(j, x[i]); + k++; + if (k >= n) + { + ASSERT(FALSE); + } + } + } + FREE(pred); + return 0; +} + +/** Solve dense sparse LAP. + */ +int lapjv_internal(const uint_t n, cost_t* cost[], int_t* x, int_t* y) +{ + int ret; + int_t* free_rows; + cost_t* v; + + NEW(free_rows, int_t, n); + NEW(v, cost_t, n); + ret = _ccrrt_dense(n, cost, free_rows, x, y, v); + int i = 0; + while (ret > 0 && i < 2) + { + ret = _carr_dense(n, cost, ret, free_rows, x, y, v); + i++; + } + if (ret > 0) + { + ret = _ca_dense(n, cost, ret, free_rows, x, y, v); + } + FREE(v); + FREE(free_rows); + return ret; +} From 7f3c96827f6cd99c3f6782fb924340776a4df825 Mon Sep 17 00:00:00 2001 From: asitaku <3482197316@qq.com> Date: Sun, 21 Jan 2024 10:21:06 +0800 Subject: [PATCH 3/5] debug --- cfg/dynamic.cfg | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cfg/dynamic.cfg b/cfg/dynamic.cfg index 36c4233..bd24084 100644 --- a/cfg/dynamic.cfg +++ b/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 \ No newline at end of file +u From 84053d66cb3f35602a187461650c0d70709c43ce Mon Sep 17 00:00:00 2001 From: asitaku <92627302+asitaku@users.noreply.github.com> Date: Wed, 24 Jan 2024 10:41:24 +0800 Subject: [PATCH 4/5] Update postprocess.cpp --- utils/postprocess.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/postprocess.cpp b/utils/postprocess.cpp index c2a7b5f..2213eaf 100644 --- a/utils/postprocess.cpp +++ b/utils/postprocess.cpp @@ -33,9 +33,9 @@ cv::Rect get_rect(cv::Mat& img, float bbox[4]) { if (rect.tl().y < 0) rect.y = 0; if (rect.br().x > img.cols) - rect.x = img.cols - rect.width - 1; + rect.width = img.cols - rect.x; if (rect.br().y > img.rows) - rect.y = img.rows - rect.height - 1; + rect.height = img.rows - rect.y; return rect; } From ada6fd008cfd4c84fdd0be4d8279fba46ed5a812 Mon Sep 17 00:00:00 2001 From: asitaku <92627302+asitaku@users.noreply.github.com> Date: Wed, 24 Jan 2024 10:42:29 +0800 Subject: [PATCH 5/5] Update STrack.cpp --- src/rm_bytetrack/STrack.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/rm_bytetrack/STrack.cpp b/src/rm_bytetrack/STrack.cpp index fd78648..56072b2 100644 --- a/src/rm_bytetrack/STrack.cpp +++ b/src/rm_bytetrack/STrack.cpp @@ -111,7 +111,7 @@ void STrack::update(STrack& new_track, int frame_id) this->is_activated_ = true; this->score_ = new_track.score_; - if (this->track_class_id_ == -1) + if (new_track.track_class_id_ != -1) { this->track_class_id_ = new_track.track_class_id_; } @@ -203,4 +203,4 @@ void STrack::multi_predict(std::vector& stracks, KalmanFilter& kalman_f strack->static_tlbr(); } } -} // namespace rm_bytetrack \ No newline at end of file +} // namespace rm_bytetrack