diff --git a/botsort/CMakeLists.txt b/botsort/CMakeLists.txt index 1aa53e9..f64966a 100644 --- a/botsort/CMakeLists.txt +++ b/botsort/CMakeLists.txt @@ -20,7 +20,10 @@ file(GLOB_RECURSE SOURCES "${PROJECT_SOURCE_DIR}/src/*.cpp") add_library(${PROJECT_NAME} SHARED ${SOURCES}) # Set include directories -target_include_directories(${PROJECT_NAME} PUBLIC ${PROJECT_SOURCE_DIR}/include) +target_include_directories(${PROJECT_NAME} PUBLIC + $ + $ +) # Find and link OpenCV find_package(OpenCV REQUIRED) diff --git a/botsort/include/BoTSORT.h b/botsort/include/BoTSORT.h old mode 100644 new mode 100755 index 9f7f7e3..684eb20 --- a/botsort/include/BoTSORT.h +++ b/botsort/include/BoTSORT.h @@ -1,19 +1,26 @@ #pragma once #include +#include #include "GlobalMotionCompensation.h" +#include "GmcParams.h" #include "ReID.h" +#include "ReIDParams.h" +#include "TrackerParams.h" #include "track.h" +template +using Config = std::variant; class BoTSORT { public: - explicit BoTSORT(const std::string &tracker_config_path, - const std::string &gmc_config_path = "", - const std::string &reid_config_path = "", + explicit BoTSORT(const Config &tracker_config, + const Config &gmc_config = {}, + const Config &reid_config = {}, const std::string &reid_onnx_model_path = ""); + ~BoTSORT() = default; @@ -79,14 +86,12 @@ class BoTSORT std::vector> &tracks_list_a, std::vector> &tracks_list_b); - /** - * @brief Load tracker parameters from the given config file + * @brief Load tracker parameters from the given config * - * @param config_path Path to the config directory + * @param config Configuration to load */ - void _load_params_from_config(const std::string &config_path); - + void _load_params_from_config(const TrackerParams &config); private: std::string _gmc_method_name; diff --git a/botsort/include/GlobalMotionCompensation.h b/botsort/include/GlobalMotionCompensation.h index e3d1918..c4f0247 100644 --- a/botsort/include/GlobalMotionCompensation.h +++ b/botsort/include/GlobalMotionCompensation.h @@ -6,6 +6,7 @@ // .clang-format off #include "DataType.h" +#include "GmcParams.h" // .clang-format on #include @@ -17,16 +18,6 @@ #include -enum GMC_Method -{ - ORB = 0, - ECC, - SparseOptFlow, - OptFlowModified, - OpenCV_VideoStab -}; - - class GMC_Algorithm { public: @@ -39,13 +30,13 @@ class GMC_Algorithm class ORB_GMC : public GMC_Algorithm { public: - explicit ORB_GMC(const std::string &config_path); + explicit ORB_GMC(const ORB_Params &orb_config); HomographyMatrix apply(const cv::Mat &frame_raw, const std::vector &detections) override; private: - void _load_params_from_config(const std::string &config_path); + void _load_params_from_config(const ORB_Params &orb_config); private: @@ -67,13 +58,13 @@ class ORB_GMC : public GMC_Algorithm class ECC_GMC : public GMC_Algorithm { public: - explicit ECC_GMC(const std::string &config_path); + explicit ECC_GMC(const ECC_Params &config); HomographyMatrix apply(const cv::Mat &frame_raw, const std::vector &detections) override; private: - void _load_params_from_config(const std::string &config_dir); + void _load_params_from_config(const ECC_Params &config); private: @@ -91,13 +82,13 @@ class ECC_GMC : public GMC_Algorithm class SparseOptFlow_GMC : public GMC_Algorithm { public: - explicit SparseOptFlow_GMC(const std::string &config_path); + explicit SparseOptFlow_GMC(const SparseOptFlow_Params &config); HomographyMatrix apply(const cv::Mat &frame_raw, const std::vector &detections) override; private: - void _load_params_from_config(const std::string &config_dir); + void _load_params_from_config(const SparseOptFlow_Params &config); private: @@ -119,13 +110,13 @@ class SparseOptFlow_GMC : public GMC_Algorithm class OptFlowModified_GMC : public GMC_Algorithm { public: - explicit OptFlowModified_GMC(const std::string &config_path); + explicit OptFlowModified_GMC(const OptFlowModified_Params &config); HomographyMatrix apply(const cv::Mat &frame_raw, const std::vector &detections) override; private: - void _load_params_from_config(const std::string &config_dir); + void _load_params_from_config(const OptFlowModified_Params &config); private: @@ -137,13 +128,13 @@ class OptFlowModified_GMC : public GMC_Algorithm class OpenCV_VideoStab_GMC : public GMC_Algorithm { public: - explicit OpenCV_VideoStab_GMC(const std::string &config_path); + explicit OpenCV_VideoStab_GMC(const OpenCV_VideoStab_GMC_Params &config); HomographyMatrix apply(const cv::Mat &frame_raw, const std::vector &detections) override; private: - void _load_params_from_config(const std::string &config_dir); + void _load_params_from_config(const OpenCV_VideoStab_GMC_Params &config); private: @@ -167,11 +158,9 @@ class GlobalMotionCompensation /** * @brief Construct a new Global Motion Compensation object * - * @param method GMC_Method enum member for GMC algorithm to use - * @param config_dir Directory containing config files for GMC algorithm + * @param gmc_params Paramerters for GMC algorithm */ - explicit GlobalMotionCompensation(GMC_Method method, - const std::string &config_path); + explicit GlobalMotionCompensation(const GMC_Params &gmc_params); ~GlobalMotionCompensation() = default; /** diff --git a/botsort/include/GmcParams.h b/botsort/include/GmcParams.h new file mode 100644 index 0000000..4257e71 --- /dev/null +++ b/botsort/include/GmcParams.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include + +enum GMC_Method +{ + ORB = 0, + ECC, + SparseOptFlow, + OptFlowModified, + OpenCV_VideoStab +}; + +struct ORB_Params +{ + float downscale{2.f}; + float inlier_ratio{0.5f}; + float ransac_conf{.99f}; + long ransac_max_iters{500}; +}; + +struct ECC_Params +{ + float downscale{5.f}; + long max_iterations{100}; + float termination_eps{1e-6}; +}; + +struct SparseOptFlow_Params +{ + long max_corners{1000}; + long block_size{3}; + long ransac_max_iters{500}; + double quality_level{0.01}; + double k{0.04}; + double min_distance{1.0}; + float downscale{2.0f}; + float inlier_ratio{0.5f}; + float ransac_conf{0.99f}; + bool use_harris_detector{false}; +}; + +struct OptFlowModified_Params +{ + float downscale{2.0f}; +}; + +struct OpenCV_VideoStab_GMC_Params +{ + float downscale{2.0f}; + float num_features{4000}; + bool detection_masking{true}; +}; + +struct GMC_Params +{ + using MethodParams = + std::variant; + + GMC_Method method_; + MethodParams method_params_; + + static GMC_Params load_config(GMC_Method method, + const std::string &config_path); +}; \ No newline at end of file diff --git a/botsort/include/INIReader.h b/botsort/include/INIReader.h index c5d0356..0ab5219 100644 --- a/botsort/include/INIReader.h +++ b/botsort/include/INIReader.h @@ -359,45 +359,74 @@ class INIReader } // Return the list of sections found in ini file - const std::set &Sections() const; + [[nodiscard]] const std::set &Sections() const; // Get a string value from INI file, returning default_value if not found. - std::string Get(const std::string §ion, const std::string &name, - const std::string &default_value) const; + [[nodiscard]] std::string Get(const std::string §ion, + const std::string &name, + const std::string &default_value) const; // Get a string value from INI file, return nullopt if not found. - std::optional Get(const std::string §ion, - const std::string &name) const; + [[nodiscard]] std::optional Get(const std::string §ion, + const std::string &name) const; // Get an integer (long) value from INI file, returning default_value if // not found or not a valid integer (decimal "1234", "-1234", or hex "0x4d2"). - long GetInteger(const std::string §ion, const std::string &name, - long default_value) const; + [[nodiscard]] long GetInteger(const std::string §ion, + const std::string &name, + long default_value) const; // Get a real (floating point double) value from INI file, returning // default_value if not found or not a valid floating point value // according to strtod(). - double GetReal(const std::string §ion, const std::string &name, - double default_value) const; + [[nodiscard]] double GetReal(const std::string §ion, + const std::string &name, + double default_value) const; // Get a single precision floating point number value from INI file, returning // default_value if not found or not a valid floating point value // according to strtof(). - float GetFloat(const std::string §ion, const std::string &name, - float default_value) const; + [[nodiscard]] float GetFloat(const std::string §ion, + const std::string &name, + float default_value) const; // Get a boolean value from INI file, returning default_value if not found or if // not a valid true/false value. Valid true values are "true", "yes", "on", "1", // and valid false values are "false", "no", "off", "0" (not case sensitive). - bool GetBoolean(const std::string §ion, const std::string &name, - bool default_value) const; + [[nodiscard]] bool GetBoolean(const std::string §ion, + const std::string &name, + bool default_value) const; template - std::vector GetList(const std::string §ion, - const std::string &name) const; - + [[nodiscard]] std::vector GetList(const std::string §ion, + const std::string &name) const; + + // Load a string value from INI file, returning its value if successfully + // loaded, or leave value unchanged + void LoadString(const std::string §ion, const std::string &name, + std::string &value) const; + + // Load an integer (long) value from INI file, returning its value if + // successfully loaded, or leave value unchanged + void LoadInteger(const std::string §ion, const std::string &name, + long &value) const; + + // Load a real (floating point double) value from INI file, returning its + // value if successfully loaded, or leave value unchanged + void LoadReal(const std::string §ion, const std::string &name, + double &value) const; + + // Load a single precision floating point number value from INI file, returning + // its value if successfully loaded, or leave value unchanged + void LoadFloat(const std::string §ion, const std::string &name, + float &value) const; + + // Load a boolean value from INI file, returning its value if successfully + // loaded, or leave value unchanged + void LoadBoolean(const std::string §ion, const std::string &name, + bool &value) const; protected: int _error; @@ -568,4 +597,35 @@ std::vector INIReader::GetList(const std::string §ion, return list_items; } +inline void INIReader::LoadString(const std::string §ion, + const std::string &name, + std::string &value) const +{ + value = Get(section, name, value); +} + +inline void INIReader::LoadInteger(const std::string §ion, + const std::string &name, long &value) const +{ + value = GetInteger(section, name, value); +} + +inline void INIReader::LoadReal(const std::string §ion, + const std::string &name, double &value) const +{ + value = GetReal(section, name, value); +} + +inline void INIReader::LoadFloat(const std::string §ion, + const std::string &name, float &value) const +{ + value = GetFloat(section, name, value); +} + +inline void INIReader::LoadBoolean(const std::string §ion, + const std::string &name, bool &value) const +{ + value = GetBoolean(section, name, value); +} + #endif// __INIREADER__ \ No newline at end of file diff --git a/botsort/include/ReID.h b/botsort/include/ReID.h index 96fcf26..9543e32 100644 --- a/botsort/include/ReID.h +++ b/botsort/include/ReID.h @@ -3,13 +3,13 @@ #include #include "DataType.h" +#include "ReIDParams.h" #include "TRT_InferenceEngine/TensorRT_InferenceEngine.h" class ReIDModel { public: - ReIDModel(const std::string &config_path, - const std::string &onnx_model_path); + ReIDModel(const ReIDParams ¶ms, const std::string &onnx_model_path); ~ReIDModel() = default; void pre_process(cv::Mat &image); @@ -21,7 +21,7 @@ class ReIDModel } private: - void _load_params_from_config(const std::string &config_path); + void _load_params_from_config(const ReIDParams ¶ms); private: diff --git a/botsort/include/ReIDParams.h b/botsort/include/ReIDParams.h new file mode 100644 index 0000000..8a6cc09 --- /dev/null +++ b/botsort/include/ReIDParams.h @@ -0,0 +1,21 @@ +#pragma once + +#include +#include + +struct ReIDParams +{ + long gpu_id{0}; + std::string distance_metric{"euclidean"}; + long trt_logging_level{1}; + long batch_size{1}; + std::string input_layer_name{""}; + std::vector input_layer_dimensions; + std::vector output_layer_names; + + bool enable_fp16{true}; + bool enable_tf32{true}; + bool swap_rb{false}; + + static ReIDParams load_config(const std::string &config_path); +}; diff --git a/botsort/include/TRT_InferenceEngine/TensorRT_InferenceEngine.h b/botsort/include/TRT_InferenceEngine/TensorRT_InferenceEngine.h index ebde60f..b11488f 100644 --- a/botsort/include/TRT_InferenceEngine/TensorRT_InferenceEngine.h +++ b/botsort/include/TRT_InferenceEngine/TensorRT_InferenceEngine.h @@ -2,11 +2,11 @@ #include #include +#include #include #include #include #include -#include #include #include @@ -55,6 +55,7 @@ struct TRTOptimizerParams { TRTOptimizerParams() = default; + int gpu_id = 0; int batch_size = 1; bool fp16 = true; bool int8 = false; @@ -69,7 +70,8 @@ struct TRTOptimizerParams std::string toStr() { - std::string str = "batch_size: " + std::to_string(batch_size) + "\n"; + std::string str = "gpu_id:" + std::to_string(gpu_id) + "\n"; + str += "batch_size: " + std::to_string(batch_size) + "\n"; str += "fp16: " + std::to_string(fp16) + "\n"; str += "int8: " + std::to_string(int8) + "\n"; str += "tf32: " + std::to_string(tf32) + "\n"; diff --git a/botsort/include/TrackerParams.h b/botsort/include/TrackerParams.h new file mode 100644 index 0000000..a7bd1fc --- /dev/null +++ b/botsort/include/TrackerParams.h @@ -0,0 +1,21 @@ +#pragma once + +#include + +struct TrackerParams +{ + bool reid_enabled{false}; + bool gmc_enabled{false}; + float track_high_thresh{0.6F}; + float track_low_thresh{0.1F}; + float new_track_thresh{0.7F}; + long track_buffer{30}; + float match_thresh{0.7F}; + float proximity_thresh{0.5F}; + float appearance_thresh{0.25F}; + std::string gmc_method_name{"sparseOptFlow"}; + long frame_rate{30}; + float lambda{0.985F}; + + static TrackerParams load_config(const std::string &config_path); +}; diff --git a/botsort/include/track.h b/botsort/include/track.h index b5b09e3..2db131a 100644 --- a/botsort/include/track.h +++ b/botsort/include/track.h @@ -77,6 +77,13 @@ class Track */ float get_score() const; + /** + * @brief Get the class ID of the track + * + * @return uint8_t Class ID of the track + */ + uint8_t get_class_id() const; + /** * @brief Activates the track * diff --git a/botsort/src/BoTSORT.cpp b/botsort/src/BoTSORT.cpp old mode 100644 new mode 100755 index 9cde758..e6c84f1 --- a/botsort/src/BoTSORT.cpp +++ b/botsort/src/BoTSORT.cpp @@ -10,12 +10,51 @@ #include "matching.h" #include "profiler.h" -BoTSORT::BoTSORT(const std::string &tracker_config_path, - const std::string &gmc_config_path, - const std::string &reid_config_path, +namespace +{ +template +bool requires_load(const Config &config) +{ + return std::holds_alternative(config) && + !std::get(config).empty(); +} + +template +bool not_empty(const Config &config) +{ + bool has_config = !std::holds_alternative(config); + bool is_non_empty = !(std::holds_alternative(config) && + std::get(config).empty()); + + return has_config && is_non_empty; +} + +template +T fetch_config(const Config &config, + std::function loader) +{ + if (std::holds_alternative(config)) + { + return std::get(config); + } + + if (requires_load(config)) + { + return loader(std::get(config)); + } + + throw std::runtime_error("Config is empty"); +} +}// namespace + +BoTSORT::BoTSORT(const Config &tracker_config, + const Config &gmc_config, + const Config &reid_config, const std::string &reid_onnx_model_path) { - _load_params_from_config(tracker_config_path); + auto tracker_params = fetch_config( + tracker_config, TrackerParams::load_config); + _load_params_from_config(tracker_params); // Tracker module _frame_id = 0; @@ -26,22 +65,31 @@ BoTSORT::BoTSORT(const std::string &tracker_config_path, // Re-ID module, load visual feature extractor here - if (_reid_enabled && reid_config_path.size() > 0 && + if (_reid_enabled && not_empty(reid_config) && reid_onnx_model_path.size() > 0) - _reid_model = std::make_unique(reid_config_path, - reid_onnx_model_path); + { + auto reid_params = + fetch_config(reid_config, ReIDParams::load_config); + _reid_model = + std::make_unique(reid_params, reid_onnx_model_path); + } else { std::cout << "Re-ID module disabled" << std::endl; _reid_enabled = false; } - // Global motion compensation module - if (_gmc_enabled && gmc_config_path.size() > 0) - _gmc_algo = std::make_unique( - GlobalMotionCompensation::GMC_method_map[_gmc_method_name], - gmc_config_path); + if (_gmc_enabled && not_empty(gmc_config)) + { + auto gmc_params = fetch_config< + GMC_Params>(gmc_config, [this](const std::string &config_path) { + return GMC_Params::load_config( + GlobalMotionCompensation::GMC_method_map[_gmc_method_name], + config_path); + }); + _gmc_algo = std::make_unique(gmc_params); + } else { std::cout << "GMC disabled" << std::endl; @@ -506,39 +554,18 @@ void BoTSORT::_remove_duplicate_tracks( } } - -void BoTSORT::_load_params_from_config(const std::string &config_path) +void BoTSORT::_load_params_from_config(const TrackerParams &config) { - const std::string tracker_name = "BoTSORT"; - - INIReader tracker_config(config_path); - if (tracker_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _reid_enabled = - tracker_config.GetBoolean(tracker_name, "enable_reid", false); - _gmc_enabled = tracker_config.GetBoolean(tracker_name, "enable_gmc", false); - _track_high_thresh = - tracker_config.GetFloat(tracker_name, "track_high_thresh", 0.6F); - _track_low_thresh = - tracker_config.GetFloat(tracker_name, "track_low_thresh", 0.1F); - _new_track_thresh = - tracker_config.GetFloat(tracker_name, "new_track_thresh", 0.7F); - - _track_buffer = tracker_config.GetInteger(tracker_name, "track_buffer", 30); - - _match_thresh = tracker_config.GetFloat(tracker_name, "match_thresh", 0.7F); - _proximity_thresh = - tracker_config.GetFloat(tracker_name, "proximity_thresh", 0.5F); - _appearance_thresh = - tracker_config.GetFloat(tracker_name, "appearance_thresh", 0.25F); - - _gmc_method_name = - tracker_config.Get(tracker_name, "gmc_method", "sparseOptFlow"); - - _frame_rate = tracker_config.GetInteger(tracker_name, "frame_rate", 30); - _lambda = tracker_config.GetFloat(tracker_name, "lambda", 0.985F); + _reid_enabled = config.reid_enabled; + _gmc_enabled = config.gmc_enabled; + _track_high_thresh = config.track_high_thresh; + _track_low_thresh = config.track_low_thresh; + _new_track_thresh = config.new_track_thresh; + _track_buffer = config.track_buffer; + _match_thresh = config.match_thresh; + _proximity_thresh = config.proximity_thresh; + _appearance_thresh = config.appearance_thresh; + _gmc_method_name = config.gmc_method_name; + _frame_rate = config.frame_rate; + _lambda = config.lambda; } \ No newline at end of file diff --git a/botsort/src/GlobalMotionCompensation.cpp b/botsort/src/GlobalMotionCompensation.cpp index ff7e78f..d3e53c0 100644 --- a/botsort/src/GlobalMotionCompensation.cpp +++ b/botsort/src/GlobalMotionCompensation.cpp @@ -3,8 +3,6 @@ #include #include -#include "INIReader.h" - std::map GlobalMotionCompensation::GMC_method_map = { {"orb", GMC_Method::ORB}, {"ecc", GMC_Method::ECC}, @@ -14,33 +12,40 @@ std::map GlobalMotionCompensation::GMC_method_map = { }; -GlobalMotionCompensation::GlobalMotionCompensation( - GMC_Method method, const std::string &config_path) +GlobalMotionCompensation::GlobalMotionCompensation(const GMC_Params &gmc_params) { + GMC_Method method = gmc_params.method_; + auto ¶ms = gmc_params.method_params_; + if (method == GMC_Method::ORB) { std::cout << "Using ORB for GMC" << std::endl; - _gmc_algorithm = std::make_unique(config_path); + _gmc_algorithm = + std::make_unique(std::get(params)); } else if (method == GMC_Method::ECC) { std::cout << "Using ECC for GMC" << std::endl; - _gmc_algorithm = std::make_unique(config_path); + _gmc_algorithm = + std::make_unique(std::get(params)); } else if (method == GMC_Method::SparseOptFlow) { std::cout << "Using SparseOptFlow for GMC" << std::endl; - _gmc_algorithm = std::make_unique(config_path); + _gmc_algorithm = std::make_unique( + std::get(params)); } else if (method == GMC_Method::OptFlowModified) { std::cout << "Using OptFlowModified for GMC" << std::endl; - _gmc_algorithm = std::make_unique(config_path); + _gmc_algorithm = std::make_unique( + std::get(params)); } else if (method == GMC_Method::OpenCV_VideoStab) { std::cout << "Using OpenCV_VideoStab for GMC" << std::endl; - _gmc_algorithm = std::make_unique(config_path); + _gmc_algorithm = std::make_unique( + std::get(params)); } else { @@ -59,9 +64,9 @@ GlobalMotionCompensation::apply(const cv::Mat &frame, // ORB -ORB_GMC::ORB_GMC(const std::string &config_path) +ORB_GMC::ORB_GMC(const ORB_Params &config) { - _load_params_from_config(config_path); + _load_params_from_config(config); _detector = cv::FastFeatureDetector::create(); _extractor = cv::ORB::create(); @@ -69,20 +74,12 @@ ORB_GMC::ORB_GMC(const std::string &config_path) } -void ORB_GMC::_load_params_from_config(const std::string &config_path) +void ORB_GMC::_load_params_from_config(const ORB_Params &config) { - INIReader gmc_config(config_path); - if (gmc_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _downscale = gmc_config.GetFloat(_algo_name, "downscale", 2.0); - _inlier_ratio = gmc_config.GetFloat(_algo_name, "inlier_ratio", 0.5); - _ransac_conf = gmc_config.GetFloat(_algo_name, "ransac_conf", 0.99); - _ransac_max_iters = - gmc_config.GetInteger(_algo_name, "ransac_max_iters", 500); + _downscale = config.downscale; + _inlier_ratio = config.inlier_ratio; + _ransac_conf = config.ransac_conf; + _ransac_max_iters = static_cast(config.ransac_max_iters); } @@ -279,9 +276,9 @@ HomographyMatrix ORB_GMC::apply(const cv::Mat &frame_raw, // ECC -ECC_GMC::ECC_GMC(const std::string &config_path) +ECC_GMC::ECC_GMC(const ECC_Params &config) { - _load_params_from_config(config_path); + _load_params_from_config(config); _termination_criteria = cv::TermCriteria(cv::TermCriteria::EPS | cv::TermCriteria::COUNT, @@ -289,18 +286,11 @@ ECC_GMC::ECC_GMC(const std::string &config_path) } -void ECC_GMC::_load_params_from_config(const std::string &config_path) +void ECC_GMC::_load_params_from_config(const ECC_Params &config) { - INIReader gmc_config(config_path); - if (gmc_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _downscale = gmc_config.GetFloat(_algo_name, "downscale", 5.0F); - _max_iterations = gmc_config.GetInteger(_algo_name, "max_iterations", 100); - _termination_eps = gmc_config.GetFloat(_algo_name, "termination_eps", 1e-6); + _downscale = config.downscale; + _max_iterations = static_cast(config.max_iterations); + _termination_eps = static_cast(config.termination_eps); } @@ -361,37 +351,29 @@ HomographyMatrix ECC_GMC::apply(const cv::Mat &frame_raw, // Optical Flow -SparseOptFlow_GMC::SparseOptFlow_GMC(const std::string &config_path) +SparseOptFlow_GMC::SparseOptFlow_GMC(const SparseOptFlow_Params &config) { - _load_params_from_config(config_path); + _load_params_from_config(config); } -void SparseOptFlow_GMC::_load_params_from_config(const std::string &config_path) +void SparseOptFlow_GMC::_load_params_from_config( + const SparseOptFlow_Params &config) { - INIReader gmc_config(config_path); - if (gmc_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _useHarrisDetector = - gmc_config.GetBoolean(_algo_name, "use_harris_detector", false); + _useHarrisDetector = config.use_harris_detector; - _maxCorners = gmc_config.GetInteger(_algo_name, "max_corners", 1000); - _blockSize = gmc_config.GetInteger(_algo_name, "block_size", 3); - _ransac_max_iters = - gmc_config.GetInteger(_algo_name, "ransac_max_iters", 500); + _maxCorners = config.max_corners; + _blockSize = config.block_size; + _ransac_max_iters = config.ransac_max_iters; - _qualityLevel = gmc_config.GetReal(_algo_name, "quality_level", 0.01); - _k = gmc_config.GetReal(_algo_name, "k", 0.04); - _minDistance = gmc_config.GetReal(_algo_name, "min_distance", 1.0); + _qualityLevel = config.quality_level; + _k = config.k; + _minDistance = config.min_distance; - _downscale = gmc_config.GetFloat(_algo_name, "downscale", 2.0F); - _inlier_ratio = gmc_config.GetFloat(_algo_name, "inlier_ratio", 0.5); - _ransac_conf = gmc_config.GetFloat(_algo_name, "ransac_conf", 0.99); + _downscale = config.downscale; + _inlier_ratio = config.inlier_ratio; + _ransac_conf = config.ransac_conf; } @@ -498,9 +480,10 @@ SparseOptFlow_GMC::apply(const cv::Mat &frame_raw, // OpenCV VideoStab -OpenCV_VideoStab_GMC::OpenCV_VideoStab_GMC(const std::string &config_path) +OpenCV_VideoStab_GMC::OpenCV_VideoStab_GMC( + const OpenCV_VideoStab_GMC_Params &config) { - _load_params_from_config(config_path); + _load_params_from_config(config); _motion_estimator = cv::makePtr( cv::videostab::MM_SIMILARITY); @@ -514,19 +497,11 @@ OpenCV_VideoStab_GMC::OpenCV_VideoStab_GMC(const std::string &config_path) void OpenCV_VideoStab_GMC::_load_params_from_config( - const std::string &config_path) + const OpenCV_VideoStab_GMC_Params &config) { - INIReader gmc_config(config_path); - if (gmc_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _downscale = gmc_config.GetFloat(_algo_name, "downscale", 2.0F); - _num_features = gmc_config.GetInteger(_algo_name, "num_features", 4000); - _detections_masking = - gmc_config.GetBoolean(_algo_name, "detections_masking", true); + _downscale = config.downscale; + _num_features = config.num_features; + _detections_masking = config.detection_masking; } @@ -596,23 +571,16 @@ OpenCV_VideoStab_GMC::apply(const cv::Mat &frame_raw, // Optical Flow Modified -OptFlowModified_GMC::OptFlowModified_GMC(const std::string &config_path) +OptFlowModified_GMC::OptFlowModified_GMC(const OptFlowModified_Params &config) { - _load_params_from_config(config_path); + _load_params_from_config(config); } void OptFlowModified_GMC::_load_params_from_config( - const std::string &config_path) + const OptFlowModified_Params &config) { - INIReader gmc_config(config_path); - if (gmc_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _downscale = gmc_config.GetFloat(_algo_name, "downscale", 2.0F); + _downscale = config.downscale; } diff --git a/botsort/src/GmcParams.cpp b/botsort/src/GmcParams.cpp new file mode 100644 index 0000000..0ff215a --- /dev/null +++ b/botsort/src/GmcParams.cpp @@ -0,0 +1,116 @@ +#include "GmcParams.h" + +#include + +#include "DataType.h" +#include "INIReader.h" + +constexpr std::size_t NUM_METHODS{5}; + +using MethodLoaderFunction = GMC_Params::MethodParams (*)(INIReader &); + +namespace +{ +GMC_Params::MethodParams load_orb_config(INIReader &gmc_config) +{ + ORB_Params params{}; + static constexpr auto SECTION = "orb"; + + gmc_config.LoadFloat(SECTION, "downscale", params.downscale); + gmc_config.LoadFloat(SECTION, "inlier_ratio", params.inlier_ratio); + gmc_config.LoadFloat(SECTION, "ransac_conf", params.ransac_conf); + gmc_config.LoadInteger(SECTION, "ransac_max_iters", + params.ransac_max_iters); + + return {params}; +} + +GMC_Params::MethodParams load_ecc_config(INIReader &gmc_config) +{ + ECC_Params params{}; + static constexpr auto SECTION = "ecc"; + + gmc_config.LoadFloat(SECTION, "downscale", params.downscale); + gmc_config.LoadInteger(SECTION, "max_iterations", params.max_iterations); + gmc_config.LoadFloat(SECTION, "termination_eps", params.termination_eps); + + return {params}; +} + +GMC_Params::MethodParams load_sparce_config(INIReader &gmc_config) +{ + SparseOptFlow_Params params{}; + static constexpr auto SECTION = "sparseOptFlow"; + + gmc_config.LoadBoolean(SECTION, "use_harris_detector", + params.use_harris_detector); + + gmc_config.LoadInteger(SECTION, "max_corners", params.max_corners); + gmc_config.LoadInteger(SECTION, "block_size", params.block_size); + gmc_config.LoadInteger(SECTION, "ransac_max_iters", + params.ransac_max_iters); + + gmc_config.LoadReal(SECTION, "quality_level", params.quality_level); + gmc_config.LoadReal(SECTION, "k", params.k); + gmc_config.LoadReal(SECTION, "min_distance", params.min_distance); + + gmc_config.LoadFloat(SECTION, "downscale", params.downscale); + gmc_config.LoadFloat(SECTION, "inlier_ratio", params.inlier_ratio); + gmc_config.LoadFloat(SECTION, "ransac_conf", params.ransac_conf); + + return {params}; +} + +GMC_Params::MethodParams load_opt_config(INIReader &gmc_config) +{ + OptFlowModified_Params params{}; + static constexpr auto SECTION = "OptFlowModified"; + + gmc_config.LoadFloat(SECTION, "downscale", params.downscale); + + return params; +} + +GMC_Params::MethodParams load_videostab_config(INIReader &gmc_config) +{ + OpenCV_VideoStab_GMC_Params params{}; + static constexpr auto SECTION = "OptFlowModified"; + + gmc_config.LoadFloat(SECTION, "downscale", params.downscale); + gmc_config.LoadFloat(SECTION, "num_features", params.num_features); + gmc_config.LoadBoolean(SECTION, "detections_masking", + params.detection_masking); + + return params; +} +}// namespace + +GMC_Params GMC_Params::load_config(GMC_Method method, + const std::string &config_path) +{ + GMC_Params config; + config.method_ = method; + + static std::array method_loaders = { + load_orb_config, load_ecc_config, load_sparce_config, + load_opt_config, load_videostab_config}; + + auto method_index = static_cast(method); + if (method_index >= NUM_METHODS) + { + throw std::runtime_error("Unknown global motion compensation method: " + + std::to_string(method)); + } + + INIReader gmc_config(config_path); + if (gmc_config.ParseError() < 0) + { + std::cout << "Can't load " << config_path << std::endl; + exit(1); + } + + config.method_params_ = + method_loaders[static_cast(method)](gmc_config); + + return config; +} \ No newline at end of file diff --git a/botsort/src/ReID.cpp b/botsort/src/ReID.cpp index 6a7d803..6b81919 100644 --- a/botsort/src/ReID.cpp +++ b/botsort/src/ReID.cpp @@ -2,11 +2,11 @@ #include "INIReader.h" -ReIDModel::ReIDModel(const std::string &config_path, +ReIDModel::ReIDModel(const ReIDParams ¶ms, const std::string &onnx_model_path) { std::cout << "Initializing ReID model" << std::endl; - _load_params_from_config(config_path); + _load_params_from_config(params); _onnx_model_path = onnx_model_path; _trt_inference_engine = @@ -44,35 +44,17 @@ void ReIDModel::pre_process(cv::Mat &image) cv::cvtColor(image, image, cv::COLOR_BGR2RGB); } - -void ReIDModel::_load_params_from_config(const std::string &config_path) +void ReIDModel::_load_params_from_config(const ReIDParams ¶ms) { - const std::string section_name = "ReID"; - INIReader reid_config(config_path); - if (reid_config.ParseError() < 0) - { - std::cout << "Can't load " << config_path << std::endl; - exit(1); - } - - _distance_metric = - reid_config.Get(section_name, "distance_metric", "euclidean"); - _trt_logging_level = - reid_config.GetInteger(section_name, "trt_log_level", 1); - - _model_optimization_params.batch_size = - reid_config.GetInteger(section_name, "batch_size", 1); - _model_optimization_params.fp16 = - reid_config.GetBoolean(section_name, "enable_FP16", true); - _model_optimization_params.tf32 = - reid_config.GetBoolean(section_name, "enable_TF32", true); - - _model_optimization_params.input_layer_name = - reid_config.Get(section_name, "input_layer_name", ""); + _distance_metric = params.distance_metric; + _trt_logging_level = params.trt_logging_level; + _model_optimization_params.batch_size = static_cast(params.batch_size); + _model_optimization_params.fp16 = params.enable_fp16; + _model_optimization_params.tf32 = params.enable_tf32; + _model_optimization_params.input_layer_name = params.input_layer_name; std::cout << "Trying to get input dims" << std::endl; - std::vector input_dims = - reid_config.GetList(section_name, "input_layer_dimensions"); + const auto &input_dims = params.input_layer_dimensions; _input_size = cv::Size(input_dims[3], input_dims[2]); std::cout << "Read input dims" << std::endl; @@ -81,10 +63,7 @@ void ReIDModel::_load_params_from_config(const std::string &config_path) _model_optimization_params.input_dims = nvinfer1::Dims4{ input_dims[0], input_dims[1], input_dims[2], input_dims[3]}; - _model_optimization_params.swapRB = - reid_config.GetBoolean(section_name, "swapRB", false); + _model_optimization_params.swapRB = params.swap_rb; - _model_optimization_params.output_layer_names = - reid_config.GetList(section_name, - "output_layer_names"); + _model_optimization_params.output_layer_names = params.output_layer_names; } \ No newline at end of file diff --git a/botsort/src/ReIDParams.cpp b/botsort/src/ReIDParams.cpp new file mode 100644 index 0000000..154f9c2 --- /dev/null +++ b/botsort/src/ReIDParams.cpp @@ -0,0 +1,41 @@ +#include "ReIDParams.h" + +#include + +#include "DataType.h" +#include "INIReader.h" + + +ReIDParams ReIDParams::load_config(const std::string &config_path) +{ + ReIDParams config{}; + + const std::string reid_name = "ReID"; + + INIReader reid_config(config_path); + if (reid_config.ParseError() < 0) + { + std::cout << "Can't load " << config_path << std::endl; + exit(1); + } + + reid_config.LoadInteger(reid_name, "gpu_id", config.gpu_id); + reid_config.LoadString(reid_name, "distance_metric", + config.distance_metric); + reid_config.LoadInteger(reid_name, "trt_logging_level", + config.trt_logging_level); + reid_config.LoadInteger(reid_name, "batch_size", config.batch_size); + reid_config.LoadString(reid_name, "input_layer_name", + config.input_layer_name); + + reid_config.LoadBoolean(reid_name, "enable_fp16", config.enable_fp16); + reid_config.LoadBoolean(reid_name, "enable_tf32", config.enable_tf32); + reid_config.LoadBoolean(reid_name, "swapRB", config.swap_rb); + + config.input_layer_dimensions = + reid_config.GetList(reid_name, "input_layer_dimensions"); + config.output_layer_names = + reid_config.GetList(reid_name, "output_layer_names"); + + return config; +} \ No newline at end of file diff --git a/botsort/src/TRT_InferenceEngine/TensorRT_InferenceEngine.cpp b/botsort/src/TRT_InferenceEngine/TensorRT_InferenceEngine.cpp index 6c12295..18f0b9b 100644 --- a/botsort/src/TRT_InferenceEngine/TensorRT_InferenceEngine.cpp +++ b/botsort/src/TRT_InferenceEngine/TensorRT_InferenceEngine.cpp @@ -2,6 +2,23 @@ #include +namespace +{ +std::string get_devicename_from_deviceid(int device_id) +{ + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, device_id); + auto device_name = std::string(prop.name); + + // Remove spaces from device name + device_name.erase( + std::remove_if(device_name.begin(), device_name.end(), ::isspace), + device_name.end()); + + return device_name; +} +}// namespace + inference_backend::TensorRTInferenceEngine::TensorRTInferenceEngine( TRTOptimizerParams &optimization_params, u_int8_t logging_level) { @@ -52,6 +69,21 @@ bool inference_backend::TensorRTInferenceEngine::load_model( return false; } + // Ensure gpu_id is valid + int numGPUs{0}; + cudaGetDeviceCount(&numGPUs); + if (_optimization_params.gpu_id >= numGPUs) + { + int numGPUs{0}; + cudaGetDeviceCount(&numGPUs); + _logger->log(nvinfer1::ILogger::Severity::kERROR, + ("Unable to set GPU device index to: " + + std::to_string(_optimization_params.gpu_id) + + ". Note, your device has " + std::to_string(numGPUs) + + " CUDA-capable GPU(s).") + .c_str()); + } + // Check if engine exists, if not build it std::string engine_path = get_engine_path(onnx_model_path); if (!file_exists(engine_path)) @@ -63,6 +95,16 @@ bool inference_backend::TensorRTInferenceEngine::load_model( _build_engine(onnx_model_path); } + // Set device index + if (auto ret = cudaSetDevice(_optimization_params.gpu_id); ret != 0) + { + _logger->log(nvinfer1::ILogger::Severity::kERROR, + ("Failed to set device index to: " + + std::to_string(_optimization_params.gpu_id)) + .c_str()); + return false; + } + // Deserialize engine _logger->log(nvinfer1::ILogger::Severity::kINFO, std::string("Deserializing engine from path: ") @@ -99,6 +141,9 @@ std::string inference_backend::TensorRTInferenceEngine::get_engine_path( gethostname(hostname, sizeof(hostname)); std::string suffix(hostname); + suffix.append("_GPU_" + + get_devicename_from_deviceid(_optimization_params.gpu_id)); + // TensorRT version suffix.append("_TRT" + std::to_string(NV_TENSORRT_VERSION)); @@ -119,7 +164,7 @@ std::string inference_backend::TensorRTInferenceEngine::get_engine_path( : suffix.append("_FP32"); } - // Engine path = parent_dir/model_name_hostname_TRT_version_CUDA_version_batch_size_int8_fp16_fp32.engine + // Engine path = parent_dir/model_name_hostname_GPU_device_name_TRT_version_CUDA_version_batch_size_int8_fp16_fp32.engine engine_path.append("_" + suffix + ".engine"); return engine_path; } diff --git a/botsort/src/TrackerParams.cpp b/botsort/src/TrackerParams.cpp new file mode 100644 index 0000000..a6ebe43 --- /dev/null +++ b/botsort/src/TrackerParams.cpp @@ -0,0 +1,44 @@ +#include "TrackerParams.h" + +#include + +#include "DataType.h" +#include "INIReader.h" + + +TrackerParams TrackerParams::load_config(const std::string &config_path) +{ + TrackerParams config{}; + + const std::string tracker_name = "BoTSORT"; + + INIReader tracker_config(config_path); + if (tracker_config.ParseError() < 0) + { + std::cout << "Can't load " << config_path << std::endl; + exit(1); + } + + tracker_config.LoadBoolean(tracker_name, "enable_reid", + config.reid_enabled); + tracker_config.LoadBoolean(tracker_name, "enable_gmc", config.gmc_enabled); + tracker_config.LoadFloat(tracker_name, "track_high_thresh", + config.track_high_thresh); + tracker_config.LoadFloat(tracker_name, "track_low_thresh", + config.track_low_thresh); + tracker_config.LoadFloat(tracker_name, "new_track_thresh", + config.new_track_thresh); + tracker_config.LoadInteger(tracker_name, "track_buffer", + config.track_buffer); + tracker_config.LoadFloat(tracker_name, "match_thresh", config.match_thresh); + tracker_config.LoadFloat(tracker_name, "proximity_thresh", + config.proximity_thresh); + tracker_config.LoadFloat(tracker_name, "appearance_thresh", + config.appearance_thresh); + tracker_config.LoadString(tracker_name, "gmc_method", + config.gmc_method_name); + tracker_config.LoadInteger(tracker_name, "frame_rate", config.frame_rate); + tracker_config.LoadFloat(tracker_name, "lambda", config.lambda); + + return config; +} \ No newline at end of file diff --git a/botsort/src/track.cpp b/botsort/src/track.cpp index 5b85d94..5e11cef 100644 --- a/botsort/src/track.cpp +++ b/botsort/src/track.cpp @@ -228,6 +228,11 @@ float Track::get_score() const return _score; } +uint8_t Track::get_class_id() const +{ + return _class_id; +} + void Track::_update_class_id(uint8_t class_id, float score) { if (!_class_hist.empty())