├── weights └── empty.txt ├── .gitignore ├── samples └── results.jpg ├── configs ├── traffic.names └── config.yaml ├── depends └── yaml-cpp │ ├── libs │ └── libyaml-cpp.a │ └── include │ └── yaml-cpp │ ├── anchor.h │ ├── emitterstyle.h │ ├── node │ ├── type.h │ ├── detail │ │ ├── iterator_fwd.h │ │ ├── bool_type.h │ │ ├── memory.h │ │ ├── iterator.h │ │ ├── node_ref.h │ │ ├── node_data.h │ │ ├── node.h │ │ ├── node_iterator.h │ │ └── impl.h │ ├── ptr.h │ ├── emit.h │ ├── iterator.h │ ├── parse.h │ ├── node.h │ ├── convert.h │ └── impl.h │ ├── emitterdef.h │ ├── yaml.h │ ├── null.h │ ├── mark.h │ ├── contrib │ ├── anchordict.h │ └── graphbuilder.h │ ├── dll.h │ ├── stlemitter.h │ ├── eventhandler.h │ ├── emitfromevents.h │ ├── ostream_wrapper.h │ ├── binary.h │ ├── parser.h │ ├── traits.h │ ├── emittermanip.h │ ├── emitter.h │ └── exceptions.h ├── includes ├── video.h ├── fast-reid.h ├── yolov5.h ├── model.h ├── KalmanTracker.h ├── tracker.h ├── common.h ├── Hungarian.h └── logging.h ├── src ├── main.cpp ├── video.cpp ├── common.cpp ├── model.cpp ├── KalmanTracker.cpp ├── fast-reid.cpp ├── yolov5.cpp ├── tracker.cpp └── Hungarian.cpp ├── CMakeLists.txt ├── Dockerfile └── README.md /weights/empty.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | cmake-build-debug/ 3 | 4 | *.onnx 5 | *.trt 6 | *.mp4 7 | *.avi 8 | *.mpg 9 | -------------------------------------------------------------------------------- /samples/results.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linghu8812/yolov5_fastreid_deepsort_tensorrt/HEAD/samples/results.jpg -------------------------------------------------------------------------------- /configs/traffic.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | bus 6 | truck 7 | SUV 8 | tricycle 9 | van 10 | muck truck 11 | oil tank 12 | -------------------------------------------------------------------------------- /depends/yaml-cpp/libs/libyaml-cpp.a: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linghu8812/yolov5_fastreid_deepsort_tensorrt/HEAD/depends/yaml-cpp/libs/libyaml-cpp.a -------------------------------------------------------------------------------- /includes/video.h: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/10/29. 3 | // 4 | 5 | #ifndef OBJECT_TRACKER_VIDEO_H 6 | #define OBJECT_TRACKER_VIDEO_H 7 | 8 | #include "yolov5.h" 9 | #include "tracker.h" 10 | #include "fast-reid.h" 11 | 12 | void InferenceVideo(const std::string &video_name, YOLOv5 &yolov5, ObjectTracker &tracker, fastreid &fastreid); 13 | 14 | #endif //OBJECT_TRACKER_VIDEO_H 15 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/anchor.h: -------------------------------------------------------------------------------- 1 | #ifndef ANCHOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define ANCHOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | 12 | namespace YAML { 13 | typedef std::size_t anchor_t; 14 | const anchor_t NullAnchor = 0; 15 | } 16 | 17 | #endif // ANCHOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 18 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/emitterstyle.h: -------------------------------------------------------------------------------- 1 | #ifndef EMITTERSTYLE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EMITTERSTYLE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | namespace YAML { 11 | struct EmitterStyle { 12 | enum value { Default, Block, Flow }; 13 | }; 14 | } 15 | 16 | #endif // EMITTERSTYLE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 17 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/type.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_TYPE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_TYPE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | namespace YAML { 11 | struct NodeType { 12 | enum value { Undefined, Null, Scalar, Sequence, Map }; 13 | }; 14 | } 15 | 16 | #endif // VALUE_TYPE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 17 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/emitterdef.h: -------------------------------------------------------------------------------- 1 | #ifndef EMITTERDEF_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EMITTERDEF_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | namespace YAML { 11 | struct EmitterNodeType { 12 | enum value { NoType, Property, Scalar, FlowSeq, BlockSeq, FlowMap, BlockMap }; 13 | }; 14 | } 15 | 16 | #endif // EMITTERDEF_H_62B23520_7C8E_11DE_8A39_0800200C9A66 17 | -------------------------------------------------------------------------------- /src/main.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/10/29. 3 | // 4 | 5 | #include 6 | #include "yaml-cpp/yaml.h" 7 | #include "video.h" 8 | 9 | int main(int argc, char **argv) { 10 | if (argc < 3) 11 | { 12 | std::cout << "Please design config file and video name!" << std::endl; 13 | return -1; 14 | } 15 | std::string config_file = argv[1]; 16 | std::string video_name = argv[2]; 17 | YAML::Node root = YAML::LoadFile(config_file); 18 | YAML::Node yolov5_config = root["yolov5"]; 19 | YAML::Node tracker_config = root["tracker"]; 20 | YAML::Node fastreid_config = root["fastreid"]; 21 | YOLOv5 yolov5(yolov5_config); 22 | yolov5.LoadEngine(); 23 | ObjectTracker tracker(tracker_config); 24 | fastreid fastreid(fastreid_config); 25 | fastreid.LoadEngine(); 26 | InferenceVideo(video_name, yolov5, tracker, fastreid); 27 | return 0; 28 | } -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/yaml.h: -------------------------------------------------------------------------------- 1 | #ifndef YAML_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define YAML_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/parser.h" 11 | #include "yaml-cpp/emitter.h" 12 | #include "yaml-cpp/emitterstyle.h" 13 | #include "yaml-cpp/stlemitter.h" 14 | #include "yaml-cpp/exceptions.h" 15 | 16 | #include "yaml-cpp/node/node.h" 17 | #include "yaml-cpp/node/impl.h" 18 | #include "yaml-cpp/node/convert.h" 19 | #include "yaml-cpp/node/iterator.h" 20 | #include "yaml-cpp/node/detail/impl.h" 21 | #include "yaml-cpp/node/parse.h" 22 | #include "yaml-cpp/node/emit.h" 23 | 24 | #endif // YAML_H_62B23520_7C8E_11DE_8A39_0800200C9A66 25 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/null.h: -------------------------------------------------------------------------------- 1 | #ifndef NULL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NULL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include 12 | 13 | namespace YAML { 14 | class Node; 15 | 16 | struct YAML_CPP_API _Null {}; 17 | inline bool operator==(const _Null&, const _Null&) { return true; } 18 | inline bool operator!=(const _Null&, const _Null&) { return false; } 19 | 20 | YAML_CPP_API bool IsNull(const Node& node); // old API only 21 | YAML_CPP_API bool IsNullString(const std::string& str); 22 | 23 | extern YAML_CPP_API _Null Null; 24 | } 25 | 26 | #endif // NULL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 27 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/mark.h: -------------------------------------------------------------------------------- 1 | #ifndef MARK_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define MARK_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | 12 | namespace YAML { 13 | struct YAML_CPP_API Mark { 14 | Mark() : pos(0), line(0), column(0) {} 15 | 16 | static const Mark null_mark() { return Mark(-1, -1, -1); } 17 | 18 | bool is_null() const { return pos == -1 && line == -1 && column == -1; } 19 | 20 | int pos; 21 | int line, column; 22 | 23 | private: 24 | Mark(int pos_, int line_, int column_) 25 | : pos(pos_), line(line_), column(column_) {} 26 | }; 27 | } 28 | 29 | #endif // MARK_H_62B23520_7C8E_11DE_8A39_0800200C9A66 30 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/iterator_fwd.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_DETAIL_ITERATOR_FWD_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_DETAIL_ITERATOR_FWD_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include 12 | #include 13 | #include 14 | 15 | namespace YAML { 16 | 17 | namespace detail { 18 | struct iterator_value; 19 | template 20 | class iterator_base; 21 | } 22 | 23 | typedef detail::iterator_base iterator; 24 | typedef detail::iterator_base const_iterator; 25 | } 26 | 27 | #endif // VALUE_DETAIL_ITERATOR_FWD_H_62B23520_7C8E_11DE_8A39_0800200C9A66 28 | -------------------------------------------------------------------------------- /configs/config.yaml: -------------------------------------------------------------------------------- 1 | yolov5: 2 | onnx_file: "../weights/yolov5s_traffic.onnx" 3 | engine_file: "../weights/yolov5s_traffic.trt" 4 | labels_file: "../configs/traffic.names" 5 | BATCH_SIZE: 1 6 | INPUT_CHANNEL: 3 7 | IMAGE_WIDTH: 640 8 | IMAGE_HEIGHT: 640 9 | obj_threshold: 0.25 10 | nms_threshold: 0.45 11 | agnostic: True 12 | strides: [8, 16, 32] 13 | num_anchors: [3, 3, 3] 14 | anchors: [[10,13], [16,30], [33,23], [30,61], [62,45], [59,119], [116,90], [156,198], [373,326]] 15 | 16 | tracker: 17 | max_age: 70 18 | iou_threshold: 0.3 19 | sim_threshold: 0.4 20 | agnostic: True 21 | labels_file: "../configs/traffic.names" 22 | 23 | fastreid: 24 | onnx_file: "../weights/fast-reid_mobilenetv2.onnx" 25 | engine_file: "../weights/fast-reid_mobilenetv2.trt" 26 | BATCH_SIZE: 32 27 | INPUT_CHANNEL: 3 28 | IMAGE_WIDTH: 128 29 | IMAGE_HEIGHT: 256 30 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/ptr.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_PTR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_PTR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include 12 | 13 | namespace YAML { 14 | namespace detail { 15 | class node; 16 | class node_ref; 17 | class node_data; 18 | class memory; 19 | class memory_holder; 20 | 21 | typedef std::shared_ptr shared_node; 22 | typedef std::shared_ptr shared_node_ref; 23 | typedef std::shared_ptr shared_node_data; 24 | typedef std::shared_ptr shared_memory_holder; 25 | typedef std::shared_ptr shared_memory; 26 | } 27 | } 28 | 29 | #endif // VALUE_PTR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 30 | -------------------------------------------------------------------------------- /includes/fast-reid.h: -------------------------------------------------------------------------------- 1 | #ifndef TRACKER_FAST_REID_H 2 | #define TRACKER_FAST_REID_H 3 | 4 | #include 5 | #include "NvInfer.h" 6 | #include "model.h" 7 | #include "yaml-cpp/yaml.h" 8 | 9 | class fastreid : public Model { 10 | public: 11 | explicit fastreid(const YAML::Node &yolov5_config); 12 | ~fastreid(); 13 | std::vector InferenceImages(std::vector &vec_img); 14 | std::vector> InferenceImages(const std::vector &vec_img, 15 | const std::vector> &detections); 16 | 17 | private: 18 | std::vector prepareImage(std::vector &image) override; 19 | static std::vector CropSubImages(const cv::Mat &org_img, const std::vector &detection); 20 | float *ModelInference(std::vector image_data) override; 21 | static void ReshapeandNormalize(float out[], std::vector &feature, const int &MAT_SIZE, const int &outSize); 22 | }; 23 | 24 | #endif //TRACKER_FAST_REID_H 25 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/bool_type.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_DETAIL_BOOL_TYPE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_DETAIL_BOOL_TYPE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | namespace YAML { 11 | namespace detail { 12 | struct unspecified_bool { 13 | struct NOT_ALLOWED; 14 | static void true_value(NOT_ALLOWED*) {} 15 | }; 16 | typedef void (*unspecified_bool_type)(unspecified_bool::NOT_ALLOWED*); 17 | } 18 | } 19 | 20 | #define YAML_CPP_OPERATOR_BOOL() \ 21 | operator YAML::detail::unspecified_bool_type() const { \ 22 | return this->operator!() ? 0 \ 23 | : &YAML::detail::unspecified_bool::true_value; \ 24 | } 25 | 26 | #endif // NODE_DETAIL_BOOL_TYPE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 27 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/emit.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_EMIT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_EMIT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | 13 | #include "yaml-cpp/dll.h" 14 | 15 | namespace YAML { 16 | class Emitter; 17 | class Node; 18 | 19 | /** 20 | * Emits the node to the given {@link Emitter}. If there is an error in writing, 21 | * {@link Emitter#good} will return false. 22 | */ 23 | YAML_CPP_API Emitter& operator<<(Emitter& out, const Node& node); 24 | 25 | /** Emits the node to the given output stream. */ 26 | YAML_CPP_API std::ostream& operator<<(std::ostream& out, const Node& node); 27 | 28 | /** Converts the node to a YAML string. */ 29 | YAML_CPP_API std::string Dump(const Node& node); 30 | } // namespace YAML 31 | 32 | #endif // NODE_EMIT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 33 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/iterator.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include "yaml-cpp/node/node.h" 12 | #include "yaml-cpp/node/detail/iterator_fwd.h" 13 | #include "yaml-cpp/node/detail/iterator.h" 14 | #include 15 | #include 16 | #include 17 | 18 | namespace YAML { 19 | namespace detail { 20 | struct iterator_value : public Node, std::pair { 21 | iterator_value() {} 22 | explicit iterator_value(const Node& rhs) 23 | : Node(rhs), 24 | std::pair(Node(Node::ZombieNode), Node(Node::ZombieNode)) {} 25 | explicit iterator_value(const Node& key, const Node& value) 26 | : Node(Node::ZombieNode), std::pair(key, value) {} 27 | }; 28 | } 29 | } 30 | 31 | #endif // VALUE_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 32 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/contrib/anchordict.h: -------------------------------------------------------------------------------- 1 | #ifndef ANCHORDICT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define ANCHORDICT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | 12 | #include "../anchor.h" 13 | 14 | namespace YAML { 15 | /** 16 | * An object that stores and retrieves values correlating to {@link anchor_t} 17 | * values. 18 | * 19 | *

Efficient implementation that can make assumptions about how 20 | * {@code anchor_t} values are assigned by the {@link Parser} class. 21 | */ 22 | template 23 | class AnchorDict { 24 | public: 25 | AnchorDict() : m_data{} {} 26 | void Register(anchor_t anchor, T value) { 27 | if (anchor > m_data.size()) { 28 | m_data.resize(anchor); 29 | } 30 | m_data[anchor - 1] = value; 31 | } 32 | 33 | T Get(anchor_t anchor) const { return m_data[anchor - 1]; } 34 | 35 | private: 36 | std::vector m_data; 37 | }; 38 | } // namespace YAML 39 | 40 | #endif // ANCHORDICT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 41 | -------------------------------------------------------------------------------- /includes/yolov5.h: -------------------------------------------------------------------------------- 1 | #ifndef TRACKER_YOLOV5_H 2 | #define TRACKER_YOLOV5_H 3 | 4 | #include 5 | #include "NvInfer.h" 6 | #include "model.h" 7 | #include "yaml-cpp/yaml.h" 8 | 9 | class YOLOv5 : public Model 10 | { 11 | public: 12 | explicit YOLOv5(const YAML::Node &yolov5_config); 13 | ~YOLOv5(); 14 | std::vector> InferenceImages(std::vector &vec_img); 15 | void DrawResults(const std::vector> &detections, std::vector &vec_img); 16 | 17 | private: 18 | std::vector prepareImage(std::vector &vec_img) override; 19 | float *ModelInference(std::vector image_data) override; 20 | std::vector> postProcess(const std::vector &vec_Mat, float *output); 21 | void NmsDetect(std::vector &detections); 22 | static float IOUCalculate(const DetectRes &det_a, const DetectRes &det_b); 23 | std::map class_labels; 24 | float obj_threshold; 25 | float nms_threshold; 26 | bool agnostic; 27 | std::vector strides; 28 | std::vector> anchors; 29 | std::vector> grids; 30 | std::vector class_colors; 31 | }; 32 | 33 | #endif //TRACKER_YOLOV5_H 34 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/memory.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_DETAIL_MEMORY_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_DETAIL_MEMORY_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | 12 | #include "yaml-cpp/dll.h" 13 | #include "yaml-cpp/node/ptr.h" 14 | 15 | namespace YAML { 16 | namespace detail { 17 | class node; 18 | } // namespace detail 19 | } // namespace YAML 20 | 21 | namespace YAML { 22 | namespace detail { 23 | class YAML_CPP_API memory { 24 | public: 25 | memory() : m_nodes{} {} 26 | node& create_node(); 27 | void merge(const memory& rhs); 28 | 29 | private: 30 | typedef std::set Nodes; 31 | Nodes m_nodes; 32 | }; 33 | 34 | class YAML_CPP_API memory_holder { 35 | public: 36 | memory_holder() : m_pMemory(new memory) {} 37 | 38 | node& create_node() { return m_pMemory->create_node(); } 39 | void merge(memory_holder& rhs); 40 | 41 | private: 42 | shared_memory m_pMemory; 43 | }; 44 | } // namespace detail 45 | } // namespace YAML 46 | 47 | #endif // VALUE_DETAIL_MEMORY_H_62B23520_7C8E_11DE_8A39_0800200C9A66 48 | -------------------------------------------------------------------------------- /src/video.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/10/29. 3 | // 4 | #include "video.h" 5 | #include 6 | 7 | void InferenceVideo(const std::string &video_name, YOLOv5 &yolov5, ObjectTracker &tracker, fastreid &fastreid) { 8 | std::cout << "Processing: " << video_name << std::endl; 9 | cv::VideoCapture video_cap(video_name); 10 | cv::Size sSize = cv::Size((int) video_cap.get(cv::CAP_PROP_FRAME_WIDTH), 11 | (int) video_cap.get(cv::CAP_PROP_FRAME_HEIGHT)); 12 | std::cout << "Frame width is: " << sSize.width << ", height is: " << sSize.height << std::endl; 13 | auto fFps = (float)video_cap.get(cv::CAP_PROP_FPS); 14 | cv::VideoWriter video_writer("result.avi", cv::VideoWriter::fourcc('M', 'J', 'P', 'G'), 15 | fFps, sSize); 16 | cv::Mat src_img; 17 | while (video_cap.read(src_img)) { 18 | cv::cvtColor(src_img, src_img, cv::COLOR_BGR2RGB); 19 | std::vector vec_org_img; 20 | vec_org_img.push_back(src_img); 21 | auto detect_boxes = yolov5.InferenceImages(vec_org_img); 22 | auto object_features = fastreid.InferenceImages(vec_org_img, detect_boxes); 23 | tracker.update(detect_boxes[0], object_features[0], vec_org_img[0].cols, vec_org_img[0].rows); 24 | tracker.DrawResults(vec_org_img[0]); 25 | cv::imwrite("../results.jpg", vec_org_img[0]); 26 | video_writer.write(vec_org_img[0]); 27 | } 28 | } -------------------------------------------------------------------------------- /includes/model.h: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/2/8. 3 | // 4 | 5 | #ifndef TENSORRT_INFERENCE_MODEL_H 6 | #define TENSORRT_INFERENCE_MODEL_H 7 | 8 | #include 9 | #include "NvInfer.h" 10 | #include "common.h" 11 | 12 | struct ClassRes{ 13 | int classes; 14 | float prob; 15 | }; 16 | 17 | struct DetectRes : ClassRes{ 18 | float x; 19 | float y; 20 | float w; 21 | float h; 22 | }; 23 | 24 | class Model 25 | { 26 | public: 27 | void LoadEngine(); 28 | virtual std::vector prepareImage(std::vector &image) = 0; 29 | // virtual float *InferenceImage(std::vector image_data) = 0; 30 | // virtual bool InferenceFolder(const std::string &folder_name) = 0; 31 | 32 | protected: 33 | bool readTrtFile(); 34 | void onnxToTRTModel(); 35 | virtual float *ModelInference(std::vector image_data) = 0; 36 | std::string onnx_file; 37 | std::string engine_file; 38 | std::string labels_file; 39 | int BATCH_SIZE; 40 | int INPUT_CHANNEL; 41 | int IMAGE_WIDTH; 42 | int IMAGE_HEIGHT; 43 | int CATEGORY; 44 | nvinfer1::ICudaEngine *engine = nullptr; 45 | nvinfer1::IExecutionContext *context = nullptr; 46 | void *buffers[2]; 47 | std::vector bufferSize; 48 | cudaStream_t stream; 49 | int outSize; 50 | std::vector img_mean; 51 | std::vector img_std; 52 | }; 53 | 54 | #endif //TENSORRT_INFERENCE_MODEL_H 55 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/dll.h: -------------------------------------------------------------------------------- 1 | #ifndef DLL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define DLL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | // The following ifdef block is the standard way of creating macros which make 11 | // exporting from a DLL simpler. All files within this DLL are compiled with the 12 | // yaml_cpp_EXPORTS symbol defined on the command line. This symbol should not 13 | // be defined on any project that uses this DLL. This way any other project 14 | // whose source files include this file see YAML_CPP_API functions as being 15 | // imported from a DLL, whereas this DLL sees symbols defined with this macro as 16 | // being exported. 17 | #undef YAML_CPP_API 18 | 19 | #ifdef YAML_CPP_DLL // Using or Building YAML-CPP DLL (definition defined 20 | // manually) 21 | #ifdef yaml_cpp_EXPORTS // Building YAML-CPP DLL (definition created by CMake 22 | // or defined manually) 23 | // #pragma message( "Defining YAML_CPP_API for DLL export" ) 24 | #define YAML_CPP_API __declspec(dllexport) 25 | #else // yaml_cpp_EXPORTS 26 | // #pragma message( "Defining YAML_CPP_API for DLL import" ) 27 | #define YAML_CPP_API __declspec(dllimport) 28 | #endif // yaml_cpp_EXPORTS 29 | #else // YAML_CPP_DLL 30 | #define YAML_CPP_API 31 | #endif // YAML_CPP_DLL 32 | 33 | #endif // DLL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 34 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/stlemitter.h: -------------------------------------------------------------------------------- 1 | #ifndef STLEMITTER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define STLEMITTER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | namespace YAML { 16 | template 17 | inline Emitter& EmitSeq(Emitter& emitter, const Seq& seq) { 18 | emitter << BeginSeq; 19 | for (typename Seq::const_iterator it = seq.begin(); it != seq.end(); ++it) 20 | emitter << *it; 21 | emitter << EndSeq; 22 | return emitter; 23 | } 24 | 25 | template 26 | inline Emitter& operator<<(Emitter& emitter, const std::vector& v) { 27 | return EmitSeq(emitter, v); 28 | } 29 | 30 | template 31 | inline Emitter& operator<<(Emitter& emitter, const std::list& v) { 32 | return EmitSeq(emitter, v); 33 | } 34 | 35 | template 36 | inline Emitter& operator<<(Emitter& emitter, const std::set& v) { 37 | return EmitSeq(emitter, v); 38 | } 39 | 40 | template 41 | inline Emitter& operator<<(Emitter& emitter, const std::map& m) { 42 | typedef typename std::map map; 43 | emitter << BeginMap; 44 | for (typename map::const_iterator it = m.begin(); it != m.end(); ++it) 45 | emitter << Key << it->first << Value << it->second; 46 | emitter << EndMap; 47 | return emitter; 48 | } 49 | } 50 | 51 | #endif // STLEMITTER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 52 | -------------------------------------------------------------------------------- /includes/KalmanTracker.h: -------------------------------------------------------------------------------- 1 | /////////////////////////////////////////////////////////////////////////////// 2 | // KalmanTracker.h: KalmanTracker Class Declaration 3 | 4 | #ifndef KALMAN_H 5 | #define KALMAN_H 2 6 | 7 | #include "opencv2/video/tracking.hpp" 8 | #include "opencv2/highgui/highgui.hpp" 9 | 10 | #define StateType cv::Rect_ 11 | 12 | 13 | // This class represents the internel state of individual tracked objects observed as bounding box. 14 | class KalmanTracker 15 | { 16 | public: 17 | KalmanTracker() 18 | { 19 | init_kf(StateType()); 20 | m_time_since_update = 0; 21 | m_hits = 0; 22 | m_hit_streak = 0; 23 | m_age = 0; 24 | m_id = kf_count; 25 | m_classes = -1; 26 | m_prob = 0.0; 27 | //kf_count++; 28 | } 29 | KalmanTracker(StateType initRect, int classes, float prob) 30 | { 31 | init_kf(initRect); 32 | m_time_since_update = 0; 33 | m_hits = 0; 34 | m_hit_streak = 0; 35 | m_age = 0; 36 | m_id = kf_count; 37 | kf_count++; 38 | m_classes = classes; 39 | m_prob = prob; 40 | } 41 | 42 | ~KalmanTracker() 43 | { 44 | m_history.clear(); 45 | } 46 | 47 | StateType predict(); 48 | void update(StateType stateMat, int classes, float prob, cv::Mat feature); 49 | 50 | StateType get_state(); 51 | StateType get_rect_xysr(float cx, float cy, float s, float r); 52 | 53 | static int kf_count; 54 | 55 | int m_time_since_update; 56 | int m_hits; 57 | int m_hit_streak; 58 | int m_age; 59 | int m_id; 60 | int m_classes; 61 | float m_prob; 62 | cv::Mat m_feature; 63 | 64 | private: 65 | void init_kf(StateType stateMat); 66 | 67 | cv::KalmanFilter kf; 68 | cv::Mat measurement; 69 | 70 | std::vector m_history; 71 | }; 72 | 73 | #endif -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/eventhandler.h: -------------------------------------------------------------------------------- 1 | #ifndef EVENTHANDLER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EVENTHANDLER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | 12 | #include "yaml-cpp/anchor.h" 13 | #include "yaml-cpp/emitterstyle.h" 14 | 15 | namespace YAML { 16 | struct Mark; 17 | 18 | class EventHandler { 19 | public: 20 | virtual ~EventHandler() {} 21 | 22 | virtual void OnDocumentStart(const Mark& mark) = 0; 23 | virtual void OnDocumentEnd() = 0; 24 | 25 | virtual void OnNull(const Mark& mark, anchor_t anchor) = 0; 26 | virtual void OnAlias(const Mark& mark, anchor_t anchor) = 0; 27 | virtual void OnScalar(const Mark& mark, const std::string& tag, 28 | anchor_t anchor, const std::string& value) = 0; 29 | 30 | virtual void OnSequenceStart(const Mark& mark, const std::string& tag, 31 | anchor_t anchor, EmitterStyle::value style) = 0; 32 | virtual void OnSequenceEnd() = 0; 33 | 34 | virtual void OnMapStart(const Mark& mark, const std::string& tag, 35 | anchor_t anchor, EmitterStyle::value style) = 0; 36 | virtual void OnMapEnd() = 0; 37 | 38 | virtual void OnAnchor(const Mark& /*mark*/, 39 | const std::string& /*anchor_name*/) { 40 | // empty default implementation for compatibility 41 | } 42 | }; 43 | } // namespace YAML 44 | 45 | #endif // EVENTHANDLER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 46 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.5) 2 | 3 | project(object_tracker) 4 | 5 | set(CMAKE_CXX_STANDARD 14) 6 | 7 | # CUDA 8 | find_package(CUDA REQUIRED) 9 | message(STATUS "Find CUDA include at ${CUDA_INCLUDE_DIRS}") 10 | message(STATUS "Find CUDA libraries: ${CUDA_LIBRARIES}") 11 | 12 | # TensorRT 13 | set(TENSORRT_ROOT /usr/src/tensorrt/) 14 | find_path(TENSORRT_INCLUDE_DIR NvInfer.h 15 | HINTS ${TENSORRT_ROOT} PATH_SUFFIXES include/) 16 | message(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}") 17 | find_library(TENSORRT_LIBRARY_INFER nvinfer 18 | HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR} 19 | PATH_SUFFIXES lib lib64 lib/x64) 20 | find_library(TENSORRT_LIBRARY_ONNXPARSER nvonnxparser 21 | HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR} 22 | PATH_SUFFIXES lib lib64 lib/x64) 23 | set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_ONNXPARSER}) 24 | message(STATUS "Find TensorRT libs: ${TENSORRT_LIBRARY}") 25 | 26 | # OpenCV 27 | find_package(OpenCV REQUIRED) 28 | message(STATUS "Find OpenCV include at ${OpenCV_INCLUDE_DIRS}") 29 | message(STATUS "Find OpenCV libraries: ${OpenCV_LIBRARIES}") 30 | 31 | set(YAML_INCLUDE ./depends/yaml-cpp/include) 32 | set(YAML_LIB_DIR ./depends/yaml-cpp/libs) 33 | 34 | include_directories(${CUDA_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${OpenCV_INCLUDE_DIRS} ${COMMON_INCLUDE} 35 | ${YAML_INCLUDE} includes) 36 | aux_source_directory(src SRC_LIST) 37 | message(STATUS "source file list: ${SRC_LIST}") 38 | link_directories(${YAML_LIB_DIR}) 39 | 40 | add_executable(object_tracker ${SRC_LIST}) 41 | target_link_libraries(object_tracker ${OpenCV_LIBRARIES} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY} yaml-cpp) 42 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/emitfromevents.h: -------------------------------------------------------------------------------- 1 | #ifndef EMITFROMEVENTS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EMITFROMEVENTS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | 12 | #include "yaml-cpp/anchor.h" 13 | #include "yaml-cpp/emitterstyle.h" 14 | #include "yaml-cpp/eventhandler.h" 15 | 16 | namespace YAML { 17 | struct Mark; 18 | } // namespace YAML 19 | 20 | namespace YAML { 21 | class Emitter; 22 | 23 | class EmitFromEvents : public EventHandler { 24 | public: 25 | EmitFromEvents(Emitter& emitter); 26 | 27 | virtual void OnDocumentStart(const Mark& mark); 28 | virtual void OnDocumentEnd(); 29 | 30 | virtual void OnNull(const Mark& mark, anchor_t anchor); 31 | virtual void OnAlias(const Mark& mark, anchor_t anchor); 32 | virtual void OnScalar(const Mark& mark, const std::string& tag, 33 | anchor_t anchor, const std::string& value); 34 | 35 | virtual void OnSequenceStart(const Mark& mark, const std::string& tag, 36 | anchor_t anchor, EmitterStyle::value style); 37 | virtual void OnSequenceEnd(); 38 | 39 | virtual void OnMapStart(const Mark& mark, const std::string& tag, 40 | anchor_t anchor, EmitterStyle::value style); 41 | virtual void OnMapEnd(); 42 | 43 | private: 44 | void BeginNode(); 45 | void EmitProps(const std::string& tag, anchor_t anchor); 46 | 47 | private: 48 | Emitter& m_emitter; 49 | 50 | struct State { 51 | enum value { WaitingForSequenceEntry, WaitingForKey, WaitingForValue }; 52 | }; 53 | std::stack m_stateStack; 54 | }; 55 | } 56 | 57 | #endif // EMITFROMEVENTS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 58 | -------------------------------------------------------------------------------- /includes/tracker.h: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/11/2. 3 | // 4 | 5 | #ifndef OBJECT_TRACKER_TRACKER_H 6 | #define OBJECT_TRACKER_TRACKER_H 7 | 8 | #include "yolov5.h" 9 | #include "KalmanTracker.h" 10 | 11 | struct TrackerRes : public DetectRes { 12 | TrackerRes(int cl, float pb, float xc, float yc, float wc, float hc, int id) : DetectRes() { 13 | classes = cl, prob = pb, x = xc, y = yc, w = wc, h = hc, object_id = id; 14 | } 15 | cv::Mat feature; 16 | int object_id; 17 | }; 18 | 19 | class ObjectTracker { 20 | public: 21 | explicit ObjectTracker(const YAML::Node &config); 22 | ~ObjectTracker(); 23 | void update(const std::vector &det_boxes, const std::vector &det_features, 24 | int width, int height); 25 | void DrawResults(cv::Mat &origin_img); 26 | 27 | public: 28 | std::vector tracker_boxes; 29 | private: 30 | static float IOUCalculate(const TrackerRes &det_a, const TrackerRes &det_b); 31 | void Alignment(std::vector> mat, std::set &unmatchedDetections, 32 | std::set &unmatchedTrajectories, std::vector &matchedPairs, 33 | int det_num, int trk_num, bool b_iou); 34 | void FeatureMatching(const std::vector &predict_boxes, std::set &unmatchedDetections, 35 | std::set &unmatchedTrajectories, std::vector &matchedPairs); 36 | void IOUMatching(const std::vector &predict_boxes, std::set &unmatchedDetections, 37 | std::set &unmatchedTrajectories, std::vector &matchedPairs); 38 | int max_age; 39 | float iou_threshold; 40 | float sim_threshold; 41 | bool agnostic; 42 | std::vector kalman_boxes; 43 | std::map class_labels; 44 | std::string labels_file; 45 | std::vector id_colors; 46 | }; 47 | 48 | #endif //OBJECT_TRACKER_TRACKER_H 49 | -------------------------------------------------------------------------------- /includes/common.h: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/2/8. 3 | // 4 | 5 | #ifndef TRACKER_COMMON_H 6 | #define TRACKER_COMMON_H 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include "NvOnnxParser.h" 13 | #include "logging.h" 14 | #include 15 | 16 | // These is necessary if we want to be able to write 1_GiB instead of 1.0_GiB. 17 | // Since the return type is signed, -1_GiB will work as expected. 18 | constexpr long long int operator"" _GiB(long long unsigned int val) 19 | { 20 | return val * (1 << 30); 21 | } 22 | constexpr long long int operator"" _MiB(long long unsigned int val) 23 | { 24 | return val * (1 << 20); 25 | } 26 | constexpr long long int operator"" _KiB(long long unsigned int val) 27 | { 28 | return val * (1 << 10); 29 | } 30 | 31 | static Logger gLogger{Logger::Severity::kINFO}; 32 | static LogStreamConsumer gLogVerbose{LOG_VERBOSE(gLogger)}; 33 | static LogStreamConsumer gLogInfo{LOG_INFO(gLogger)}; 34 | static LogStreamConsumer gLogWarning{LOG_WARN(gLogger)}; 35 | static LogStreamConsumer gLogError{LOG_ERROR(gLogger)}; 36 | static LogStreamConsumer gLogFatal{LOG_FATAL(gLogger)}; 37 | 38 | inline unsigned int getElementSize(nvinfer1::DataType t) 39 | { 40 | switch (t) 41 | { 42 | case nvinfer1::DataType::kINT32: return 4; 43 | case nvinfer1::DataType::kFLOAT: return 4; 44 | case nvinfer1::DataType::kHALF: return 2; 45 | case nvinfer1::DataType::kBOOL: 46 | case nvinfer1::DataType::kINT8: return 1; 47 | } 48 | throw std::runtime_error("Invalid DataType."); 49 | return 0; 50 | } 51 | 52 | inline int64_t volume(const nvinfer1::Dims& d) 53 | { 54 | return std::accumulate(d.d, d.d + d.nbDims, 1, std::multiplies()); 55 | } 56 | 57 | void setReportableSeverity(Logger::Severity severity); 58 | std::vectorreadFolder(const std::string &image_path); 59 | std::map readImageNetLabel(const std::string &fileName); 60 | std::map readClassLabel(const std::string &fileName); 61 | 62 | #endif //TRACKER_COMMON_H 63 | -------------------------------------------------------------------------------- /includes/Hungarian.h: -------------------------------------------------------------------------------- 1 | /////////////////////////////////////////////////////////////////////////////// 2 | // Hungarian.h: Header file for Class HungarianAlgorithm. 3 | // 4 | // This is a C++ wrapper with slight modification of a hungarian algorithm implementation by Markus Buehren. 5 | // The original implementation is a few mex-functions for use in MATLAB, found here: 6 | // http://www.mathworks.com/matlabcentral/fileexchange/6543-functions-for-the-rectangular-assignment-problem 7 | // 8 | // Both this code and the orignal code are published under the BSD license. 9 | // by Cong Ma, 2016 10 | // 11 | 12 | #include 13 | #include 14 | 15 | 16 | class HungarianAlgorithm 17 | { 18 | public: 19 | HungarianAlgorithm(); 20 | ~HungarianAlgorithm(); 21 | double Solve(std::vector>& DistMatrix, std::vector& Assignment); 22 | 23 | private: 24 | void assignmentoptimal(int *assignment, double *cost, double *distMatrix, int nOfRows, int nOfColumns); 25 | void buildassignmentvector(int *assignment, bool *starMatrix, int nOfRows, int nOfColumns); 26 | void computeassignmentcost(int *assignment, double *cost, double *distMatrix, int nOfRows); 27 | void step2a(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim); 28 | void step2b(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim); 29 | void step3(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim); 30 | void step4(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim, int row, int col); 31 | void step5(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim); 32 | }; 33 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/ostream_wrapper.h: -------------------------------------------------------------------------------- 1 | #ifndef OSTREAM_WRAPPER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define OSTREAM_WRAPPER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | 13 | #include "yaml-cpp/dll.h" 14 | 15 | namespace YAML { 16 | class YAML_CPP_API ostream_wrapper { 17 | public: 18 | ostream_wrapper(); 19 | explicit ostream_wrapper(std::ostream& stream); 20 | ostream_wrapper(const ostream_wrapper&) = delete; 21 | ostream_wrapper(ostream_wrapper&&) = delete; 22 | ostream_wrapper& operator=(const ostream_wrapper&) = delete; 23 | ostream_wrapper& operator=(ostream_wrapper&&) = delete; 24 | ~ostream_wrapper(); 25 | 26 | void write(const std::string& str); 27 | void write(const char* str, std::size_t size); 28 | 29 | void set_comment() { m_comment = true; } 30 | 31 | const char* str() const { 32 | if (m_pStream) { 33 | return nullptr; 34 | } else { 35 | m_buffer[m_pos] = '\0'; 36 | return &m_buffer[0]; 37 | } 38 | } 39 | 40 | std::size_t row() const { return m_row; } 41 | std::size_t col() const { return m_col; } 42 | std::size_t pos() const { return m_pos; } 43 | bool comment() const { return m_comment; } 44 | 45 | private: 46 | void update_pos(char ch); 47 | 48 | private: 49 | mutable std::vector m_buffer; 50 | std::ostream* const m_pStream; 51 | 52 | std::size_t m_pos; 53 | std::size_t m_row, m_col; 54 | bool m_comment; 55 | }; 56 | 57 | template 58 | inline ostream_wrapper& operator<<(ostream_wrapper& stream, 59 | const char (&str)[N]) { 60 | stream.write(str, N - 1); 61 | return stream; 62 | } 63 | 64 | inline ostream_wrapper& operator<<(ostream_wrapper& stream, 65 | const std::string& str) { 66 | stream.write(str); 67 | return stream; 68 | } 69 | 70 | inline ostream_wrapper& operator<<(ostream_wrapper& stream, char ch) { 71 | stream.write(&ch, 1); 72 | return stream; 73 | } 74 | } // namespace YAML 75 | 76 | #endif // OSTREAM_WRAPPER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 77 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvidia/cuda:10.2-cudnn8-devel-ubuntu18.04 2 | 3 | WORKDIR /home/install/ 4 | 5 | # apt-get 安装 6 | RUN apt-get update 7 | RUN apt-get install software-properties-common -y 8 | RUN add-apt-repository "deb http://security.ubuntu.com/ubuntu xenial-security main" 9 | RUN rm /etc/apt/sources.list.d/cuda.list /etc/apt/sources.list.d/nvidia-ml.list && \ 10 | apt-get update && apt-get -y upgrade && apt-get -y install ssh vim build-essential cmake git libgtk2.0-dev pkg-config \ 11 | libavcodec-dev libavformat-dev libswscale-dev python-dev python-numpy libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev \ 12 | libjasper-dev libdc1394-22-dev qtbase5-dev qtdeclarative5-dev python3-pip zip 13 | 14 | # opencv 安装 15 | RUN bash -xc "curl -O https://github.com/opencv/opencv/archive/4.3.0.zip && unzip opencv-4.3.0.zip && mv opencv-4.3.0 opencv && \ 16 | curl -O https://github.com/opencv/opencv_contrib/archive/4.3.0.zip && unzip opencv_contrib-4.3.0.zip && mv opencv_contrib-4.3.0 opencv_contrib && \ 17 | pushd opencv>&1 > /dev/null && mkdir build && pushd build>&1 > /dev/null && \ 18 | cmake -D WITH_QT=ON \ 19 | -D WITH_CUDA=ON \ 20 | -D BUILD_TIFF=ON \ 21 | -D BUILD_TESTS=OFF \ 22 | -D BUILD_PERF_TESTS=OFF \ 23 | -D OPENCV_GENERATE_PKGCONFIG=ON \ 24 | -D CMAKE_INSTALL_PREFIX=/usr/local \ 25 | -D OPENCV_EXTRA_MODULES_PATH=../../opencv_contrib/modules/ \ 26 | -D BUILD_opencv_xfeatures2d=OFF .. && \ 27 | make -j4 && make -j4 install && pkg-config --cflags opencv4 && echo '/usr/local/lib' > /etc/ld.so.conf.d/opencv.conf && \ 28 | popd 2>&1 > /dev/null && popd 2>&1 > /dev/null && rm -rf opencv-4.3.0.zip && rm -rf opencv_contrib-4.3.0.zip" 29 | 30 | # pip3 安装 31 | RUN pip3 install -U pip && pip3 install torch torchvision mxnet-cu102 onnx-simplifier \ 32 | && pip3 install --ignore-installed -U PyYAML 33 | 34 | # tensorrt 安装 35 | RUN bash -xc "curl -O https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.1/local_repo/nv-tensorrt-repo-ubuntu1804-cuda10.2-trt7.1.3.4-ga-20200617_1-1_amd64.deb && \ 36 | dpkg -i nv-tensorrt-repo-ubuntu1804-cuda10.2-trt7.1.3.4-ga-20200617_1-1_amd64.deb \ 37 | && apt-key add /var/nv-tensorrt-repo-cuda10.2-trt7.1.3.4-ga-20200617/7fa2af80.pub && apt-get update\ 38 | && apt-get -y install tensorrt && rm -rf nv-tensorrt-repo-ubuntu1804-cuda10.2-trt7.1.3.4-ga-20200617_1-1_amd64.deb" 39 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/binary.h: -------------------------------------------------------------------------------- 1 | #ifndef BASE64_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define BASE64_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | 13 | #include "yaml-cpp/dll.h" 14 | 15 | namespace YAML { 16 | YAML_CPP_API std::string EncodeBase64(const unsigned char *data, 17 | std::size_t size); 18 | YAML_CPP_API std::vector DecodeBase64(const std::string &input); 19 | 20 | class YAML_CPP_API Binary { 21 | public: 22 | Binary(const unsigned char *data_, std::size_t size_) 23 | : m_data{}, m_unownedData(data_), m_unownedSize(size_) {} 24 | Binary() : Binary(nullptr, 0) {} 25 | Binary(const Binary &) = default; 26 | Binary(Binary &&) = default; 27 | Binary &operator=(const Binary &) = default; 28 | Binary &operator=(Binary &&) = default; 29 | 30 | bool owned() const { return !m_unownedData; } 31 | std::size_t size() const { return owned() ? m_data.size() : m_unownedSize; } 32 | const unsigned char *data() const { 33 | return owned() ? &m_data[0] : m_unownedData; 34 | } 35 | 36 | void swap(std::vector &rhs) { 37 | if (m_unownedData) { 38 | m_data.swap(rhs); 39 | rhs.clear(); 40 | rhs.resize(m_unownedSize); 41 | std::copy(m_unownedData, m_unownedData + m_unownedSize, rhs.begin()); 42 | m_unownedData = nullptr; 43 | m_unownedSize = 0; 44 | } else { 45 | m_data.swap(rhs); 46 | } 47 | } 48 | 49 | bool operator==(const Binary &rhs) const { 50 | const std::size_t s = size(); 51 | if (s != rhs.size()) 52 | return false; 53 | const unsigned char *d1 = data(); 54 | const unsigned char *d2 = rhs.data(); 55 | for (std::size_t i = 0; i < s; i++) { 56 | if (*d1++ != *d2++) 57 | return false; 58 | } 59 | return true; 60 | } 61 | 62 | bool operator!=(const Binary &rhs) const { return !(*this == rhs); } 63 | 64 | private: 65 | std::vector m_data; 66 | const unsigned char *m_unownedData; 67 | std::size_t m_unownedSize; 68 | }; 69 | } // namespace YAML 70 | 71 | #endif // BASE64_H_62B23520_7C8E_11DE_8A39_0800200C9A66 72 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/parse.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_PARSE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_PARSE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | #include 13 | 14 | #include "yaml-cpp/dll.h" 15 | 16 | namespace YAML { 17 | class Node; 18 | 19 | /** 20 | * Loads the input string as a single YAML document. 21 | * 22 | * @throws {@link ParserException} if it is malformed. 23 | */ 24 | YAML_CPP_API Node Load(const std::string& input); 25 | 26 | /** 27 | * Loads the input string as a single YAML document. 28 | * 29 | * @throws {@link ParserException} if it is malformed. 30 | */ 31 | YAML_CPP_API Node Load(const char* input); 32 | 33 | /** 34 | * Loads the input stream as a single YAML document. 35 | * 36 | * @throws {@link ParserException} if it is malformed. 37 | */ 38 | YAML_CPP_API Node Load(std::istream& input); 39 | 40 | /** 41 | * Loads the input file as a single YAML document. 42 | * 43 | * @throws {@link ParserException} if it is malformed. 44 | * @throws {@link BadFile} if the file cannot be loaded. 45 | */ 46 | YAML_CPP_API Node LoadFile(const std::string& filename); 47 | 48 | /** 49 | * Loads the input string as a list of YAML documents. 50 | * 51 | * @throws {@link ParserException} if it is malformed. 52 | */ 53 | YAML_CPP_API std::vector LoadAll(const std::string& input); 54 | 55 | /** 56 | * Loads the input string as a list of YAML documents. 57 | * 58 | * @throws {@link ParserException} if it is malformed. 59 | */ 60 | YAML_CPP_API std::vector LoadAll(const char* input); 61 | 62 | /** 63 | * Loads the input stream as a list of YAML documents. 64 | * 65 | * @throws {@link ParserException} if it is malformed. 66 | */ 67 | YAML_CPP_API std::vector LoadAll(std::istream& input); 68 | 69 | /** 70 | * Loads the input file as a list of YAML documents. 71 | * 72 | * @throws {@link ParserException} if it is malformed. 73 | * @throws {@link BadFile} if the file cannot be loaded. 74 | */ 75 | YAML_CPP_API std::vector LoadAllFromFile(const std::string& filename); 76 | } // namespace YAML 77 | 78 | #endif // VALUE_PARSE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 79 | -------------------------------------------------------------------------------- /src/common.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/2/8. 3 | // 4 | 5 | #include "common.h" 6 | #include 7 | 8 | void setReportableSeverity(Logger::Severity severity) 9 | { 10 | gLogger.setReportableSeverity(severity); 11 | gLogVerbose.setReportableSeverity(severity); 12 | gLogInfo.setReportableSeverity(severity); 13 | gLogWarning.setReportableSeverity(severity); 14 | gLogError.setReportableSeverity(severity); 15 | gLogFatal.setReportableSeverity(severity); 16 | } 17 | 18 | std::vectorreadFolder(const std::string &image_path) 19 | { 20 | std::vector image_names; 21 | auto dir = opendir(image_path.c_str()); 22 | 23 | if ((dir) != nullptr) 24 | { 25 | struct dirent *entry; 26 | entry = readdir(dir); 27 | while (entry) 28 | { 29 | auto temp = image_path + "/" + entry->d_name; 30 | if (strcmp(entry->d_name, "") == 0 || strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) 31 | { 32 | entry = readdir(dir); 33 | continue; 34 | } 35 | image_names.push_back(temp); 36 | entry = readdir(dir); 37 | } 38 | } 39 | return image_names; 40 | } 41 | 42 | std::map readImageNetLabel(const std::string &fileName) 43 | { 44 | std::map imagenet_label; 45 | std::ifstream file(fileName); 46 | if (!file.is_open()) 47 | { 48 | std::cout << "read file error: " << fileName << std::endl; 49 | } 50 | std::string strLine; 51 | while (getline(file, strLine)) 52 | { 53 | int pos1 = strLine.find(":"); 54 | std::string first = strLine.substr(0, pos1); 55 | int pos2 = strLine.find_last_of("'"); 56 | std::string second = strLine.substr(pos1 + 3, pos2 - pos1 - 3); 57 | imagenet_label.insert({atoi(first.c_str()), second}); 58 | } 59 | file.close(); 60 | return imagenet_label; 61 | } 62 | 63 | std::map readClassLabel(const std::string &fileName) 64 | { 65 | std::map class_label; 66 | std::ifstream file(fileName); 67 | if (!file.is_open()) 68 | { 69 | std::cout << "read file error: " << fileName << std::endl; 70 | } 71 | std::string strLine; 72 | int index = 0; 73 | while (getline(file, strLine)) 74 | { 75 | class_label.insert({index, strLine}); 76 | index++; 77 | } 78 | file.close(); 79 | return class_label; 80 | } -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/parser.h: -------------------------------------------------------------------------------- 1 | #ifndef PARSER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define PARSER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | 13 | #include "yaml-cpp/dll.h" 14 | 15 | namespace YAML { 16 | class EventHandler; 17 | class Node; 18 | class Scanner; 19 | struct Directives; 20 | struct Token; 21 | 22 | /** 23 | * A parser turns a stream of bytes into one stream of "events" per YAML 24 | * document in the input stream. 25 | */ 26 | class YAML_CPP_API Parser { 27 | public: 28 | /** Constructs an empty parser (with no input. */ 29 | Parser(); 30 | 31 | Parser(const Parser&) = delete; 32 | Parser(Parser&&) = delete; 33 | Parser& operator=(const Parser&) = delete; 34 | Parser& operator=(Parser&&) = delete; 35 | 36 | /** 37 | * Constructs a parser from the given input stream. The input stream must 38 | * live as long as the parser. 39 | */ 40 | explicit Parser(std::istream& in); 41 | 42 | ~Parser(); 43 | 44 | /** Evaluates to true if the parser has some valid input to be read. */ 45 | explicit operator bool() const; 46 | 47 | /** 48 | * Resets the parser with the given input stream. Any existing state is 49 | * erased. 50 | */ 51 | void Load(std::istream& in); 52 | 53 | /** 54 | * Handles the next document by calling events on the {@code eventHandler}. 55 | * 56 | * @throw a ParserException on error. 57 | * @return false if there are no more documents 58 | */ 59 | bool HandleNextDocument(EventHandler& eventHandler); 60 | 61 | void PrintTokens(std::ostream& out); 62 | 63 | private: 64 | /** 65 | * Reads any directives that are next in the queue, setting the internal 66 | * {@code m_pDirectives} state. 67 | */ 68 | void ParseDirectives(); 69 | 70 | void HandleDirective(const Token& token); 71 | 72 | /** 73 | * Handles a "YAML" directive, which should be of the form 'major.minor' (like 74 | * a version number). 75 | */ 76 | void HandleYamlDirective(const Token& token); 77 | 78 | /** 79 | * Handles a "TAG" directive, which should be of the form 'handle prefix', 80 | * where 'handle' is converted to 'prefix' in the file. 81 | */ 82 | void HandleTagDirective(const Token& token); 83 | 84 | private: 85 | std::unique_ptr m_pScanner; 86 | std::unique_ptr m_pDirectives; 87 | }; 88 | } // namespace YAML 89 | 90 | #endif // PARSER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 91 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/iterator.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_DETAIL_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_DETAIL_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include "yaml-cpp/node/detail/node_iterator.h" 12 | #include "yaml-cpp/node/node.h" 13 | #include "yaml-cpp/node/ptr.h" 14 | #include 15 | #include 16 | 17 | 18 | namespace YAML { 19 | namespace detail { 20 | struct iterator_value; 21 | 22 | template 23 | class iterator_base { 24 | 25 | private: 26 | template 27 | friend class iterator_base; 28 | struct enabler {}; 29 | typedef node_iterator base_type; 30 | 31 | struct proxy { 32 | explicit proxy(const V& x) : m_ref(x) {} 33 | V* operator->() { return std::addressof(m_ref); } 34 | operator V*() { return std::addressof(m_ref); } 35 | 36 | V m_ref; 37 | }; 38 | 39 | public: 40 | using iterator_category = std::forward_iterator_tag; 41 | using value_type = V; 42 | using difference_type = std::ptrdiff_t; 43 | using pointer = V*; 44 | using reference = V; 45 | 46 | public: 47 | iterator_base() : m_iterator(), m_pMemory() {} 48 | explicit iterator_base(base_type rhs, shared_memory_holder pMemory) 49 | : m_iterator(rhs), m_pMemory(pMemory) {} 50 | 51 | template 52 | iterator_base(const iterator_base& rhs, 53 | typename std::enable_if::value, 54 | enabler>::type = enabler()) 55 | : m_iterator(rhs.m_iterator), m_pMemory(rhs.m_pMemory) {} 56 | 57 | iterator_base& operator++() { 58 | ++m_iterator; 59 | return *this; 60 | } 61 | 62 | iterator_base operator++(int) { 63 | iterator_base iterator_pre(*this); 64 | ++(*this); 65 | return iterator_pre; 66 | } 67 | 68 | template 69 | bool operator==(const iterator_base& rhs) const { 70 | return m_iterator == rhs.m_iterator; 71 | } 72 | 73 | template 74 | bool operator!=(const iterator_base& rhs) const { 75 | return m_iterator != rhs.m_iterator; 76 | } 77 | 78 | value_type operator*() const { 79 | const typename base_type::value_type& v = *m_iterator; 80 | if (v.pNode) 81 | return value_type(Node(*v, m_pMemory)); 82 | if (v.first && v.second) 83 | return value_type(Node(*v.first, m_pMemory), Node(*v.second, m_pMemory)); 84 | return value_type(); 85 | } 86 | 87 | proxy operator->() const { return proxy(**this); } 88 | 89 | private: 90 | base_type m_iterator; 91 | shared_memory_holder m_pMemory; 92 | }; 93 | } // namespace detail 94 | } // namespace YAML 95 | 96 | #endif // VALUE_DETAIL_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 97 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/traits.h: -------------------------------------------------------------------------------- 1 | #ifndef TRAITS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define TRAITS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | namespace YAML { 16 | template 17 | struct is_numeric { 18 | enum { value = false }; 19 | }; 20 | 21 | template <> 22 | struct is_numeric { 23 | enum { value = true }; 24 | }; 25 | template <> 26 | struct is_numeric { 27 | enum { value = true }; 28 | }; 29 | template <> 30 | struct is_numeric { 31 | enum { value = true }; 32 | }; 33 | template <> 34 | struct is_numeric { 35 | enum { value = true }; 36 | }; 37 | template <> 38 | struct is_numeric { 39 | enum { value = true }; 40 | }; 41 | template <> 42 | struct is_numeric { 43 | enum { value = true }; 44 | }; 45 | template <> 46 | struct is_numeric { 47 | enum { value = true }; 48 | }; 49 | template <> 50 | struct is_numeric { 51 | enum { value = true }; 52 | }; 53 | #if defined(_MSC_VER) && (_MSC_VER < 1310) 54 | template <> 55 | struct is_numeric<__int64> { 56 | enum { value = true }; 57 | }; 58 | template <> 59 | struct is_numeric { 60 | enum { value = true }; 61 | }; 62 | #else 63 | template <> 64 | struct is_numeric { 65 | enum { value = true }; 66 | }; 67 | template <> 68 | struct is_numeric { 69 | enum { value = true }; 70 | }; 71 | #endif 72 | template <> 73 | struct is_numeric { 74 | enum { value = true }; 75 | }; 76 | template <> 77 | struct is_numeric { 78 | enum { value = true }; 79 | }; 80 | template <> 81 | struct is_numeric { 82 | enum { value = true }; 83 | }; 84 | 85 | template 86 | struct enable_if_c { 87 | typedef T type; 88 | }; 89 | 90 | template 91 | struct enable_if_c {}; 92 | 93 | template 94 | struct enable_if : public enable_if_c {}; 95 | 96 | template 97 | struct disable_if_c { 98 | typedef T type; 99 | }; 100 | 101 | template 102 | struct disable_if_c {}; 103 | 104 | template 105 | struct disable_if : public disable_if_c {}; 106 | } 107 | 108 | template 109 | struct is_streamable { 110 | template 111 | static auto test(int) 112 | -> decltype(std::declval() << std::declval(), std::true_type()); 113 | 114 | template 115 | static auto test(...) -> std::false_type; 116 | 117 | static const bool value = decltype(test(0))::value; 118 | }; 119 | 120 | template 121 | struct streamable_to_string { 122 | static std::string impl(const Key& key) { 123 | std::stringstream ss; 124 | ss << key; 125 | return ss.str(); 126 | } 127 | }; 128 | 129 | template 130 | struct streamable_to_string { 131 | static std::string impl(const Key&) { 132 | return ""; 133 | } 134 | }; 135 | #endif // TRAITS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 136 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/node_ref.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_DETAIL_NODE_REF_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_DETAIL_NODE_REF_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include "yaml-cpp/node/type.h" 12 | #include "yaml-cpp/node/ptr.h" 13 | #include "yaml-cpp/node/detail/node_data.h" 14 | 15 | namespace YAML { 16 | namespace detail { 17 | class node_ref { 18 | public: 19 | node_ref() : m_pData(new node_data) {} 20 | node_ref(const node_ref&) = delete; 21 | node_ref& operator=(const node_ref&) = delete; 22 | 23 | bool is_defined() const { return m_pData->is_defined(); } 24 | const Mark& mark() const { return m_pData->mark(); } 25 | NodeType::value type() const { return m_pData->type(); } 26 | const std::string& scalar() const { return m_pData->scalar(); } 27 | const std::string& tag() const { return m_pData->tag(); } 28 | EmitterStyle::value style() const { return m_pData->style(); } 29 | 30 | void mark_defined() { m_pData->mark_defined(); } 31 | void set_data(const node_ref& rhs) { m_pData = rhs.m_pData; } 32 | 33 | void set_mark(const Mark& mark) { m_pData->set_mark(mark); } 34 | void set_type(NodeType::value type) { m_pData->set_type(type); } 35 | void set_tag(const std::string& tag) { m_pData->set_tag(tag); } 36 | void set_null() { m_pData->set_null(); } 37 | void set_scalar(const std::string& scalar) { m_pData->set_scalar(scalar); } 38 | void set_style(EmitterStyle::value style) { m_pData->set_style(style); } 39 | 40 | // size/iterator 41 | std::size_t size() const { return m_pData->size(); } 42 | 43 | const_node_iterator begin() const { 44 | return static_cast(*m_pData).begin(); 45 | } 46 | node_iterator begin() { return m_pData->begin(); } 47 | 48 | const_node_iterator end() const { 49 | return static_cast(*m_pData).end(); 50 | } 51 | node_iterator end() { return m_pData->end(); } 52 | 53 | // sequence 54 | void push_back(node& node, shared_memory_holder pMemory) { 55 | m_pData->push_back(node, pMemory); 56 | } 57 | void insert(node& key, node& value, shared_memory_holder pMemory) { 58 | m_pData->insert(key, value, pMemory); 59 | } 60 | 61 | // indexing 62 | template 63 | node* get(const Key& key, shared_memory_holder pMemory) const { 64 | return static_cast(*m_pData).get(key, pMemory); 65 | } 66 | template 67 | node& get(const Key& key, shared_memory_holder pMemory) { 68 | return m_pData->get(key, pMemory); 69 | } 70 | template 71 | bool remove(const Key& key, shared_memory_holder pMemory) { 72 | return m_pData->remove(key, pMemory); 73 | } 74 | 75 | node* get(node& key, shared_memory_holder pMemory) const { 76 | return static_cast(*m_pData).get(key, pMemory); 77 | } 78 | node& get(node& key, shared_memory_holder pMemory) { 79 | return m_pData->get(key, pMemory); 80 | } 81 | bool remove(node& key, shared_memory_holder pMemory) { 82 | return m_pData->remove(key, pMemory); 83 | } 84 | 85 | // map 86 | template 87 | void force_insert(const Key& key, const Value& value, 88 | shared_memory_holder pMemory) { 89 | m_pData->force_insert(key, value, pMemory); 90 | } 91 | 92 | private: 93 | shared_node_data m_pData; 94 | }; 95 | } 96 | } 97 | 98 | #endif // VALUE_DETAIL_NODE_REF_H_62B23520_7C8E_11DE_8A39_0800200C9A66 99 | -------------------------------------------------------------------------------- /src/model.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/2/8. 3 | // 4 | 5 | #include "model.h" 6 | #include "common.h" 7 | 8 | void Model::onnxToTRTModel() { 9 | // create the builder 10 | nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(gLogger.getTRTLogger()); 11 | assert(builder != nullptr); 12 | 13 | const auto explicitBatch = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 14 | auto network = builder->createNetworkV2(explicitBatch); 15 | auto config = builder->createBuilderConfig(); 16 | 17 | auto parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger()); 18 | if (!parser->parseFromFile(onnx_file.c_str(), static_cast(gLogger.getReportableSeverity()))) { 19 | gLogError << "Failure while parsing ONNX file" << std::endl; 20 | } 21 | // Build the engine 22 | builder->setMaxBatchSize(BATCH_SIZE); 23 | config->setMaxWorkspaceSize(1_GiB); 24 | config->setFlag(nvinfer1::BuilderFlag::kFP16); 25 | 26 | std::cout << "start building engine" << std::endl; 27 | engine = builder->buildEngineWithConfig(*network, *config); 28 | std::cout << "build engine done" << std::endl; 29 | assert(engine); 30 | // we can destroy the parser 31 | parser->destroy(); 32 | // save engine 33 | nvinfer1::IHostMemory *data = engine->serialize(); 34 | std::ofstream file; 35 | file.open(engine_file, std::ios::binary | std::ios::out); 36 | std::cout << "writing engine file..." << std::endl; 37 | file.write((const char *) data->data(), data->size()); 38 | std::cout << "save engine file done" << std::endl; 39 | file.close(); 40 | // then close everything down 41 | network->destroy(); 42 | builder->destroy(); 43 | } 44 | 45 | bool Model::readTrtFile() { 46 | std::string cached_engine; 47 | std::fstream file; 48 | std::cout << "loading filename from:" << engine_file << std::endl; 49 | nvinfer1::IRuntime *trtRuntime; 50 | file.open(engine_file, std::ios::binary | std::ios::in); 51 | 52 | if (!file.is_open()) { 53 | std::cout << "read file error: " << engine_file << std::endl; 54 | cached_engine = ""; 55 | } 56 | 57 | while (file.peek() != EOF) { 58 | std::stringstream buffer; 59 | buffer << file.rdbuf(); 60 | cached_engine.append(buffer.str()); 61 | } 62 | file.close(); 63 | 64 | trtRuntime = nvinfer1::createInferRuntime(gLogger.getTRTLogger()); 65 | engine = trtRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr); 66 | std::cout << "deserialize done" << std::endl; 67 | 68 | } 69 | 70 | void Model::LoadEngine() { 71 | // create and load engine 72 | std::fstream existEngine; 73 | existEngine.open(engine_file, std::ios::in); 74 | if (existEngine) { 75 | readTrtFile(); 76 | assert(engine != nullptr); 77 | } else { 78 | onnxToTRTModel(); 79 | assert(engine != nullptr); 80 | } 81 | 82 | context = engine->createExecutionContext(); 83 | assert(context != nullptr); 84 | 85 | //get buffers 86 | assert(engine->getNbBindings() == 2); 87 | int nbBindings = engine->getNbBindings(); 88 | bufferSize.resize(nbBindings); 89 | for (int i = 0; i < nbBindings; ++i) { 90 | nvinfer1::Dims dims = engine->getBindingDimensions(i); 91 | nvinfer1::DataType dtype = engine->getBindingDataType(i); 92 | int64_t totalSize = volume(dims) * getElementSize(dtype); 93 | bufferSize[i] = totalSize; 94 | std::cout << "binding" << i << ": " << totalSize << std::endl; 95 | cudaMalloc(&buffers[i], totalSize); 96 | } 97 | //get stream 98 | cudaStreamCreate(&stream); 99 | outSize = int(bufferSize[1] / sizeof(float) / BATCH_SIZE); 100 | } 101 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/emittermanip.h: -------------------------------------------------------------------------------- 1 | #ifndef EMITTERMANIP_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EMITTERMANIP_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | 12 | namespace YAML { 13 | enum EMITTER_MANIP { 14 | // general manipulators 15 | Auto, 16 | TagByKind, 17 | Newline, 18 | 19 | // output character set 20 | EmitNonAscii, 21 | EscapeNonAscii, 22 | 23 | // string manipulators 24 | // Auto, // duplicate 25 | SingleQuoted, 26 | DoubleQuoted, 27 | Literal, 28 | 29 | // bool manipulators 30 | YesNoBool, // yes, no 31 | TrueFalseBool, // true, false 32 | OnOffBool, // on, off 33 | UpperCase, // TRUE, N 34 | LowerCase, // f, yes 35 | CamelCase, // No, Off 36 | LongBool, // yes, On 37 | ShortBool, // y, t 38 | 39 | // int manipulators 40 | Dec, 41 | Hex, 42 | Oct, 43 | 44 | // document manipulators 45 | BeginDoc, 46 | EndDoc, 47 | 48 | // sequence manipulators 49 | BeginSeq, 50 | EndSeq, 51 | Flow, 52 | Block, 53 | 54 | // map manipulators 55 | BeginMap, 56 | EndMap, 57 | Key, 58 | Value, 59 | // Flow, // duplicate 60 | // Block, // duplicate 61 | // Auto, // duplicate 62 | LongKey 63 | }; 64 | 65 | struct _Indent { 66 | _Indent(int value_) : value(value_) {} 67 | int value; 68 | }; 69 | 70 | inline _Indent Indent(int value) { return _Indent(value); } 71 | 72 | struct _Alias { 73 | _Alias(const std::string& content_) : content(content_) {} 74 | std::string content; 75 | }; 76 | 77 | inline _Alias Alias(const std::string content) { return _Alias(content); } 78 | 79 | struct _Anchor { 80 | _Anchor(const std::string& content_) : content(content_) {} 81 | std::string content; 82 | }; 83 | 84 | inline _Anchor Anchor(const std::string content) { return _Anchor(content); } 85 | 86 | struct _Tag { 87 | struct Type { 88 | enum value { Verbatim, PrimaryHandle, NamedHandle }; 89 | }; 90 | 91 | explicit _Tag(const std::string& prefix_, const std::string& content_, 92 | Type::value type_) 93 | : prefix(prefix_), content(content_), type(type_) {} 94 | std::string prefix; 95 | std::string content; 96 | Type::value type; 97 | }; 98 | 99 | inline _Tag VerbatimTag(const std::string content) { 100 | return _Tag("", content, _Tag::Type::Verbatim); 101 | } 102 | 103 | inline _Tag LocalTag(const std::string content) { 104 | return _Tag("", content, _Tag::Type::PrimaryHandle); 105 | } 106 | 107 | inline _Tag LocalTag(const std::string& prefix, const std::string content) { 108 | return _Tag(prefix, content, _Tag::Type::NamedHandle); 109 | } 110 | 111 | inline _Tag SecondaryTag(const std::string content) { 112 | return _Tag("", content, _Tag::Type::NamedHandle); 113 | } 114 | 115 | struct _Comment { 116 | _Comment(const std::string& content_) : content(content_) {} 117 | std::string content; 118 | }; 119 | 120 | inline _Comment Comment(const std::string content) { return _Comment(content); } 121 | 122 | struct _Precision { 123 | _Precision(int floatPrecision_, int doublePrecision_) 124 | : floatPrecision(floatPrecision_), doublePrecision(doublePrecision_) {} 125 | 126 | int floatPrecision; 127 | int doublePrecision; 128 | }; 129 | 130 | inline _Precision FloatPrecision(int n) { return _Precision(n, -1); } 131 | 132 | inline _Precision DoublePrecision(int n) { return _Precision(-1, n); } 133 | 134 | inline _Precision Precision(int n) { return _Precision(n, n); } 135 | } 136 | 137 | #endif // EMITTERMANIP_H_62B23520_7C8E_11DE_8A39_0800200C9A66 138 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/node_data.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_DETAIL_NODE_DATA_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_DETAIL_NODE_DATA_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | #include "yaml-cpp/dll.h" 17 | #include "yaml-cpp/node/detail/node_iterator.h" 18 | #include "yaml-cpp/node/iterator.h" 19 | #include "yaml-cpp/node/ptr.h" 20 | #include "yaml-cpp/node/type.h" 21 | 22 | namespace YAML { 23 | namespace detail { 24 | class node; 25 | } // namespace detail 26 | } // namespace YAML 27 | 28 | namespace YAML { 29 | namespace detail { 30 | class YAML_CPP_API node_data { 31 | public: 32 | node_data(); 33 | node_data(const node_data&) = delete; 34 | node_data& operator=(const node_data&) = delete; 35 | 36 | void mark_defined(); 37 | void set_mark(const Mark& mark); 38 | void set_type(NodeType::value type); 39 | void set_tag(const std::string& tag); 40 | void set_null(); 41 | void set_scalar(const std::string& scalar); 42 | void set_style(EmitterStyle::value style); 43 | 44 | bool is_defined() const { return m_isDefined; } 45 | const Mark& mark() const { return m_mark; } 46 | NodeType::value type() const { 47 | return m_isDefined ? m_type : NodeType::Undefined; 48 | } 49 | const std::string& scalar() const { return m_scalar; } 50 | const std::string& tag() const { return m_tag; } 51 | EmitterStyle::value style() const { return m_style; } 52 | 53 | // size/iterator 54 | std::size_t size() const; 55 | 56 | const_node_iterator begin() const; 57 | node_iterator begin(); 58 | 59 | const_node_iterator end() const; 60 | node_iterator end(); 61 | 62 | // sequence 63 | void push_back(node& node, shared_memory_holder pMemory); 64 | void insert(node& key, node& value, shared_memory_holder pMemory); 65 | 66 | // indexing 67 | template 68 | node* get(const Key& key, shared_memory_holder pMemory) const; 69 | template 70 | node& get(const Key& key, shared_memory_holder pMemory); 71 | template 72 | bool remove(const Key& key, shared_memory_holder pMemory); 73 | 74 | node* get(node& key, shared_memory_holder pMemory) const; 75 | node& get(node& key, shared_memory_holder pMemory); 76 | bool remove(node& key, shared_memory_holder pMemory); 77 | 78 | // map 79 | template 80 | void force_insert(const Key& key, const Value& value, 81 | shared_memory_holder pMemory); 82 | 83 | public: 84 | static const std::string& empty_scalar(); 85 | 86 | private: 87 | void compute_seq_size() const; 88 | void compute_map_size() const; 89 | 90 | void reset_sequence(); 91 | void reset_map(); 92 | 93 | void insert_map_pair(node& key, node& value); 94 | void convert_to_map(shared_memory_holder pMemory); 95 | void convert_sequence_to_map(shared_memory_holder pMemory); 96 | 97 | template 98 | static node& convert_to_node(const T& rhs, shared_memory_holder pMemory); 99 | 100 | private: 101 | bool m_isDefined; 102 | Mark m_mark; 103 | NodeType::value m_type; 104 | std::string m_tag; 105 | EmitterStyle::value m_style; 106 | 107 | // scalar 108 | std::string m_scalar; 109 | 110 | // sequence 111 | typedef std::vector node_seq; 112 | node_seq m_sequence; 113 | 114 | mutable std::size_t m_seqSize; 115 | 116 | // map 117 | typedef std::vector> node_map; 118 | node_map m_map; 119 | 120 | typedef std::pair kv_pair; 121 | typedef std::list kv_pairs; 122 | mutable kv_pairs m_undefinedPairs; 123 | }; 124 | } 125 | } 126 | 127 | #endif // VALUE_DETAIL_NODE_DATA_H_62B23520_7C8E_11DE_8A39_0800200C9A66 128 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # **Object Tracking with TensorRT** 2 | 3 | ## **Introduction** 4 | 5 | This is an implementation for object tracking in cplusplus code. The object detector uses [yolov5s](https://github.com/ultralytics/yolov5/blob/master/models/yolov5s.yaml) model. The idea of [deepsort](https://arxiv.org/abs/1703.07402) is adopted in object tracking. The implementation of [sort](https://arxiv.org/abs/1602.00763) refers to the [sort-cpp](https://github.com/mcximing/sort-cpp). When extracting object features, it is extracted through the [fast-reid](https://github.com/JDAI-CV/fast-reid) trained model, and the person ReID uses [mobilenetv2](https://github.com/JDAI-CV/fast-reid/blob/master/fastreid/modeling/backbones/mobilenet.py). The purpose of using these lightweight models is to ensure the real-time efficiency of video processing. The model inference base on [TensorRT](https://developer.nvidia.com/zh-cn/tensorrt) engine. 6 | 7 | ## **How to run?** 8 | 9 | **0. Build environments** 10 | 11 | The TenosrRT environments build from Dockerfile, run with the following command. 12 | 13 | ```bash 14 | docker build -t tensorrt_tracker:0.1.0_rc . 15 | ``` 16 | 17 | Following yolov5 and fast-reid requirements file to install their depends packages. 18 | 19 | **1. Transform PyTorch weights to ONNX** 20 | 21 | - **Transform yolov5 weights** 22 | 23 | Use this [yolov5 repo](https://github.com/linghu8812/yolov5) to transform yolov5 *.pt weights to ONNX models. Run the following command 24 | ``` 25 | git clone https://github.com/linghu8812/yolov5.git 26 | python3 export.py ---weights weights/yolov5s.pt --batch-size 1 --imgsz 640 --include onnx --simplify 27 | ``` 28 | A pretrained yolov5 ***ONNX*** detection model can be downloaded form here, link: [https://pan.baidu.com/s/1RUz7Xk78lvKCeZNk_BBvoQ](https://pan.baidu.com/s/1RUz7Xk78lvKCeZNk_BBvoQ), code: *jung*. download this model and put it to the `weights` folder. 29 | 30 | - **Transform fastreid weights** 31 | 32 | Use official [fast-reid](https://github.com/JDAI-CV/fast-reid) to transform PyTorch weights to ONNX model. Run the following command 33 | ``` 34 | https://github.com/JDAI-CV/fast-reid.git 35 | python3 tools/deploy/onnx_export.py --config-file configs/Market1501/mgn_R50-ibn.yml --name mgn_R50-ibn --output outputs/onnx_model --batch-size 32 --opts MODEL.WEIGHTS market_mgn_R50-ibn.pth 36 | ``` 37 | A pretrained fast-reid ***ONNX*** detection model can be downloaded form here, link: [https://pan.baidu.com/s/19TuHxxuVYLBzie5_Vu0cCQ](https://pan.baidu.com/s/19TuHxxuVYLBzie5_Vu0cCQ), code: *1e35*. download this model and put it to the `weights` folder. 38 | 39 | **2. Get video for inference ready** 40 | 41 | Put video file for inference to `samples` folder. Here is a video demo for inference can be used: [https://pan.baidu.com/s/1Yyh1lwmzNl_gjvNz9EVI5w](https://pan.baidu.com/s/1Yyh1lwmzNl_gjvNz9EVI5w), code: *fpi0*. 42 | 43 | **3. Build project** 44 | 45 | Run the following command 46 | ``` 47 | git clone git@github.com:linghu8812/tensorrt_tracker.git 48 | mkdir build && cd build 49 | cmake .. 50 | make -j 51 | ``` 52 | 53 | **4. Run project** 54 | 55 | Run the following command 56 | ``` 57 | ./object_tracker ../configs/config.yaml ../samples/test.mpg 58 | ``` 59 | 60 | results demo: 61 | 62 | [![tensorrt_tracker](samples/results.jpg)](https://www.bilibili.com/video/BV1qg411K74p?t=0.0) 63 | 64 | ## **Reference** 65 | 66 | - **yolov5:** [https://github.com/ultralytics/yolov5](https://github.com/ultralytics/yolov5) 67 | - **FastReID: A Pytorch Toolbox for General Instance Re-identification:** [https://arxiv.org/abs/2006.02631](https://arxiv.org/abs/2006.02631) 68 | - **fast-reid:** [https://github.com/JDAI-CV/fast-reid](https://github.com/JDAI-CV/fast-reid) 69 | - **Simple Online and Realtime Tracking:** [https://arxiv.org/abs/1602.00763](https://arxiv.org/abs/1602.00763) 70 | - **sort-cpp:** [https://github.com/mcximing/sort-cpp](https://github.com/mcximing/sort-cpp) 71 | - **Simple Online and Realtime Tracking with a Deep Association Metric:** [https://arxiv.org/abs/1703.07402](https://arxiv.org/abs/1703.07402) 72 | - **tensorrt_inference:** [https://github.com/linghu8812/tensorrt_inference](https://github.com/linghu8812/tensorrt_inference) 73 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/node.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_NODE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_NODE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | 13 | #include "yaml-cpp/dll.h" 14 | #include "yaml-cpp/emitterstyle.h" 15 | #include "yaml-cpp/mark.h" 16 | #include "yaml-cpp/node/detail/bool_type.h" 17 | #include "yaml-cpp/node/detail/iterator_fwd.h" 18 | #include "yaml-cpp/node/ptr.h" 19 | #include "yaml-cpp/node/type.h" 20 | 21 | namespace YAML { 22 | namespace detail { 23 | class node; 24 | class node_data; 25 | struct iterator_value; 26 | } // namespace detail 27 | } // namespace YAML 28 | 29 | namespace YAML { 30 | class YAML_CPP_API Node { 31 | public: 32 | friend class NodeBuilder; 33 | friend class NodeEvents; 34 | friend struct detail::iterator_value; 35 | friend class detail::node; 36 | friend class detail::node_data; 37 | template 38 | friend class detail::iterator_base; 39 | template 40 | friend struct as_if; 41 | 42 | typedef YAML::iterator iterator; 43 | typedef YAML::const_iterator const_iterator; 44 | 45 | Node(); 46 | explicit Node(NodeType::value type); 47 | template 48 | explicit Node(const T& rhs); 49 | explicit Node(const detail::iterator_value& rhs); 50 | Node(const Node& rhs); 51 | ~Node(); 52 | 53 | YAML::Mark Mark() const; 54 | NodeType::value Type() const; 55 | bool IsDefined() const; 56 | bool IsNull() const { return Type() == NodeType::Null; } 57 | bool IsScalar() const { return Type() == NodeType::Scalar; } 58 | bool IsSequence() const { return Type() == NodeType::Sequence; } 59 | bool IsMap() const { return Type() == NodeType::Map; } 60 | 61 | // bool conversions 62 | YAML_CPP_OPERATOR_BOOL() 63 | bool operator!() const { return !IsDefined(); } 64 | 65 | // access 66 | template 67 | T as() const; 68 | template 69 | T as(const S& fallback) const; 70 | const std::string& Scalar() const; 71 | 72 | const std::string& Tag() const; 73 | void SetTag(const std::string& tag); 74 | 75 | // style 76 | // WARNING: This API might change in future releases. 77 | EmitterStyle::value Style() const; 78 | void SetStyle(EmitterStyle::value style); 79 | 80 | // assignment 81 | bool is(const Node& rhs) const; 82 | template 83 | Node& operator=(const T& rhs); 84 | Node& operator=(const Node& rhs); 85 | void reset(const Node& rhs = Node()); 86 | 87 | // size/iterator 88 | std::size_t size() const; 89 | 90 | const_iterator begin() const; 91 | iterator begin(); 92 | 93 | const_iterator end() const; 94 | iterator end(); 95 | 96 | // sequence 97 | template 98 | void push_back(const T& rhs); 99 | void push_back(const Node& rhs); 100 | 101 | // indexing 102 | template 103 | const Node operator[](const Key& key) const; 104 | template 105 | Node operator[](const Key& key); 106 | template 107 | bool remove(const Key& key); 108 | 109 | const Node operator[](const Node& key) const; 110 | Node operator[](const Node& key); 111 | bool remove(const Node& key); 112 | 113 | // map 114 | template 115 | void force_insert(const Key& key, const Value& value); 116 | 117 | private: 118 | enum Zombie { ZombieNode }; 119 | explicit Node(Zombie); 120 | explicit Node(Zombie, const std::string&); 121 | explicit Node(detail::node& node, detail::shared_memory_holder pMemory); 122 | 123 | void EnsureNodeExists() const; 124 | 125 | template 126 | void Assign(const T& rhs); 127 | void Assign(const char* rhs); 128 | void Assign(char* rhs); 129 | 130 | void AssignData(const Node& rhs); 131 | void AssignNode(const Node& rhs); 132 | 133 | private: 134 | bool m_isValid; 135 | // String representation of invalid key, if the node is invalid. 136 | std::string m_invalidKey; 137 | mutable detail::shared_memory_holder m_pMemory; 138 | mutable detail::node* m_pNode; 139 | }; 140 | 141 | YAML_CPP_API bool operator==(const Node& lhs, const Node& rhs); 142 | 143 | YAML_CPP_API Node Clone(const Node& node); 144 | 145 | template 146 | struct convert; 147 | } 148 | 149 | #endif // NODE_NODE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 150 | -------------------------------------------------------------------------------- /src/KalmanTracker.cpp: -------------------------------------------------------------------------------- 1 | /////////////////////////////////////////////////////////////////////////////// 2 | // KalmanTracker.cpp: KalmanTracker Class Implementation Declaration 3 | 4 | #include "KalmanTracker.h" 5 | 6 | 7 | int KalmanTracker::kf_count = 0; 8 | 9 | 10 | // initialize Kalman filter 11 | void KalmanTracker::init_kf(StateType stateMat) 12 | { 13 | int stateNum = 7; 14 | int measureNum = 4; 15 | kf = cv::KalmanFilter(stateNum, measureNum, 0); 16 | 17 | measurement = cv::Mat::zeros(measureNum, 1, CV_32F); 18 | 19 | kf.transitionMatrix = (cv::Mat_(stateNum, stateNum) << 20 | 1, 0, 0, 0, 1, 0, 0, 21 | 0, 1, 0, 0, 0, 1, 0, 22 | 0, 0, 1, 0, 0, 0, 1, 23 | 0, 0, 0, 1, 0, 0, 0, 24 | 0, 0, 0, 0, 1, 0, 0, 25 | 0, 0, 0, 0, 0, 1, 0, 26 | 0, 0, 0, 0, 0, 0, 1); 27 | 28 | setIdentity(kf.measurementMatrix); 29 | setIdentity(kf.processNoiseCov, cv::Scalar::all(1e-2)); 30 | setIdentity(kf.measurementNoiseCov, cv::Scalar::all(1e-1)); 31 | setIdentity(kf.errorCovPost, cv::Scalar::all(1)); 32 | 33 | // initialize state vector with bounding box in [cx,cy,s,r] style 34 | kf.statePost.at(0, 0) = stateMat.x + stateMat.width / 2; 35 | kf.statePost.at(1, 0) = stateMat.y + stateMat.height / 2; 36 | kf.statePost.at(2, 0) = stateMat.area(); 37 | kf.statePost.at(3, 0) = stateMat.width / stateMat.height; 38 | } 39 | 40 | 41 | // Predict the estimated bounding box. 42 | StateType KalmanTracker::predict() 43 | { 44 | // predict 45 | cv::Mat p = kf.predict(); 46 | m_age += 1; 47 | 48 | if (m_time_since_update > 0) 49 | m_hit_streak = 0; 50 | m_time_since_update += 1; 51 | 52 | StateType predictBox = get_rect_xysr(p.at(0, 0), p.at(1, 0), p.at(2, 0), p.at(3, 0)); 53 | 54 | m_history.push_back(predictBox); 55 | return m_history.back(); 56 | } 57 | 58 | 59 | // Update the state vector with observed bounding box. 60 | void KalmanTracker::update(StateType stateMat, int classes, float prob, cv::Mat feature) 61 | { 62 | m_time_since_update = 0; 63 | m_history.clear(); 64 | m_hits += 1; 65 | m_hit_streak += 1; 66 | m_classes = classes; 67 | m_prob = prob; 68 | m_feature = feature.clone(); 69 | 70 | // measurement 71 | measurement.at(0, 0) = stateMat.x + stateMat.width / 2; 72 | measurement.at(1, 0) = stateMat.y + stateMat.height / 2; 73 | measurement.at(2, 0) = stateMat.area(); 74 | measurement.at(3, 0) = stateMat.width / stateMat.height; 75 | 76 | // update 77 | kf.correct(measurement); 78 | } 79 | 80 | 81 | // Return the current state vector 82 | StateType KalmanTracker::get_state() 83 | { 84 | cv::Mat s = kf.statePost; 85 | return get_rect_xysr(s.at(0, 0), s.at(1, 0), s.at(2, 0), s.at(3, 0)); 86 | } 87 | 88 | 89 | // Convert bounding box from [cx,cy,s,r] to [x,y,w,h] style. 90 | StateType KalmanTracker::get_rect_xysr(float cx, float cy, float s, float r) 91 | { 92 | float w = sqrt(s * r); 93 | float h = s / w; 94 | float x = (cx - w / 2); 95 | float y = (cy - h / 2); 96 | 97 | if (x < 0 && cx > 0) 98 | x = 0; 99 | if (y < 0 && cy > 0) 100 | y = 0; 101 | 102 | return StateType(x, y, w, h); 103 | } 104 | 105 | 106 | 107 | /* 108 | // -------------------------------------------------------------------- 109 | // Kalman Filter Demonstrating, a 2-d ball demo 110 | // -------------------------------------------------------------------- 111 | 112 | const int winHeight = 600; 113 | const int winWidth = 800; 114 | 115 | Point mousePosition = Point(winWidth >> 1, winHeight >> 1); 116 | 117 | // mouse event callback 118 | void mouseEvent(int event, int x, int y, int flags, void *param) 119 | { 120 | if (event == CV_EVENT_MOUSEMOVE) { 121 | mousePosition = Point(x, y); 122 | } 123 | } 124 | 125 | void TestKF(); 126 | 127 | void main() 128 | { 129 | TestKF(); 130 | } 131 | 132 | 133 | void TestKF() 134 | { 135 | int stateNum = 4; 136 | int measureNum = 2; 137 | KalmanFilter kf = KalmanFilter(stateNum, measureNum, 0); 138 | 139 | // initialization 140 | Mat processNoise(stateNum, 1, CV_32F); 141 | Mat measurement = Mat::zeros(measureNum, 1, CV_32F); 142 | 143 | kf.transitionMatrix = *(Mat_(stateNum, stateNum) << 144 | 1, 0, 1, 0, 145 | 0, 1, 0, 1, 146 | 0, 0, 1, 0, 147 | 0, 0, 0, 1); 148 | 149 | setIdentity(kf.measurementMatrix); 150 | setIdentity(kf.processNoiseCov, Scalar::all(1e-2)); 151 | setIdentity(kf.measurementNoiseCov, Scalar::all(1e-1)); 152 | setIdentity(kf.errorCovPost, Scalar::all(1)); 153 | 154 | randn(kf.statePost, Scalar::all(0), Scalar::all(winHeight)); 155 | 156 | namedWindow("Kalman"); 157 | setMouseCallback("Kalman", mouseEvent); 158 | Mat img(winHeight, winWidth, CV_8UC3); 159 | 160 | while (1) 161 | { 162 | // predict 163 | Mat prediction = kf.predict(); 164 | Point predictPt = Point(prediction.at(0, 0), prediction.at(1, 0)); 165 | 166 | // generate measurement 167 | Point statePt = mousePosition; 168 | measurement.at(0, 0) = statePt.x; 169 | measurement.at(1, 0) = statePt.y; 170 | 171 | // update 172 | kf.correct(measurement); 173 | 174 | // visualization 175 | img.setTo(Scalar(255, 255, 255)); 176 | circle(img, predictPt, 8, CV_RGB(0, 255, 0), -1); // predicted point as green 177 | circle(img, statePt, 8, CV_RGB(255, 0, 0), -1); // current position as red 178 | 179 | imshow("Kalman", img); 180 | char code = (char)waitKey(100); 181 | if (code == 27 || code == 'q' || code == 'Q') 182 | break; 183 | } 184 | destroyWindow("Kalman"); 185 | } 186 | */ 187 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/node.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_DETAIL_NODE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_DETAIL_NODE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include "yaml-cpp/emitterstyle.h" 12 | #include "yaml-cpp/node/detail/node_ref.h" 13 | #include "yaml-cpp/node/ptr.h" 14 | #include "yaml-cpp/node/type.h" 15 | #include 16 | 17 | namespace YAML { 18 | namespace detail { 19 | class node { 20 | public: 21 | node() : m_pRef(new node_ref), m_dependencies{} {} 22 | node(const node&) = delete; 23 | node& operator=(const node&) = delete; 24 | 25 | bool is(const node& rhs) const { return m_pRef == rhs.m_pRef; } 26 | const node_ref* ref() const { return m_pRef.get(); } 27 | 28 | bool is_defined() const { return m_pRef->is_defined(); } 29 | const Mark& mark() const { return m_pRef->mark(); } 30 | NodeType::value type() const { return m_pRef->type(); } 31 | 32 | const std::string& scalar() const { return m_pRef->scalar(); } 33 | const std::string& tag() const { return m_pRef->tag(); } 34 | EmitterStyle::value style() const { return m_pRef->style(); } 35 | 36 | template 37 | bool equals(const T& rhs, shared_memory_holder pMemory); 38 | bool equals(const char* rhs, shared_memory_holder pMemory); 39 | 40 | void mark_defined() { 41 | if (is_defined()) 42 | return; 43 | 44 | m_pRef->mark_defined(); 45 | for (nodes::iterator it = m_dependencies.begin(); 46 | it != m_dependencies.end(); ++it) 47 | (*it)->mark_defined(); 48 | m_dependencies.clear(); 49 | } 50 | 51 | void add_dependency(node& rhs) { 52 | if (is_defined()) 53 | rhs.mark_defined(); 54 | else 55 | m_dependencies.insert(&rhs); 56 | } 57 | 58 | void set_ref(const node& rhs) { 59 | if (rhs.is_defined()) 60 | mark_defined(); 61 | m_pRef = rhs.m_pRef; 62 | } 63 | void set_data(const node& rhs) { 64 | if (rhs.is_defined()) 65 | mark_defined(); 66 | m_pRef->set_data(*rhs.m_pRef); 67 | } 68 | 69 | void set_mark(const Mark& mark) { m_pRef->set_mark(mark); } 70 | 71 | void set_type(NodeType::value type) { 72 | if (type != NodeType::Undefined) 73 | mark_defined(); 74 | m_pRef->set_type(type); 75 | } 76 | void set_null() { 77 | mark_defined(); 78 | m_pRef->set_null(); 79 | } 80 | void set_scalar(const std::string& scalar) { 81 | mark_defined(); 82 | m_pRef->set_scalar(scalar); 83 | } 84 | void set_tag(const std::string& tag) { 85 | mark_defined(); 86 | m_pRef->set_tag(tag); 87 | } 88 | 89 | // style 90 | void set_style(EmitterStyle::value style) { 91 | mark_defined(); 92 | m_pRef->set_style(style); 93 | } 94 | 95 | // size/iterator 96 | std::size_t size() const { return m_pRef->size(); } 97 | 98 | const_node_iterator begin() const { 99 | return static_cast(*m_pRef).begin(); 100 | } 101 | node_iterator begin() { return m_pRef->begin(); } 102 | 103 | const_node_iterator end() const { 104 | return static_cast(*m_pRef).end(); 105 | } 106 | node_iterator end() { return m_pRef->end(); } 107 | 108 | // sequence 109 | void push_back(node& input, shared_memory_holder pMemory) { 110 | m_pRef->push_back(input, pMemory); 111 | input.add_dependency(*this); 112 | } 113 | void insert(node& key, node& value, shared_memory_holder pMemory) { 114 | m_pRef->insert(key, value, pMemory); 115 | key.add_dependency(*this); 116 | value.add_dependency(*this); 117 | } 118 | 119 | // indexing 120 | template 121 | node* get(const Key& key, shared_memory_holder pMemory) const { 122 | // NOTE: this returns a non-const node so that the top-level Node can wrap 123 | // it, and returns a pointer so that it can be NULL (if there is no such 124 | // key). 125 | return static_cast(*m_pRef).get(key, pMemory); 126 | } 127 | template 128 | node& get(const Key& key, shared_memory_holder pMemory) { 129 | node& value = m_pRef->get(key, pMemory); 130 | value.add_dependency(*this); 131 | return value; 132 | } 133 | template 134 | bool remove(const Key& key, shared_memory_holder pMemory) { 135 | return m_pRef->remove(key, pMemory); 136 | } 137 | 138 | node* get(node& key, shared_memory_holder pMemory) const { 139 | // NOTE: this returns a non-const node so that the top-level Node can wrap 140 | // it, and returns a pointer so that it can be NULL (if there is no such 141 | // key). 142 | return static_cast(*m_pRef).get(key, pMemory); 143 | } 144 | node& get(node& key, shared_memory_holder pMemory) { 145 | node& value = m_pRef->get(key, pMemory); 146 | key.add_dependency(*this); 147 | value.add_dependency(*this); 148 | return value; 149 | } 150 | bool remove(node& key, shared_memory_holder pMemory) { 151 | return m_pRef->remove(key, pMemory); 152 | } 153 | 154 | // map 155 | template 156 | void force_insert(const Key& key, const Value& value, 157 | shared_memory_holder pMemory) { 158 | m_pRef->force_insert(key, value, pMemory); 159 | } 160 | 161 | private: 162 | shared_node_ref m_pRef; 163 | typedef std::set nodes; 164 | nodes m_dependencies; 165 | }; 166 | } // namespace detail 167 | } // namespace YAML 168 | 169 | #endif // NODE_DETAIL_NODE_H_62B23520_7C8E_11DE_8A39_0800200C9A66 170 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/node_iterator.h: -------------------------------------------------------------------------------- 1 | #ifndef VALUE_DETAIL_NODE_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define VALUE_DETAIL_NODE_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/dll.h" 11 | #include "yaml-cpp/node/ptr.h" 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | 19 | namespace YAML { 20 | namespace detail { 21 | struct iterator_type { 22 | enum value { NoneType, Sequence, Map }; 23 | }; 24 | 25 | template 26 | struct node_iterator_value : public std::pair { 27 | typedef std::pair kv; 28 | 29 | node_iterator_value() : kv(), pNode(nullptr) {} 30 | explicit node_iterator_value(V& rhs) : kv(), pNode(&rhs) {} 31 | explicit node_iterator_value(V& key, V& value) : kv(&key, &value), pNode(nullptr) {} 32 | 33 | V& operator*() const { return *pNode; } 34 | V& operator->() const { return *pNode; } 35 | 36 | V* pNode; 37 | }; 38 | 39 | typedef std::vector node_seq; 40 | typedef std::vector> node_map; 41 | 42 | template 43 | struct node_iterator_type { 44 | typedef node_seq::iterator seq; 45 | typedef node_map::iterator map; 46 | }; 47 | 48 | template 49 | struct node_iterator_type { 50 | typedef node_seq::const_iterator seq; 51 | typedef node_map::const_iterator map; 52 | }; 53 | 54 | template 55 | class node_iterator_base 56 | : public std::iterator, 57 | std::ptrdiff_t, node_iterator_value*, 58 | node_iterator_value> { 59 | private: 60 | struct enabler {}; 61 | 62 | struct proxy { 63 | explicit proxy(const node_iterator_value& x) : m_ref(x) {} 64 | node_iterator_value* operator->() { return std::addressof(m_ref); } 65 | operator node_iterator_value*() { return std::addressof(m_ref); } 66 | 67 | node_iterator_value m_ref; 68 | }; 69 | 70 | public: 71 | typedef typename node_iterator_type::seq SeqIter; 72 | typedef typename node_iterator_type::map MapIter; 73 | typedef node_iterator_value value_type; 74 | 75 | node_iterator_base() 76 | : m_type(iterator_type::NoneType), m_seqIt(), m_mapIt(), m_mapEnd() {} 77 | explicit node_iterator_base(SeqIter seqIt) 78 | : m_type(iterator_type::Sequence), 79 | m_seqIt(seqIt), 80 | m_mapIt(), 81 | m_mapEnd() {} 82 | explicit node_iterator_base(MapIter mapIt, MapIter mapEnd) 83 | : m_type(iterator_type::Map), 84 | m_seqIt(), 85 | m_mapIt(mapIt), 86 | m_mapEnd(mapEnd) { 87 | m_mapIt = increment_until_defined(m_mapIt); 88 | } 89 | 90 | template 91 | node_iterator_base(const node_iterator_base& rhs, 92 | typename std::enable_if::value, 93 | enabler>::type = enabler()) 94 | : m_type(rhs.m_type), 95 | m_seqIt(rhs.m_seqIt), 96 | m_mapIt(rhs.m_mapIt), 97 | m_mapEnd(rhs.m_mapEnd) {} 98 | 99 | template 100 | friend class node_iterator_base; 101 | 102 | template 103 | bool operator==(const node_iterator_base& rhs) const { 104 | if (m_type != rhs.m_type) 105 | return false; 106 | 107 | switch (m_type) { 108 | case iterator_type::NoneType: 109 | return true; 110 | case iterator_type::Sequence: 111 | return m_seqIt == rhs.m_seqIt; 112 | case iterator_type::Map: 113 | return m_mapIt == rhs.m_mapIt; 114 | } 115 | return true; 116 | } 117 | 118 | template 119 | bool operator!=(const node_iterator_base& rhs) const { 120 | return !(*this == rhs); 121 | } 122 | 123 | node_iterator_base& operator++() { 124 | switch (m_type) { 125 | case iterator_type::NoneType: 126 | break; 127 | case iterator_type::Sequence: 128 | ++m_seqIt; 129 | break; 130 | case iterator_type::Map: 131 | ++m_mapIt; 132 | m_mapIt = increment_until_defined(m_mapIt); 133 | break; 134 | } 135 | return *this; 136 | } 137 | 138 | node_iterator_base operator++(int) { 139 | node_iterator_base iterator_pre(*this); 140 | ++(*this); 141 | return iterator_pre; 142 | } 143 | 144 | value_type operator*() const { 145 | switch (m_type) { 146 | case iterator_type::NoneType: 147 | return value_type(); 148 | case iterator_type::Sequence: 149 | return value_type(**m_seqIt); 150 | case iterator_type::Map: 151 | return value_type(*m_mapIt->first, *m_mapIt->second); 152 | } 153 | return value_type(); 154 | } 155 | 156 | proxy operator->() const { return proxy(**this); } 157 | 158 | MapIter increment_until_defined(MapIter it) { 159 | while (it != m_mapEnd && !is_defined(it)) 160 | ++it; 161 | return it; 162 | } 163 | 164 | bool is_defined(MapIter it) const { 165 | return it->first->is_defined() && it->second->is_defined(); 166 | } 167 | 168 | private: 169 | typename iterator_type::value m_type; 170 | 171 | SeqIter m_seqIt; 172 | MapIter m_mapIt, m_mapEnd; 173 | }; 174 | 175 | typedef node_iterator_base node_iterator; 176 | typedef node_iterator_base const_node_iterator; 177 | } 178 | } 179 | 180 | #endif // VALUE_DETAIL_NODE_ITERATOR_H_62B23520_7C8E_11DE_8A39_0800200C9A66 181 | -------------------------------------------------------------------------------- /src/fast-reid.cpp: -------------------------------------------------------------------------------- 1 | #include "fast-reid.h" 2 | #include "yaml-cpp/yaml.h" 3 | 4 | fastreid::fastreid(const YAML::Node &config) { 5 | onnx_file = config["onnx_file"].as(); 6 | engine_file = config["engine_file"].as(); 7 | BATCH_SIZE = config["BATCH_SIZE"].as(); 8 | INPUT_CHANNEL = config["INPUT_CHANNEL"].as(); 9 | IMAGE_WIDTH = config["IMAGE_WIDTH"].as(); 10 | IMAGE_HEIGHT = config["IMAGE_HEIGHT"].as(); 11 | } 12 | 13 | fastreid::~fastreid() = default; 14 | 15 | std::vector fastreid::InferenceImages(std::vector &vec_img) { 16 | std::vector res_feature; 17 | int start_index = 0, end_index = 0; 18 | while (end_index < (int)vec_img.size()) { 19 | end_index = std::min(int(start_index + BATCH_SIZE), int(vec_img.size())); 20 | std::vector::const_iterator iter_1 = vec_img.begin() + start_index; 21 | std::vector::const_iterator iter_2 = vec_img.begin() + end_index; 22 | std::vector sub_vec_img(iter_1, iter_2); 23 | auto t_start_pre = std::chrono::high_resolution_clock::now(); 24 | std::vector image_data = prepareImage(sub_vec_img); 25 | auto t_end_pre = std::chrono::high_resolution_clock::now(); 26 | float total_pre = std::chrono::duration(t_end_pre - t_start_pre).count(); 27 | std::cout << "fast-reid prepare image take: " << total_pre << " ms." << std::endl; 28 | auto t_start = std::chrono::high_resolution_clock::now(); 29 | auto *output = ModelInference(image_data); 30 | auto t_end = std::chrono::high_resolution_clock::now(); 31 | float total_inf = std::chrono::duration(t_end - t_start).count(); 32 | std::cout << "fast-reid inference take: " << total_inf << " ms." << std::endl; 33 | auto r_start = std::chrono::high_resolution_clock::now(); 34 | // cv::Mat feature(BATCH_SIZE, outSize, CV_32FC1); 35 | std::vector feature; 36 | ReshapeandNormalize(output, feature, iter_2 - iter_1, outSize); 37 | res_feature.insert(res_feature.end(), feature.begin(), feature.end()); 38 | auto r_end = std::chrono::high_resolution_clock::now(); 39 | float total_res = std::chrono::duration(r_end - r_start).count(); 40 | std::cout << "fast-reid post process take: " << total_res << " ms." << std::endl; 41 | delete output; 42 | start_index += BATCH_SIZE; 43 | } 44 | return res_feature; 45 | } 46 | 47 | std::vector fastreid::CropSubImages(const cv::Mat &org_img, const std::vector &detection) { 48 | std::vector vec_mat; 49 | for (const auto &bbox : detection) { 50 | int x1 = std::max(int(bbox.x - bbox.w / 2), 0), y1 = std::max(int(bbox.y - bbox.h / 2), 0), 51 | x2 = std::min(int(bbox.x + bbox.w / 2), org_img.cols), y2 = std::min(int(bbox.y + bbox.h / 2), org_img.rows); 52 | cv::Rect rect{x1, y1, x2 - x1, y2 - y1}; 53 | cv::Mat sub_img = org_img(rect); 54 | vec_mat.push_back(sub_img); 55 | } 56 | return vec_mat; 57 | } 58 | 59 | std::vector> fastreid::InferenceImages(const std::vector &vec_img, 60 | const std::vector> &detections) { 61 | assert(vec_img.size() == detections.size()); 62 | std::vector> res_feature; 63 | for (int i = 0; i < (int)vec_img.size(); i++) { 64 | const cv::Mat &org_img = vec_img[i]; 65 | const std::vector detection = detections[i]; 66 | auto vec_mat = CropSubImages(org_img, detection); 67 | auto vec_features = InferenceImages(vec_mat); 68 | res_feature.push_back(vec_features); 69 | } 70 | return res_feature; 71 | } 72 | 73 | std::vector fastreid::prepareImage(std::vector &vec_img) { 74 | std::vector result(BATCH_SIZE * IMAGE_WIDTH * IMAGE_HEIGHT * INPUT_CHANNEL); 75 | float *data = result.data(); 76 | for (const cv::Mat &src_img : vec_img) 77 | { 78 | if (!src_img.data) 79 | continue; 80 | cv::Mat flt_img; 81 | cv::resize(src_img, flt_img, cv::Size(IMAGE_WIDTH, IMAGE_HEIGHT)); 82 | flt_img.convertTo(flt_img, CV_32FC3); 83 | 84 | //HWC TO CHW 85 | std::vector split_img(INPUT_CHANNEL); 86 | cv::split(flt_img, split_img); 87 | 88 | int channelLength = IMAGE_WIDTH * IMAGE_HEIGHT; 89 | for (int i = 0; i < INPUT_CHANNEL; ++i) 90 | { 91 | memcpy(data, split_img[i].data, channelLength * sizeof(float)); 92 | data += channelLength; 93 | } 94 | } 95 | return result; 96 | } 97 | 98 | float *fastreid::ModelInference(std::vector image_data) { 99 | auto *out = new float[outSize * BATCH_SIZE]; 100 | if (!image_data.data()) { 101 | std::cout << "prepare images ERROR!" << std::endl; 102 | return out; 103 | } 104 | // DMA the input to the GPU, execute the batch asynchronously, and DMA it back: 105 | cudaMemcpyAsync(buffers[0], image_data.data(), bufferSize[0], cudaMemcpyHostToDevice, stream); 106 | 107 | // do inference 108 | context->execute(BATCH_SIZE, buffers); 109 | cudaMemcpyAsync(out, buffers[1], bufferSize[1], cudaMemcpyDeviceToHost, stream); 110 | cudaStreamSynchronize(stream); 111 | return out; 112 | } 113 | 114 | void fastreid::ReshapeandNormalize(float *out, std::vector &feature, const int &MAT_SIZE, const int &outSize) { 115 | for (int i = 0; i < MAT_SIZE; i++) 116 | { 117 | cv::Mat onefeature(1, outSize, CV_32FC1, out + i * outSize); 118 | cv::normalize(onefeature, onefeature); 119 | feature.push_back(onefeature.clone()); 120 | } 121 | } 122 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/contrib/graphbuilder.h: -------------------------------------------------------------------------------- 1 | #ifndef GRAPHBUILDER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define GRAPHBUILDER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/mark.h" 11 | #include 12 | 13 | namespace YAML { 14 | class Parser; 15 | 16 | // GraphBuilderInterface 17 | // . Abstraction of node creation 18 | // . pParentNode is always NULL or the return value of one of the NewXXX() 19 | // functions. 20 | class GraphBuilderInterface { 21 | public: 22 | virtual ~GraphBuilderInterface() = 0; 23 | 24 | // Create and return a new node with a null value. 25 | virtual void *NewNull(const Mark &mark, void *pParentNode) = 0; 26 | 27 | // Create and return a new node with the given tag and value. 28 | virtual void *NewScalar(const Mark &mark, const std::string &tag, 29 | void *pParentNode, const std::string &value) = 0; 30 | 31 | // Create and return a new sequence node 32 | virtual void *NewSequence(const Mark &mark, const std::string &tag, 33 | void *pParentNode) = 0; 34 | 35 | // Add pNode to pSequence. pNode was created with one of the NewXxx() 36 | // functions and pSequence with NewSequence(). 37 | virtual void AppendToSequence(void *pSequence, void *pNode) = 0; 38 | 39 | // Note that no moew entries will be added to pSequence 40 | virtual void SequenceComplete(void *pSequence) { (void)pSequence; } 41 | 42 | // Create and return a new map node 43 | virtual void *NewMap(const Mark &mark, const std::string &tag, 44 | void *pParentNode) = 0; 45 | 46 | // Add the pKeyNode => pValueNode mapping to pMap. pKeyNode and pValueNode 47 | // were created with one of the NewXxx() methods and pMap with NewMap(). 48 | virtual void AssignInMap(void *pMap, void *pKeyNode, void *pValueNode) = 0; 49 | 50 | // Note that no more assignments will be made in pMap 51 | virtual void MapComplete(void *pMap) { (void)pMap; } 52 | 53 | // Return the node that should be used in place of an alias referencing 54 | // pNode (pNode by default) 55 | virtual void *AnchorReference(const Mark &mark, void *pNode) { 56 | (void)mark; 57 | return pNode; 58 | } 59 | }; 60 | 61 | // Typesafe wrapper for GraphBuilderInterface. Assumes that Impl defines 62 | // Node, Sequence, and Map types. Sequence and Map must derive from Node 63 | // (unless Node is defined as void). Impl must also implement function with 64 | // all of the same names as the virtual functions in GraphBuilderInterface 65 | // -- including the ones with default implementations -- but with the 66 | // prototypes changed to accept an explicit Node*, Sequence*, or Map* where 67 | // appropriate. 68 | template 69 | class GraphBuilder : public GraphBuilderInterface { 70 | public: 71 | typedef typename Impl::Node Node; 72 | typedef typename Impl::Sequence Sequence; 73 | typedef typename Impl::Map Map; 74 | 75 | GraphBuilder(Impl &impl) : m_impl(impl) { 76 | Map *pMap = NULL; 77 | Sequence *pSeq = NULL; 78 | Node *pNode = NULL; 79 | 80 | // Type consistency checks 81 | pNode = pMap; 82 | pNode = pSeq; 83 | } 84 | 85 | GraphBuilderInterface &AsBuilderInterface() { return *this; } 86 | 87 | virtual void *NewNull(const Mark &mark, void *pParentNode) { 88 | return CheckType(m_impl.NewNull(mark, AsNode(pParentNode))); 89 | } 90 | 91 | virtual void *NewScalar(const Mark &mark, const std::string &tag, 92 | void *pParentNode, const std::string &value) { 93 | return CheckType( 94 | m_impl.NewScalar(mark, tag, AsNode(pParentNode), value)); 95 | } 96 | 97 | virtual void *NewSequence(const Mark &mark, const std::string &tag, 98 | void *pParentNode) { 99 | return CheckType( 100 | m_impl.NewSequence(mark, tag, AsNode(pParentNode))); 101 | } 102 | virtual void AppendToSequence(void *pSequence, void *pNode) { 103 | m_impl.AppendToSequence(AsSequence(pSequence), AsNode(pNode)); 104 | } 105 | virtual void SequenceComplete(void *pSequence) { 106 | m_impl.SequenceComplete(AsSequence(pSequence)); 107 | } 108 | 109 | virtual void *NewMap(const Mark &mark, const std::string &tag, 110 | void *pParentNode) { 111 | return CheckType(m_impl.NewMap(mark, tag, AsNode(pParentNode))); 112 | } 113 | virtual void AssignInMap(void *pMap, void *pKeyNode, void *pValueNode) { 114 | m_impl.AssignInMap(AsMap(pMap), AsNode(pKeyNode), AsNode(pValueNode)); 115 | } 116 | virtual void MapComplete(void *pMap) { m_impl.MapComplete(AsMap(pMap)); } 117 | 118 | virtual void *AnchorReference(const Mark &mark, void *pNode) { 119 | return CheckType(m_impl.AnchorReference(mark, AsNode(pNode))); 120 | } 121 | 122 | private: 123 | Impl &m_impl; 124 | 125 | // Static check for pointer to T 126 | template 127 | static T *CheckType(U *p) { 128 | return p; 129 | } 130 | 131 | static Node *AsNode(void *pNode) { return static_cast(pNode); } 132 | static Sequence *AsSequence(void *pSeq) { 133 | return static_cast(pSeq); 134 | } 135 | static Map *AsMap(void *pMap) { return static_cast(pMap); } 136 | }; 137 | 138 | void *BuildGraphOfNextDocument(Parser &parser, 139 | GraphBuilderInterface &graphBuilder); 140 | 141 | template 142 | typename Impl::Node *BuildGraphOfNextDocument(Parser &parser, Impl &impl) { 143 | GraphBuilder graphBuilder(impl); 144 | return static_cast( 145 | BuildGraphOfNextDocument(parser, graphBuilder)); 146 | } 147 | } 148 | 149 | #endif // GRAPHBUILDER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 150 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/detail/impl.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_DETAIL_IMPL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_DETAIL_IMPL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/node/detail/node.h" 11 | #include "yaml-cpp/node/detail/node_data.h" 12 | #include 13 | 14 | namespace YAML { 15 | namespace detail { 16 | template 17 | struct get_idx { 18 | static node* get(const std::vector& /* sequence */, 19 | const Key& /* key */, shared_memory_holder /* pMemory */) { 20 | return nullptr; 21 | } 22 | }; 23 | 24 | template 25 | struct get_idx::value && 27 | !std::is_same::value>::type> { 28 | static node* get(const std::vector& sequence, const Key& key, 29 | shared_memory_holder /* pMemory */) { 30 | return key < sequence.size() ? sequence[key] : nullptr; 31 | } 32 | 33 | static node* get(std::vector& sequence, const Key& key, 34 | shared_memory_holder pMemory) { 35 | if (key > sequence.size() || (key > 0 && !sequence[key - 1]->is_defined())) 36 | return 0; 37 | if (key == sequence.size()) 38 | sequence.push_back(&pMemory->create_node()); 39 | return sequence[key]; 40 | } 41 | }; 42 | 43 | template 44 | struct get_idx::value>::type> { 45 | static node* get(const std::vector& sequence, const Key& key, 46 | shared_memory_holder pMemory) { 47 | return key >= 0 ? get_idx::get( 48 | sequence, static_cast(key), pMemory) 49 | : nullptr; 50 | } 51 | static node* get(std::vector& sequence, const Key& key, 52 | shared_memory_holder pMemory) { 53 | return key >= 0 ? get_idx::get( 54 | sequence, static_cast(key), pMemory) 55 | : nullptr; 56 | } 57 | }; 58 | 59 | template 60 | struct remove_idx { 61 | static bool remove(std::vector&, const Key&) { return false; } 62 | }; 63 | 64 | template 65 | struct remove_idx< 66 | Key, typename std::enable_if::value && 67 | !std::is_same::value>::type> { 68 | 69 | static bool remove(std::vector& sequence, const Key& key) { 70 | if (key >= sequence.size()) { 71 | return false; 72 | } else { 73 | sequence.erase(sequence.begin() + key); 74 | return true; 75 | } 76 | } 77 | }; 78 | 79 | template 80 | struct remove_idx::value>::type> { 82 | 83 | static bool remove(std::vector& sequence, const Key& key) { 84 | return key >= 0 ? remove_idx::remove( 85 | sequence, static_cast(key)) 86 | : false; 87 | } 88 | }; 89 | 90 | template 91 | inline bool node::equals(const T& rhs, shared_memory_holder pMemory) { 92 | T lhs; 93 | if (convert::decode(Node(*this, pMemory), lhs)) { 94 | return lhs == rhs; 95 | } 96 | return false; 97 | } 98 | 99 | inline bool node::equals(const char* rhs, shared_memory_holder pMemory) { 100 | return equals(rhs, pMemory); 101 | } 102 | 103 | // indexing 104 | template 105 | inline node* node_data::get(const Key& key, 106 | shared_memory_holder pMemory) const { 107 | switch (m_type) { 108 | case NodeType::Map: 109 | break; 110 | case NodeType::Undefined: 111 | case NodeType::Null: 112 | return nullptr; 113 | case NodeType::Sequence: 114 | if (node* pNode = get_idx::get(m_sequence, key, pMemory)) 115 | return pNode; 116 | return nullptr; 117 | case NodeType::Scalar: 118 | throw BadSubscript(key); 119 | } 120 | 121 | for (node_map::const_iterator it = m_map.begin(); it != m_map.end(); ++it) { 122 | if (it->first->equals(key, pMemory)) { 123 | return it->second; 124 | } 125 | } 126 | 127 | return nullptr; 128 | } 129 | 130 | template 131 | inline node& node_data::get(const Key& key, shared_memory_holder pMemory) { 132 | switch (m_type) { 133 | case NodeType::Map: 134 | break; 135 | case NodeType::Undefined: 136 | case NodeType::Null: 137 | case NodeType::Sequence: 138 | if (node* pNode = get_idx::get(m_sequence, key, pMemory)) { 139 | m_type = NodeType::Sequence; 140 | return *pNode; 141 | } 142 | 143 | convert_to_map(pMemory); 144 | break; 145 | case NodeType::Scalar: 146 | throw BadSubscript(key); 147 | } 148 | 149 | for (node_map::const_iterator it = m_map.begin(); it != m_map.end(); ++it) { 150 | if (it->first->equals(key, pMemory)) { 151 | return *it->second; 152 | } 153 | } 154 | 155 | node& k = convert_to_node(key, pMemory); 156 | node& v = pMemory->create_node(); 157 | insert_map_pair(k, v); 158 | return v; 159 | } 160 | 161 | template 162 | inline bool node_data::remove(const Key& key, shared_memory_holder pMemory) { 163 | if (m_type == NodeType::Sequence) { 164 | return remove_idx::remove(m_sequence, key); 165 | } else if (m_type == NodeType::Map) { 166 | kv_pairs::iterator it = m_undefinedPairs.begin(); 167 | while (it != m_undefinedPairs.end()) { 168 | kv_pairs::iterator jt = std::next(it); 169 | if (it->first->equals(key, pMemory)) { 170 | m_undefinedPairs.erase(it); 171 | } 172 | it = jt; 173 | } 174 | 175 | for (node_map::iterator iter = m_map.begin(); iter != m_map.end(); ++iter) { 176 | if (iter->first->equals(key, pMemory)) { 177 | m_map.erase(iter); 178 | return true; 179 | } 180 | } 181 | } 182 | 183 | return false; 184 | } 185 | 186 | // map 187 | template 188 | inline void node_data::force_insert(const Key& key, const Value& value, 189 | shared_memory_holder pMemory) { 190 | switch (m_type) { 191 | case NodeType::Map: 192 | break; 193 | case NodeType::Undefined: 194 | case NodeType::Null: 195 | case NodeType::Sequence: 196 | convert_to_map(pMemory); 197 | break; 198 | case NodeType::Scalar: 199 | throw BadInsert(); 200 | } 201 | 202 | node& k = convert_to_node(key, pMemory); 203 | node& v = convert_to_node(value, pMemory); 204 | insert_map_pair(k, v); 205 | } 206 | 207 | template 208 | inline node& node_data::convert_to_node(const T& rhs, 209 | shared_memory_holder pMemory) { 210 | Node value = convert::encode(rhs); 211 | value.EnsureNodeExists(); 212 | pMemory->merge(*value.m_pMemory); 213 | return *value.m_pNode; 214 | } 215 | } 216 | } 217 | 218 | #endif // NODE_DETAIL_IMPL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 219 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/emitter.h: -------------------------------------------------------------------------------- 1 | #ifndef EMITTER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EMITTER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | #include "yaml-cpp/binary.h" 19 | #include "yaml-cpp/dll.h" 20 | #include "yaml-cpp/emitterdef.h" 21 | #include "yaml-cpp/emittermanip.h" 22 | #include "yaml-cpp/null.h" 23 | #include "yaml-cpp/ostream_wrapper.h" 24 | 25 | namespace YAML { 26 | class Binary; 27 | struct _Null; 28 | } // namespace YAML 29 | 30 | namespace YAML { 31 | class EmitterState; 32 | 33 | class YAML_CPP_API Emitter { 34 | public: 35 | Emitter(); 36 | explicit Emitter(std::ostream& stream); 37 | Emitter(const Emitter&) = delete; 38 | Emitter& operator=(const Emitter&) = delete; 39 | ~Emitter(); 40 | 41 | // output 42 | const char* c_str() const; 43 | std::size_t size() const; 44 | 45 | // state checking 46 | bool good() const; 47 | const std::string GetLastError() const; 48 | 49 | // global setters 50 | bool SetOutputCharset(EMITTER_MANIP value); 51 | bool SetStringFormat(EMITTER_MANIP value); 52 | bool SetBoolFormat(EMITTER_MANIP value); 53 | bool SetIntBase(EMITTER_MANIP value); 54 | bool SetSeqFormat(EMITTER_MANIP value); 55 | bool SetMapFormat(EMITTER_MANIP value); 56 | bool SetIndent(std::size_t n); 57 | bool SetPreCommentIndent(std::size_t n); 58 | bool SetPostCommentIndent(std::size_t n); 59 | bool SetFloatPrecision(std::size_t n); 60 | bool SetDoublePrecision(std::size_t n); 61 | 62 | // local setters 63 | Emitter& SetLocalValue(EMITTER_MANIP value); 64 | Emitter& SetLocalIndent(const _Indent& indent); 65 | Emitter& SetLocalPrecision(const _Precision& precision); 66 | 67 | // overloads of write 68 | Emitter& Write(const std::string& str); 69 | Emitter& Write(bool b); 70 | Emitter& Write(char ch); 71 | Emitter& Write(const _Alias& alias); 72 | Emitter& Write(const _Anchor& anchor); 73 | Emitter& Write(const _Tag& tag); 74 | Emitter& Write(const _Comment& comment); 75 | Emitter& Write(const _Null& n); 76 | Emitter& Write(const Binary& binary); 77 | 78 | template 79 | Emitter& WriteIntegralType(T value); 80 | 81 | template 82 | Emitter& WriteStreamable(T value); 83 | 84 | private: 85 | template 86 | void SetStreamablePrecision(std::stringstream&) {} 87 | std::size_t GetFloatPrecision() const; 88 | std::size_t GetDoublePrecision() const; 89 | 90 | void PrepareIntegralStream(std::stringstream& stream) const; 91 | void StartedScalar(); 92 | 93 | private: 94 | void EmitBeginDoc(); 95 | void EmitEndDoc(); 96 | void EmitBeginSeq(); 97 | void EmitEndSeq(); 98 | void EmitBeginMap(); 99 | void EmitEndMap(); 100 | void EmitNewline(); 101 | void EmitKindTag(); 102 | void EmitTag(bool verbatim, const _Tag& tag); 103 | 104 | void PrepareNode(EmitterNodeType::value child); 105 | void PrepareTopNode(EmitterNodeType::value child); 106 | void FlowSeqPrepareNode(EmitterNodeType::value child); 107 | void BlockSeqPrepareNode(EmitterNodeType::value child); 108 | 109 | void FlowMapPrepareNode(EmitterNodeType::value child); 110 | 111 | void FlowMapPrepareLongKey(EmitterNodeType::value child); 112 | void FlowMapPrepareLongKeyValue(EmitterNodeType::value child); 113 | void FlowMapPrepareSimpleKey(EmitterNodeType::value child); 114 | void FlowMapPrepareSimpleKeyValue(EmitterNodeType::value child); 115 | 116 | void BlockMapPrepareNode(EmitterNodeType::value child); 117 | 118 | void BlockMapPrepareLongKey(EmitterNodeType::value child); 119 | void BlockMapPrepareLongKeyValue(EmitterNodeType::value child); 120 | void BlockMapPrepareSimpleKey(EmitterNodeType::value child); 121 | void BlockMapPrepareSimpleKeyValue(EmitterNodeType::value child); 122 | 123 | void SpaceOrIndentTo(bool requireSpace, std::size_t indent); 124 | 125 | const char* ComputeFullBoolName(bool b) const; 126 | bool CanEmitNewline() const; 127 | 128 | private: 129 | std::unique_ptr m_pState; 130 | ostream_wrapper m_stream; 131 | }; 132 | 133 | template 134 | inline Emitter& Emitter::WriteIntegralType(T value) { 135 | if (!good()) 136 | return *this; 137 | 138 | PrepareNode(EmitterNodeType::Scalar); 139 | 140 | std::stringstream stream; 141 | PrepareIntegralStream(stream); 142 | stream << value; 143 | m_stream << stream.str(); 144 | 145 | StartedScalar(); 146 | 147 | return *this; 148 | } 149 | 150 | template 151 | inline Emitter& Emitter::WriteStreamable(T value) { 152 | if (!good()) 153 | return *this; 154 | 155 | PrepareNode(EmitterNodeType::Scalar); 156 | 157 | std::stringstream stream; 158 | SetStreamablePrecision(stream); 159 | 160 | bool special = false; 161 | if (std::is_floating_point::value) { 162 | if ((std::numeric_limits::has_quiet_NaN || 163 | std::numeric_limits::has_signaling_NaN) && 164 | std::isnan(value)) { 165 | special = true; 166 | stream << ".nan"; 167 | } else if (std::numeric_limits::has_infinity) { 168 | if (value == std::numeric_limits::infinity()) { 169 | special = true; 170 | stream << ".inf"; 171 | } else if (value == -std::numeric_limits::infinity()) { 172 | special = true; 173 | stream << "-.inf"; 174 | } 175 | } 176 | } 177 | 178 | if (!special) { 179 | stream << value; 180 | } 181 | m_stream << stream.str(); 182 | 183 | StartedScalar(); 184 | 185 | return *this; 186 | } 187 | 188 | template <> 189 | inline void Emitter::SetStreamablePrecision(std::stringstream& stream) { 190 | stream.precision(static_cast(GetFloatPrecision())); 191 | } 192 | 193 | template <> 194 | inline void Emitter::SetStreamablePrecision(std::stringstream& stream) { 195 | stream.precision(static_cast(GetDoublePrecision())); 196 | } 197 | 198 | // overloads of insertion 199 | inline Emitter& operator<<(Emitter& emitter, const std::string& v) { 200 | return emitter.Write(v); 201 | } 202 | inline Emitter& operator<<(Emitter& emitter, bool v) { 203 | return emitter.Write(v); 204 | } 205 | inline Emitter& operator<<(Emitter& emitter, char v) { 206 | return emitter.Write(v); 207 | } 208 | inline Emitter& operator<<(Emitter& emitter, unsigned char v) { 209 | return emitter.Write(static_cast(v)); 210 | } 211 | inline Emitter& operator<<(Emitter& emitter, const _Alias& v) { 212 | return emitter.Write(v); 213 | } 214 | inline Emitter& operator<<(Emitter& emitter, const _Anchor& v) { 215 | return emitter.Write(v); 216 | } 217 | inline Emitter& operator<<(Emitter& emitter, const _Tag& v) { 218 | return emitter.Write(v); 219 | } 220 | inline Emitter& operator<<(Emitter& emitter, const _Comment& v) { 221 | return emitter.Write(v); 222 | } 223 | inline Emitter& operator<<(Emitter& emitter, const _Null& v) { 224 | return emitter.Write(v); 225 | } 226 | inline Emitter& operator<<(Emitter& emitter, const Binary& b) { 227 | return emitter.Write(b); 228 | } 229 | 230 | inline Emitter& operator<<(Emitter& emitter, const char* v) { 231 | return emitter.Write(std::string(v)); 232 | } 233 | 234 | inline Emitter& operator<<(Emitter& emitter, int v) { 235 | return emitter.WriteIntegralType(v); 236 | } 237 | inline Emitter& operator<<(Emitter& emitter, unsigned int v) { 238 | return emitter.WriteIntegralType(v); 239 | } 240 | inline Emitter& operator<<(Emitter& emitter, short v) { 241 | return emitter.WriteIntegralType(v); 242 | } 243 | inline Emitter& operator<<(Emitter& emitter, unsigned short v) { 244 | return emitter.WriteIntegralType(v); 245 | } 246 | inline Emitter& operator<<(Emitter& emitter, long v) { 247 | return emitter.WriteIntegralType(v); 248 | } 249 | inline Emitter& operator<<(Emitter& emitter, unsigned long v) { 250 | return emitter.WriteIntegralType(v); 251 | } 252 | inline Emitter& operator<<(Emitter& emitter, long long v) { 253 | return emitter.WriteIntegralType(v); 254 | } 255 | inline Emitter& operator<<(Emitter& emitter, unsigned long long v) { 256 | return emitter.WriteIntegralType(v); 257 | } 258 | 259 | inline Emitter& operator<<(Emitter& emitter, float v) { 260 | return emitter.WriteStreamable(v); 261 | } 262 | inline Emitter& operator<<(Emitter& emitter, double v) { 263 | return emitter.WriteStreamable(v); 264 | } 265 | 266 | inline Emitter& operator<<(Emitter& emitter, EMITTER_MANIP value) { 267 | return emitter.SetLocalValue(value); 268 | } 269 | 270 | inline Emitter& operator<<(Emitter& emitter, _Indent indent) { 271 | return emitter.SetLocalIndent(indent); 272 | } 273 | 274 | inline Emitter& operator<<(Emitter& emitter, _Precision precision) { 275 | return emitter.SetLocalPrecision(precision); 276 | } 277 | } // namespace YAML 278 | 279 | #endif // EMITTER_H_62B23520_7C8E_11DE_8A39_0800200C9A66 280 | -------------------------------------------------------------------------------- /src/yolov5.cpp: -------------------------------------------------------------------------------- 1 | #include "yolov5.h" 2 | #include "yaml-cpp/yaml.h" 3 | 4 | YOLOv5::YOLOv5(const YAML::Node &config) { 5 | onnx_file = config["onnx_file"].as(); 6 | engine_file = config["engine_file"].as(); 7 | labels_file = config["labels_file"].as(); 8 | BATCH_SIZE = config["BATCH_SIZE"].as(); 9 | INPUT_CHANNEL = config["INPUT_CHANNEL"].as(); 10 | IMAGE_WIDTH = config["IMAGE_WIDTH"].as(); 11 | IMAGE_HEIGHT = config["IMAGE_HEIGHT"].as(); 12 | obj_threshold = config["obj_threshold"].as(); 13 | nms_threshold = config["nms_threshold"].as(); 14 | agnostic = config["agnostic"].as(); 15 | strides = config["strides"].as>(); 16 | anchors = config["anchors"].as>>(); 17 | class_labels = readClassLabel(labels_file); 18 | CATEGORY = class_labels.size(); 19 | grids = { 20 | {3, int(IMAGE_WIDTH / strides[0]), int(IMAGE_HEIGHT / strides[0])}, 21 | {3, int(IMAGE_WIDTH / strides[1]), int(IMAGE_HEIGHT / strides[1])}, 22 | {3, int(IMAGE_WIDTH / strides[2]), int(IMAGE_HEIGHT / strides[2])}, 23 | }; 24 | class_colors.resize(CATEGORY); 25 | srand((int) time(nullptr)); 26 | for (cv::Scalar &class_color : class_colors) 27 | class_color = cv::Scalar(rand() % 255, rand() % 255, rand() % 255); 28 | } 29 | 30 | YOLOv5::~YOLOv5() = default; 31 | 32 | std::vector> YOLOv5::InferenceImages(std::vector &vec_img) { 33 | auto t_start_pre = std::chrono::high_resolution_clock::now(); 34 | std::vector image_data = prepareImage(vec_img); 35 | auto t_end_pre = std::chrono::high_resolution_clock::now(); 36 | float total_pre = std::chrono::duration(t_end_pre - t_start_pre).count(); 37 | std::cout << "yolov5 prepare image take: " << total_pre << " ms." << std::endl; 38 | auto t_start = std::chrono::high_resolution_clock::now(); 39 | auto *output = ModelInference(image_data); 40 | auto t_end = std::chrono::high_resolution_clock::now(); 41 | float total_inf = std::chrono::duration(t_end - t_start).count(); 42 | std::cout << "yolov5 inference take: " << total_inf << " ms." << std::endl; 43 | auto r_start = std::chrono::high_resolution_clock::now(); 44 | auto boxes = postProcess(vec_img, output); 45 | auto r_end = std::chrono::high_resolution_clock::now(); 46 | float total_res = std::chrono::duration(r_end - r_start).count(); 47 | std::cout << "yolov5 postprocess take: " << total_res << " ms." << std::endl; 48 | delete output; 49 | return boxes; 50 | } 51 | 52 | std::vector YOLOv5::prepareImage(std::vector &vec_img) { 53 | std::vector result(BATCH_SIZE * IMAGE_WIDTH * IMAGE_HEIGHT * INPUT_CHANNEL); 54 | float *data = result.data(); 55 | int index = 0; 56 | for (const cv::Mat &src_img : vec_img) 57 | { 58 | if (!src_img.data) 59 | continue; 60 | float ratio = float(IMAGE_WIDTH) / float(src_img.cols) < float(IMAGE_HEIGHT) / float(src_img.rows) ? float(IMAGE_WIDTH) / float(src_img.cols) : float(IMAGE_HEIGHT) / float(src_img.rows); 61 | cv::Mat flt_img = cv::Mat::zeros(cv::Size(IMAGE_WIDTH, IMAGE_HEIGHT), CV_8UC3); 62 | cv::Mat rsz_img; 63 | cv::resize(src_img, rsz_img, cv::Size(), ratio, ratio); 64 | rsz_img.copyTo(flt_img(cv::Rect(0, 0, rsz_img.cols, rsz_img.rows))); 65 | flt_img.convertTo(flt_img, CV_32FC3, 1.0 / 255); 66 | 67 | //HWC TO CHW 68 | int channelLength = IMAGE_WIDTH * IMAGE_HEIGHT; 69 | std::vector split_img = { 70 | cv::Mat(IMAGE_WIDTH, IMAGE_HEIGHT, CV_32FC1, data + channelLength * (index + 2)), 71 | cv::Mat(IMAGE_WIDTH, IMAGE_HEIGHT, CV_32FC1, data + channelLength * (index + 1)), 72 | cv::Mat(IMAGE_WIDTH, IMAGE_HEIGHT, CV_32FC1, data + channelLength * index) 73 | }; 74 | index += 3; 75 | cv::split(flt_img, split_img); 76 | } 77 | return result; 78 | } 79 | 80 | float *YOLOv5::ModelInference(std::vector image_data) { 81 | auto *out = new float[outSize * BATCH_SIZE]; 82 | if (!image_data.data()) { 83 | std::cout << "prepare images ERROR!" << std::endl; 84 | return out; 85 | } 86 | // DMA the input to the GPU, execute the batch asynchronously, and DMA it back: 87 | cudaMemcpyAsync(buffers[0], image_data.data(), bufferSize[0], cudaMemcpyHostToDevice, stream); 88 | 89 | // do inference 90 | context->execute(BATCH_SIZE, buffers); 91 | cudaMemcpyAsync(out, buffers[1], bufferSize[1], cudaMemcpyDeviceToHost, stream); 92 | cudaStreamSynchronize(stream); 93 | return out; 94 | } 95 | 96 | std::vector> YOLOv5::postProcess(const std::vector &vec_Mat, float *output) { 97 | std::vector> vec_result; 98 | int index = 0; 99 | for (const cv::Mat &src_img : vec_Mat) 100 | { 101 | std::vector result; 102 | float ratio = float(src_img.cols) / float(IMAGE_WIDTH) > float(src_img.rows) / float(IMAGE_HEIGHT) ? float(src_img.cols) / float(IMAGE_WIDTH) : float(src_img.rows) / float(IMAGE_HEIGHT); 103 | float *out = output + index * outSize; 104 | int position = 0; 105 | for (int n = 0; n < (int)grids.size(); n++) 106 | { 107 | for (int c = 0; c < grids[n][0]; c++) 108 | { 109 | std::vector anchor = anchors[n * grids[n][0] + c]; 110 | for (int h = 0; h < grids[n][1]; h++) 111 | for (int w = 0; w < grids[n][2]; w++) 112 | { 113 | float *row = out + position * (CATEGORY + 5); 114 | position++; 115 | DetectRes box; 116 | auto max_pos = std::max_element(row + 5, row + CATEGORY + 5); 117 | box.prob = row[4] * row[max_pos - row]; 118 | if (box.prob < obj_threshold) 119 | continue; 120 | box.classes = max_pos - row - 5; 121 | box.x = (row[0] * 2 - 0.5 + w) / grids[n][1] * IMAGE_WIDTH * ratio; 122 | box.y = (row[1] * 2 - 0.5 + h) / grids[n][2] * IMAGE_HEIGHT * ratio; 123 | box.w = pow(row[2] * 2, 2) * anchor[0] * ratio; 124 | box.h = pow(row[3] * 2, 2) * anchor[1] * ratio; 125 | result.push_back(box); 126 | } 127 | } 128 | } 129 | NmsDetect(result); 130 | vec_result.push_back(result); 131 | index++; 132 | } 133 | return vec_result; 134 | } 135 | 136 | void YOLOv5::NmsDetect(std::vector &detections) { 137 | sort(detections.begin(), detections.end(), [=](const DetectRes &left, const DetectRes &right) { 138 | return left.prob > right.prob; 139 | }); 140 | 141 | for (int i = 0; i < (int)detections.size(); i++) 142 | for (int j = i + 1; j < (int)detections.size(); j++) 143 | { 144 | if (detections[i].classes == detections[j].classes or agnostic) 145 | { 146 | float iou = IOUCalculate(detections[i], detections[j]); 147 | if (iou > nms_threshold) 148 | detections[j].prob = 0; 149 | } 150 | } 151 | 152 | detections.erase(std::remove_if(detections.begin(), detections.end(), [](const DetectRes &det) 153 | { return det.prob == 0; }), detections.end()); 154 | } 155 | 156 | float YOLOv5::IOUCalculate(const DetectRes &det_a, const DetectRes &det_b) { 157 | cv::Point2f center_a(det_a.x, det_a.y); 158 | cv::Point2f center_b(det_b.x, det_b.y); 159 | cv::Point2f left_up(std::min(det_a.x - det_a.w / 2, det_b.x - det_b.w / 2), 160 | std::min(det_a.y - det_a.h / 2, det_b.y - det_b.h / 2)); 161 | cv::Point2f right_down(std::max(det_a.x + det_a.w / 2, det_b.x + det_b.w / 2), 162 | std::max(det_a.y + det_a.h / 2, det_b.y + det_b.h / 2)); 163 | float distance_d = (center_a - center_b).x * (center_a - center_b).x + (center_a - center_b).y * (center_a - center_b).y; 164 | float distance_c = (left_up - right_down).x * (left_up - right_down).x + (left_up - right_down).y * (left_up - right_down).y; 165 | float inter_l = det_a.x - det_a.w / 2 > det_b.x - det_b.w / 2 ? det_a.x - det_a.w / 2 : det_b.x - det_b.w / 2; 166 | float inter_t = det_a.y - det_a.h / 2 > det_b.y - det_b.h / 2 ? det_a.y - det_a.h / 2 : det_b.y - det_b.h / 2; 167 | float inter_r = det_a.x + det_a.w / 2 < det_b.x + det_b.w / 2 ? det_a.x + det_a.w / 2 : det_b.x + det_b.w / 2; 168 | float inter_b = det_a.y + det_a.h / 2 < det_b.y + det_b.h / 2 ? det_a.y + det_a.h / 2 : det_b.y + det_b.h / 2; 169 | if (inter_b < inter_t || inter_r < inter_l) 170 | return 0; 171 | float inter_area = (inter_b - inter_t) * (inter_r - inter_l); 172 | float union_area = det_a.w * det_a.h + det_b.w * det_b.h - inter_area; 173 | if (union_area == 0) 174 | return 0; 175 | else 176 | return inter_area / union_area - distance_d / distance_c; 177 | } 178 | 179 | void YOLOv5::DrawResults(const std::vector> &detections, std::vector &vec_img) { 180 | for (int i = 0; i < (int)vec_img.size(); i++) { 181 | auto org_img = vec_img[i]; 182 | if (!org_img.data) 183 | continue; 184 | auto rects = detections[i]; 185 | cv::cvtColor(org_img, org_img, cv::COLOR_BGR2RGB); 186 | for(const auto &rect : rects) { 187 | char t[256]; 188 | sprintf(t, "%.2f", rect.prob); 189 | std::string name = class_labels[rect.classes] + "-" + t; 190 | cv::putText(org_img, name, cv::Point(rect.x - rect.w / 2, rect.y - rect.h / 2 - 5), 191 | cv::FONT_HERSHEY_COMPLEX, 0.7, class_colors[rect.classes], 2); 192 | cv::Rect rst(rect.x - rect.w / 2, rect.y - rect.h / 2, rect.w, rect.h); 193 | cv::rectangle(org_img, rst, class_colors[rect.classes], 2, cv::LINE_8, 0); 194 | } 195 | } 196 | } 197 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/convert.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_CONVERT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_CONVERT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | #include "yaml-cpp/binary.h" 18 | #include "yaml-cpp/node/impl.h" 19 | #include "yaml-cpp/node/iterator.h" 20 | #include "yaml-cpp/node/node.h" 21 | #include "yaml-cpp/node/type.h" 22 | #include "yaml-cpp/null.h" 23 | 24 | namespace YAML { 25 | class Binary; 26 | struct _Null; 27 | template 28 | struct convert; 29 | } // namespace YAML 30 | 31 | namespace YAML { 32 | namespace conversion { 33 | inline bool IsInfinity(const std::string& input) { 34 | return input == ".inf" || input == ".Inf" || input == ".INF" || 35 | input == "+.inf" || input == "+.Inf" || input == "+.INF"; 36 | } 37 | 38 | inline bool IsNegativeInfinity(const std::string& input) { 39 | return input == "-.inf" || input == "-.Inf" || input == "-.INF"; 40 | } 41 | 42 | inline bool IsNaN(const std::string& input) { 43 | return input == ".nan" || input == ".NaN" || input == ".NAN"; 44 | } 45 | } 46 | 47 | // Node 48 | template <> 49 | struct convert { 50 | static Node encode(const Node& rhs) { return rhs; } 51 | 52 | static bool decode(const Node& node, Node& rhs) { 53 | rhs.reset(node); 54 | return true; 55 | } 56 | }; 57 | 58 | // std::string 59 | template <> 60 | struct convert { 61 | static Node encode(const std::string& rhs) { return Node(rhs); } 62 | 63 | static bool decode(const Node& node, std::string& rhs) { 64 | if (!node.IsScalar()) 65 | return false; 66 | rhs = node.Scalar(); 67 | return true; 68 | } 69 | }; 70 | 71 | // C-strings can only be encoded 72 | template <> 73 | struct convert { 74 | static Node encode(const char*& rhs) { return Node(rhs); } 75 | }; 76 | 77 | template 78 | struct convert { 79 | static Node encode(const char(&rhs)[N]) { return Node(rhs); } 80 | }; 81 | 82 | template <> 83 | struct convert<_Null> { 84 | static Node encode(const _Null& /* rhs */) { return Node(); } 85 | 86 | static bool decode(const Node& node, _Null& /* rhs */) { 87 | return node.IsNull(); 88 | } 89 | }; 90 | 91 | #define YAML_DEFINE_CONVERT_STREAMABLE(type, negative_op) \ 92 | template <> \ 93 | struct convert { \ 94 | static Node encode(const type& rhs) { \ 95 | std::stringstream stream; \ 96 | stream.precision(std::numeric_limits::max_digits10); \ 97 | stream << rhs; \ 98 | return Node(stream.str()); \ 99 | } \ 100 | \ 101 | static bool decode(const Node& node, type& rhs) { \ 102 | if (node.Type() != NodeType::Scalar) \ 103 | return false; \ 104 | const std::string& input = node.Scalar(); \ 105 | std::stringstream stream(input); \ 106 | stream.unsetf(std::ios::dec); \ 107 | if ((stream >> std::noskipws >> rhs) && (stream >> std::ws).eof()) \ 108 | return true; \ 109 | if (std::numeric_limits::has_infinity) { \ 110 | if (conversion::IsInfinity(input)) { \ 111 | rhs = std::numeric_limits::infinity(); \ 112 | return true; \ 113 | } else if (conversion::IsNegativeInfinity(input)) { \ 114 | rhs = negative_op std::numeric_limits::infinity(); \ 115 | return true; \ 116 | } \ 117 | } \ 118 | \ 119 | if (std::numeric_limits::has_quiet_NaN) { \ 120 | if (conversion::IsNaN(input)) { \ 121 | rhs = std::numeric_limits::quiet_NaN(); \ 122 | return true; \ 123 | } \ 124 | } \ 125 | \ 126 | return false; \ 127 | } \ 128 | } 129 | 130 | #define YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(type) \ 131 | YAML_DEFINE_CONVERT_STREAMABLE(type, -) 132 | 133 | #define YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED(type) \ 134 | YAML_DEFINE_CONVERT_STREAMABLE(type, +) 135 | 136 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(int); 137 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(short); 138 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(long); 139 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(long long); 140 | YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED(unsigned); 141 | YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED(unsigned short); 142 | YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED(unsigned long); 143 | YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED(unsigned long long); 144 | 145 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(char); 146 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(signed char); 147 | YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED(unsigned char); 148 | 149 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(float); 150 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(double); 151 | YAML_DEFINE_CONVERT_STREAMABLE_SIGNED(long double); 152 | 153 | #undef YAML_DEFINE_CONVERT_STREAMABLE_SIGNED 154 | #undef YAML_DEFINE_CONVERT_STREAMABLE_UNSIGNED 155 | #undef YAML_DEFINE_CONVERT_STREAMABLE 156 | 157 | // bool 158 | template <> 159 | struct convert { 160 | static Node encode(bool rhs) { return rhs ? Node("true") : Node("false"); } 161 | 162 | YAML_CPP_API static bool decode(const Node& node, bool& rhs); 163 | }; 164 | 165 | // std::map 166 | template 167 | struct convert> { 168 | static Node encode(const std::map& rhs) { 169 | Node node(NodeType::Map); 170 | for (typename std::map::const_iterator it = rhs.begin(); 171 | it != rhs.end(); ++it) 172 | node.force_insert(it->first, it->second); 173 | return node; 174 | } 175 | 176 | static bool decode(const Node& node, std::map& rhs) { 177 | if (!node.IsMap()) 178 | return false; 179 | 180 | rhs.clear(); 181 | for (const_iterator it = node.begin(); it != node.end(); ++it) 182 | #if defined(__GNUC__) && __GNUC__ < 4 183 | // workaround for GCC 3: 184 | rhs[it->first.template as()] = it->second.template as(); 185 | #else 186 | rhs[it->first.as()] = it->second.as(); 187 | #endif 188 | return true; 189 | } 190 | }; 191 | 192 | // std::vector 193 | template 194 | struct convert> { 195 | static Node encode(const std::vector& rhs) { 196 | Node node(NodeType::Sequence); 197 | for (typename std::vector::const_iterator it = rhs.begin(); 198 | it != rhs.end(); ++it) 199 | node.push_back(*it); 200 | return node; 201 | } 202 | 203 | static bool decode(const Node& node, std::vector& rhs) { 204 | if (!node.IsSequence()) 205 | return false; 206 | 207 | rhs.clear(); 208 | for (const_iterator it = node.begin(); it != node.end(); ++it) 209 | #if defined(__GNUC__) && __GNUC__ < 4 210 | // workaround for GCC 3: 211 | rhs.push_back(it->template as()); 212 | #else 213 | rhs.push_back(it->as()); 214 | #endif 215 | return true; 216 | } 217 | }; 218 | 219 | // std::list 220 | template 221 | struct convert> { 222 | static Node encode(const std::list& rhs) { 223 | Node node(NodeType::Sequence); 224 | for (typename std::list::const_iterator it = rhs.begin(); 225 | it != rhs.end(); ++it) 226 | node.push_back(*it); 227 | return node; 228 | } 229 | 230 | static bool decode(const Node& node, std::list& rhs) { 231 | if (!node.IsSequence()) 232 | return false; 233 | 234 | rhs.clear(); 235 | for (const_iterator it = node.begin(); it != node.end(); ++it) 236 | #if defined(__GNUC__) && __GNUC__ < 4 237 | // workaround for GCC 3: 238 | rhs.push_back(it->template as()); 239 | #else 240 | rhs.push_back(it->as()); 241 | #endif 242 | return true; 243 | } 244 | }; 245 | 246 | // std::array 247 | template 248 | struct convert> { 249 | static Node encode(const std::array& rhs) { 250 | Node node(NodeType::Sequence); 251 | for (const auto& element : rhs) { 252 | node.push_back(element); 253 | } 254 | return node; 255 | } 256 | 257 | static bool decode(const Node& node, std::array& rhs) { 258 | if (!isNodeValid(node)) { 259 | return false; 260 | } 261 | 262 | for (auto i = 0u; i < node.size(); ++i) { 263 | #if defined(__GNUC__) && __GNUC__ < 4 264 | // workaround for GCC 3: 265 | rhs[i] = node[i].template as(); 266 | #else 267 | rhs[i] = node[i].as(); 268 | #endif 269 | } 270 | return true; 271 | } 272 | 273 | private: 274 | static bool isNodeValid(const Node& node) { 275 | return node.IsSequence() && node.size() == N; 276 | } 277 | }; 278 | 279 | // std::pair 280 | template 281 | struct convert> { 282 | static Node encode(const std::pair& rhs) { 283 | Node node(NodeType::Sequence); 284 | node.push_back(rhs.first); 285 | node.push_back(rhs.second); 286 | return node; 287 | } 288 | 289 | static bool decode(const Node& node, std::pair& rhs) { 290 | if (!node.IsSequence()) 291 | return false; 292 | if (node.size() != 2) 293 | return false; 294 | 295 | #if defined(__GNUC__) && __GNUC__ < 4 296 | // workaround for GCC 3: 297 | rhs.first = node[0].template as(); 298 | #else 299 | rhs.first = node[0].as(); 300 | #endif 301 | #if defined(__GNUC__) && __GNUC__ < 4 302 | // workaround for GCC 3: 303 | rhs.second = node[1].template as(); 304 | #else 305 | rhs.second = node[1].as(); 306 | #endif 307 | return true; 308 | } 309 | }; 310 | 311 | // binary 312 | template <> 313 | struct convert { 314 | static Node encode(const Binary& rhs) { 315 | return Node(EncodeBase64(rhs.data(), rhs.size())); 316 | } 317 | 318 | static bool decode(const Node& node, Binary& rhs) { 319 | if (!node.IsScalar()) 320 | return false; 321 | 322 | std::vector data = DecodeBase64(node.Scalar()); 323 | if (data.empty() && !node.Scalar().empty()) 324 | return false; 325 | 326 | rhs.swap(data); 327 | return true; 328 | } 329 | }; 330 | } 331 | 332 | #endif // NODE_CONVERT_H_62B23520_7C8E_11DE_8A39_0800200C9A66 333 | -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/exceptions.h: -------------------------------------------------------------------------------- 1 | #ifndef EXCEPTIONS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define EXCEPTIONS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/mark.h" 11 | #include "yaml-cpp/traits.h" 12 | #include 13 | #include 14 | #include 15 | 16 | // This is here for compatibility with older versions of Visual Studio 17 | // which don't support noexcept 18 | #if defined(_MSC_VER) && _MSC_VER < 1900 19 | #define YAML_CPP_NOEXCEPT _NOEXCEPT 20 | #else 21 | #define YAML_CPP_NOEXCEPT noexcept 22 | #endif 23 | 24 | namespace YAML { 25 | // error messages 26 | namespace ErrorMsg { 27 | const char* const YAML_DIRECTIVE_ARGS = 28 | "YAML directives must have exactly one argument"; 29 | const char* const YAML_VERSION = "bad YAML version: "; 30 | const char* const YAML_MAJOR_VERSION = "YAML major version too large"; 31 | const char* const REPEATED_YAML_DIRECTIVE = "repeated YAML directive"; 32 | const char* const TAG_DIRECTIVE_ARGS = 33 | "TAG directives must have exactly two arguments"; 34 | const char* const REPEATED_TAG_DIRECTIVE = "repeated TAG directive"; 35 | const char* const CHAR_IN_TAG_HANDLE = 36 | "illegal character found while scanning tag handle"; 37 | const char* const TAG_WITH_NO_SUFFIX = "tag handle with no suffix"; 38 | const char* const END_OF_VERBATIM_TAG = "end of verbatim tag not found"; 39 | const char* const END_OF_MAP = "end of map not found"; 40 | const char* const END_OF_MAP_FLOW = "end of map flow not found"; 41 | const char* const END_OF_SEQ = "end of sequence not found"; 42 | const char* const END_OF_SEQ_FLOW = "end of sequence flow not found"; 43 | const char* const MULTIPLE_TAGS = 44 | "cannot assign multiple tags to the same node"; 45 | const char* const MULTIPLE_ANCHORS = 46 | "cannot assign multiple anchors to the same node"; 47 | const char* const MULTIPLE_ALIASES = 48 | "cannot assign multiple aliases to the same node"; 49 | const char* const ALIAS_CONTENT = 50 | "aliases can't have any content, *including* tags"; 51 | const char* const INVALID_HEX = "bad character found while scanning hex number"; 52 | const char* const INVALID_UNICODE = "invalid unicode: "; 53 | const char* const INVALID_ESCAPE = "unknown escape character: "; 54 | const char* const UNKNOWN_TOKEN = "unknown token"; 55 | const char* const DOC_IN_SCALAR = "illegal document indicator in scalar"; 56 | const char* const EOF_IN_SCALAR = "illegal EOF in scalar"; 57 | const char* const CHAR_IN_SCALAR = "illegal character in scalar"; 58 | const char* const TAB_IN_INDENTATION = 59 | "illegal tab when looking for indentation"; 60 | const char* const FLOW_END = "illegal flow end"; 61 | const char* const BLOCK_ENTRY = "illegal block entry"; 62 | const char* const MAP_KEY = "illegal map key"; 63 | const char* const MAP_VALUE = "illegal map value"; 64 | const char* const ALIAS_NOT_FOUND = "alias not found after *"; 65 | const char* const ANCHOR_NOT_FOUND = "anchor not found after &"; 66 | const char* const CHAR_IN_ALIAS = 67 | "illegal character found while scanning alias"; 68 | const char* const CHAR_IN_ANCHOR = 69 | "illegal character found while scanning anchor"; 70 | const char* const ZERO_INDENT_IN_BLOCK = 71 | "cannot set zero indentation for a block scalar"; 72 | const char* const CHAR_IN_BLOCK = "unexpected character in block scalar"; 73 | const char* const AMBIGUOUS_ANCHOR = 74 | "cannot assign the same alias to multiple nodes"; 75 | const char* const UNKNOWN_ANCHOR = "the referenced anchor is not defined"; 76 | 77 | const char* const INVALID_NODE = 78 | "invalid node; this may result from using a map iterator as a sequence " 79 | "iterator, or vice-versa"; 80 | const char* const INVALID_SCALAR = "invalid scalar"; 81 | const char* const KEY_NOT_FOUND = "key not found"; 82 | const char* const BAD_CONVERSION = "bad conversion"; 83 | const char* const BAD_DEREFERENCE = "bad dereference"; 84 | const char* const BAD_SUBSCRIPT = "operator[] call on a scalar"; 85 | const char* const BAD_PUSHBACK = "appending to a non-sequence"; 86 | const char* const BAD_INSERT = "inserting in a non-convertible-to-map"; 87 | 88 | const char* const UNMATCHED_GROUP_TAG = "unmatched group tag"; 89 | const char* const UNEXPECTED_END_SEQ = "unexpected end sequence token"; 90 | const char* const UNEXPECTED_END_MAP = "unexpected end map token"; 91 | const char* const SINGLE_QUOTED_CHAR = 92 | "invalid character in single-quoted string"; 93 | const char* const INVALID_ANCHOR = "invalid anchor"; 94 | const char* const INVALID_ALIAS = "invalid alias"; 95 | const char* const INVALID_TAG = "invalid tag"; 96 | const char* const BAD_FILE = "bad file"; 97 | 98 | template 99 | inline const std::string KEY_NOT_FOUND_WITH_KEY( 100 | const T&, typename disable_if>::type* = 0) { 101 | return KEY_NOT_FOUND; 102 | } 103 | 104 | inline const std::string KEY_NOT_FOUND_WITH_KEY(const std::string& key) { 105 | std::stringstream stream; 106 | stream << KEY_NOT_FOUND << ": " << key; 107 | return stream.str(); 108 | } 109 | 110 | template 111 | inline const std::string KEY_NOT_FOUND_WITH_KEY( 112 | const T& key, typename enable_if>::type* = 0) { 113 | std::stringstream stream; 114 | stream << KEY_NOT_FOUND << ": " << key; 115 | return stream.str(); 116 | } 117 | 118 | template 119 | inline const std::string BAD_SUBSCRIPT_WITH_KEY( 120 | const T&, typename disable_if>::type* = nullptr) { 121 | return BAD_SUBSCRIPT; 122 | } 123 | 124 | inline const std::string BAD_SUBSCRIPT_WITH_KEY(const std::string& key) { 125 | std::stringstream stream; 126 | stream << BAD_SUBSCRIPT << " (key: \"" << key << "\")"; 127 | return stream.str(); 128 | } 129 | 130 | template 131 | inline const std::string BAD_SUBSCRIPT_WITH_KEY( 132 | const T& key, typename enable_if>::type* = nullptr) { 133 | std::stringstream stream; 134 | stream << BAD_SUBSCRIPT << " (key: \"" << key << "\")"; 135 | return stream.str(); 136 | } 137 | 138 | inline const std::string INVALID_NODE_WITH_KEY(const std::string& key) { 139 | std::stringstream stream; 140 | if (key.empty()) { 141 | return INVALID_NODE; 142 | } 143 | stream << "invalid node; first invalid key: \"" << key << "\""; 144 | return stream.str(); 145 | } 146 | } 147 | 148 | class YAML_CPP_API Exception : public std::runtime_error { 149 | public: 150 | Exception(const Mark& mark_, const std::string& msg_) 151 | : std::runtime_error(build_what(mark_, msg_)), mark(mark_), msg(msg_) {} 152 | virtual ~Exception() YAML_CPP_NOEXCEPT; 153 | 154 | Exception(const Exception&) = default; 155 | 156 | Mark mark; 157 | std::string msg; 158 | 159 | private: 160 | static const std::string build_what(const Mark& mark, 161 | const std::string& msg) { 162 | if (mark.is_null()) { 163 | return msg; 164 | } 165 | 166 | std::stringstream output; 167 | output << "yaml-cpp: error at line " << mark.line + 1 << ", column " 168 | << mark.column + 1 << ": " << msg; 169 | return output.str(); 170 | } 171 | }; 172 | 173 | class YAML_CPP_API ParserException : public Exception { 174 | public: 175 | ParserException(const Mark& mark_, const std::string& msg_) 176 | : Exception(mark_, msg_) {} 177 | ParserException(const ParserException&) = default; 178 | virtual ~ParserException() YAML_CPP_NOEXCEPT; 179 | }; 180 | 181 | class YAML_CPP_API RepresentationException : public Exception { 182 | public: 183 | RepresentationException(const Mark& mark_, const std::string& msg_) 184 | : Exception(mark_, msg_) {} 185 | RepresentationException(const RepresentationException&) = default; 186 | virtual ~RepresentationException() YAML_CPP_NOEXCEPT; 187 | }; 188 | 189 | // representation exceptions 190 | class YAML_CPP_API InvalidScalar : public RepresentationException { 191 | public: 192 | InvalidScalar(const Mark& mark_) 193 | : RepresentationException(mark_, ErrorMsg::INVALID_SCALAR) {} 194 | InvalidScalar(const InvalidScalar&) = default; 195 | virtual ~InvalidScalar() YAML_CPP_NOEXCEPT; 196 | }; 197 | 198 | class YAML_CPP_API KeyNotFound : public RepresentationException { 199 | public: 200 | template 201 | KeyNotFound(const Mark& mark_, const T& key_) 202 | : RepresentationException(mark_, ErrorMsg::KEY_NOT_FOUND_WITH_KEY(key_)) { 203 | } 204 | KeyNotFound(const KeyNotFound&) = default; 205 | virtual ~KeyNotFound() YAML_CPP_NOEXCEPT; 206 | }; 207 | 208 | template 209 | class YAML_CPP_API TypedKeyNotFound : public KeyNotFound { 210 | public: 211 | TypedKeyNotFound(const Mark& mark_, const T& key_) 212 | : KeyNotFound(mark_, key_), key(key_) {} 213 | virtual ~TypedKeyNotFound() YAML_CPP_NOEXCEPT {} 214 | 215 | T key; 216 | }; 217 | 218 | template 219 | inline TypedKeyNotFound MakeTypedKeyNotFound(const Mark& mark, 220 | const T& key) { 221 | return TypedKeyNotFound(mark, key); 222 | } 223 | 224 | class YAML_CPP_API InvalidNode : public RepresentationException { 225 | public: 226 | InvalidNode(std::string key) 227 | : RepresentationException(Mark::null_mark(), 228 | ErrorMsg::INVALID_NODE_WITH_KEY(key)) {} 229 | InvalidNode(const InvalidNode&) = default; 230 | virtual ~InvalidNode() YAML_CPP_NOEXCEPT; 231 | }; 232 | 233 | class YAML_CPP_API BadConversion : public RepresentationException { 234 | public: 235 | explicit BadConversion(const Mark& mark_) 236 | : RepresentationException(mark_, ErrorMsg::BAD_CONVERSION) {} 237 | BadConversion(const BadConversion&) = default; 238 | virtual ~BadConversion() YAML_CPP_NOEXCEPT; 239 | }; 240 | 241 | template 242 | class TypedBadConversion : public BadConversion { 243 | public: 244 | explicit TypedBadConversion(const Mark& mark_) : BadConversion(mark_) {} 245 | }; 246 | 247 | class YAML_CPP_API BadDereference : public RepresentationException { 248 | public: 249 | BadDereference() 250 | : RepresentationException(Mark::null_mark(), ErrorMsg::BAD_DEREFERENCE) {} 251 | BadDereference(const BadDereference&) = default; 252 | virtual ~BadDereference() YAML_CPP_NOEXCEPT; 253 | }; 254 | 255 | class YAML_CPP_API BadSubscript : public RepresentationException { 256 | public: 257 | template 258 | BadSubscript(const Key& key) 259 | : RepresentationException(Mark::null_mark(), 260 | ErrorMsg::BAD_SUBSCRIPT_WITH_KEY(key)) {} 261 | BadSubscript(const BadSubscript&) = default; 262 | virtual ~BadSubscript() YAML_CPP_NOEXCEPT; 263 | }; 264 | 265 | class YAML_CPP_API BadPushback : public RepresentationException { 266 | public: 267 | BadPushback() 268 | : RepresentationException(Mark::null_mark(), ErrorMsg::BAD_PUSHBACK) {} 269 | BadPushback(const BadPushback&) = default; 270 | virtual ~BadPushback() YAML_CPP_NOEXCEPT; 271 | }; 272 | 273 | class YAML_CPP_API BadInsert : public RepresentationException { 274 | public: 275 | BadInsert() 276 | : RepresentationException(Mark::null_mark(), ErrorMsg::BAD_INSERT) {} 277 | BadInsert(const BadInsert&) = default; 278 | virtual ~BadInsert() YAML_CPP_NOEXCEPT; 279 | }; 280 | 281 | class YAML_CPP_API EmitterException : public Exception { 282 | public: 283 | EmitterException(const std::string& msg_) 284 | : Exception(Mark::null_mark(), msg_) {} 285 | EmitterException(const EmitterException&) = default; 286 | virtual ~EmitterException() YAML_CPP_NOEXCEPT; 287 | }; 288 | 289 | class YAML_CPP_API BadFile : public Exception { 290 | public: 291 | BadFile() : Exception(Mark::null_mark(), ErrorMsg::BAD_FILE) {} 292 | BadFile(const BadFile&) = default; 293 | virtual ~BadFile() YAML_CPP_NOEXCEPT; 294 | }; 295 | } 296 | 297 | #undef YAML_CPP_NOEXCEPT 298 | 299 | #endif // EXCEPTIONS_H_62B23520_7C8E_11DE_8A39_0800200C9A66 300 | -------------------------------------------------------------------------------- /src/tracker.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by linghu8812 on 2021/11/2. 3 | // 4 | 5 | #include "tracker.h" 6 | #include "Hungarian.h" 7 | 8 | ObjectTracker::ObjectTracker(const YAML::Node &config) { 9 | max_age = config["max_age"].as(); 10 | iou_threshold = config["iou_threshold"].as(); 11 | sim_threshold = config["sim_threshold"].as(); 12 | agnostic = config["agnostic"].as(); 13 | labels_file = config["labels_file"].as(); 14 | class_labels = readClassLabel(labels_file); 15 | id_colors.resize(100); 16 | srand((int) time(nullptr)); 17 | for (cv::Scalar &id_color : id_colors) 18 | id_color = cv::Scalar(rand() % 255, rand() % 255, rand() % 255); 19 | } 20 | 21 | ObjectTracker::~ObjectTracker() = default; 22 | 23 | float ObjectTracker::IOUCalculate(const TrackerRes &det_a, const TrackerRes &det_b) { 24 | cv::Point2f center_a(det_a.x, det_a.y); 25 | cv::Point2f center_b(det_b.x, det_b.y); 26 | cv::Point2f left_up(std::min(det_a.x - det_a.w / 2, det_b.x - det_b.w / 2), 27 | std::min(det_a.y - det_a.h / 2, det_b.y - det_b.h / 2)); 28 | cv::Point2f right_down(std::max(det_a.x + det_a.w / 2, det_b.x + det_b.w / 2), 29 | std::max(det_a.y + det_a.h / 2, det_b.y + det_b.h / 2)); 30 | float distance_d = (center_a - center_b).x * (center_a - center_b).x + (center_a - center_b).y * (center_a - center_b).y; 31 | float distance_c = (left_up - right_down).x * (left_up - right_down).x + (left_up - right_down).y * (left_up - right_down).y; 32 | float inter_l = det_a.x - det_a.w / 2 > det_b.x - det_b.w / 2 ? det_a.x - det_a.w / 2 : det_b.x - det_b.w / 2; 33 | float inter_t = det_a.y - det_a.h / 2 > det_b.y - det_b.h / 2 ? det_a.y - det_a.h / 2 : det_b.y - det_b.h / 2; 34 | float inter_r = det_a.x + det_a.w / 2 < det_b.x + det_b.w / 2 ? det_a.x + det_a.w / 2 : det_b.x + det_b.w / 2; 35 | float inter_b = det_a.y + det_a.h / 2 < det_b.y + det_b.h / 2 ? det_a.y + det_a.h / 2 : det_b.y + det_b.h / 2; 36 | if (inter_b < inter_t || inter_r < inter_l) 37 | return 0; 38 | float inter_area = (inter_b - inter_t) * (inter_r - inter_l); 39 | float union_area = det_a.w * det_a.h + det_b.w * det_b.h - inter_area; 40 | if (union_area == 0) 41 | return 0; 42 | else 43 | return inter_area / union_area - distance_d / distance_c; 44 | } 45 | 46 | void ObjectTracker::Alignment(std::vector> mat, std::set &unmatchedDetections, 47 | std::set &unmatchedTrajectories, std::vector &matchedPairs, 48 | int det_num, int trk_num, bool b_iou) { 49 | std::vector assignment; 50 | HungarianAlgorithm HungAlgo; 51 | HungAlgo.Solve(mat, assignment); 52 | 53 | std::set allItems; 54 | std::set matchedItems; 55 | 56 | if (b_iou) { 57 | std::vector detection_index(unmatchedDetections.size()); 58 | std::vector tracker_index(unmatchedTrajectories.size()); 59 | int idx = 0; 60 | for (const int &umd:unmatchedDetections) { 61 | detection_index[idx] = umd; 62 | idx++; 63 | } 64 | idx = 0; 65 | for (const int &umt:unmatchedTrajectories) { 66 | tracker_index[idx] = umt; 67 | idx++; 68 | } 69 | unmatchedDetections.clear(); 70 | unmatchedTrajectories.clear(); 71 | if (det_num > trk_num) { // there are unmatched detections 72 | for (unsigned int n = 0; n < det_num; n++) 73 | allItems.insert(detection_index[n]); 74 | 75 | for (unsigned int i = 0; i < trk_num; ++i) 76 | matchedItems.insert(detection_index[assignment[i]]); 77 | 78 | set_difference(allItems.begin(), allItems.end(), 79 | matchedItems.begin(), matchedItems.end(), 80 | std::insert_iterator>(unmatchedDetections, unmatchedDetections.begin())); 81 | } 82 | else if (det_num < trk_num) { // there are unmatched trajectory/predictions 83 | for (unsigned int i = 0; i < trk_num; ++i) 84 | if (assignment[i] == -1) // unassigned label will be set as -1 in the assignment algorithm 85 | unmatchedTrajectories.insert(tracker_index[i]); 86 | } 87 | for (unsigned int i = 0; i < trk_num; ++i) { 88 | if (assignment[i] == -1) // pass over invalid values 89 | continue; 90 | if (1 - mat[i][assignment[i]] < iou_threshold) { 91 | unmatchedTrajectories.insert(tracker_index[i]); 92 | unmatchedDetections.insert(detection_index[assignment[i]]); 93 | } 94 | else 95 | matchedPairs.emplace_back(cv::Point(tracker_index[i], detection_index[assignment[i]])); 96 | } 97 | } else { 98 | if (det_num > trk_num) { // there are unmatched detections 99 | for (unsigned int n = 0; n < det_num; n++) 100 | allItems.insert(n); 101 | 102 | for (unsigned int i = 0; i < trk_num; ++i) 103 | matchedItems.insert(assignment[i]); 104 | 105 | set_difference(allItems.begin(), allItems.end(), 106 | matchedItems.begin(), matchedItems.end(), 107 | std::insert_iterator>(unmatchedDetections, unmatchedDetections.begin())); 108 | } 109 | else if (det_num < trk_num) { // there are unmatched trajectory/predictions 110 | for (unsigned int i = 0; i < trk_num; ++i) 111 | if (assignment[i] == -1) // unassigned label will be set as -1 in the assignment algorithm 112 | unmatchedTrajectories.insert(i); 113 | } 114 | for (unsigned int i = 0; i < trk_num; ++i) { 115 | if (assignment[i] == -1) // pass over invalid values 116 | continue; 117 | if (1 - mat[i][assignment[i]] < sim_threshold) { 118 | unmatchedTrajectories.insert(i); 119 | unmatchedDetections.insert(assignment[i]); 120 | } 121 | else 122 | matchedPairs.emplace_back(cv::Point(i, assignment[i])); 123 | } 124 | } 125 | } 126 | 127 | void ObjectTracker::FeatureMatching(const std::vector &predict_boxes, std::set &unmatchedDetections, 128 | std::set &unmatchedTrajectories, std::vector &matchedPairs) { 129 | int det_num = tracker_boxes.size(); 130 | int trk_num = predict_boxes.size(); 131 | std::vector> similar_mat(trk_num, std::vector(det_num, 0)); 132 | for (unsigned int i = 0; i < trk_num; i++) { // compute iou matrix as a distance matrix 133 | for (unsigned int j = 0; j < det_num; j++) { 134 | // use 1-iou because the hungarian algorithm computes a minimum-cost assignment. 135 | if (predict_boxes[i].classes == tracker_boxes[j].classes or agnostic) { 136 | similar_mat[i][j] = 1 - predict_boxes[i].feature.dot(tracker_boxes[j].feature); 137 | } else 138 | similar_mat[i][j] = 1; 139 | } 140 | } 141 | Alignment(similar_mat, unmatchedDetections, unmatchedTrajectories, matchedPairs, det_num, trk_num, false); 142 | } 143 | 144 | void ObjectTracker::IOUMatching(const std::vector &predict_boxes, std::set &unmatchedDetections, 145 | std::set &unmatchedTrajectories, std::vector &matchedPairs) { 146 | int det_num = unmatchedDetections.size(); 147 | int trk_num = unmatchedTrajectories.size(); 148 | if (det_num == 0 or trk_num == 0) 149 | return; 150 | std::vector> iou_mat(trk_num, std::vector(det_num, 0)); 151 | int i = 0; 152 | for (const int &umt : unmatchedTrajectories) { // compute iou matrix as a distance matrix 153 | int j = 0; 154 | for (const int &umd : unmatchedDetections) { 155 | if (predict_boxes[umt].classes == tracker_boxes[umd].classes or agnostic) { 156 | iou_mat[i][j] = 1 - IOUCalculate(predict_boxes[umt], tracker_boxes[umd]); 157 | } else 158 | iou_mat[i][j] = 1; 159 | j++; 160 | } 161 | i++; 162 | } 163 | Alignment(iou_mat, unmatchedDetections, unmatchedTrajectories, matchedPairs, det_num, trk_num, true); 164 | } 165 | 166 | void ObjectTracker::update(const std::vector &det_boxes, const std::vector &det_features, 167 | int width, int height) { 168 | tracker_boxes.clear(); 169 | int index = 0; 170 | for (const auto &det_box : det_boxes) { 171 | TrackerRes tracker_box(det_box.classes, det_box.prob, det_box.x, det_box.y, det_box.w, det_box.h, -1); 172 | tracker_box.feature = det_features[index]; 173 | index++; 174 | tracker_boxes.push_back(tracker_box); 175 | } 176 | if (kalman_boxes.empty()) { 177 | for (auto &tracker_box : tracker_boxes) { 178 | StateType rect_box = { tracker_box.x, tracker_box.y, tracker_box.w, tracker_box.h }; 179 | KalmanTracker tracker = KalmanTracker(rect_box, tracker_box.classes, tracker_box.prob); 180 | tracker.m_feature = tracker_box.feature.clone(); 181 | tracker_box.object_id = tracker.m_id; 182 | kalman_boxes.push_back(tracker); 183 | } 184 | return; 185 | } 186 | std::vector predict_boxes; 187 | for (auto it = kalman_boxes.begin(); it != kalman_boxes.end();) 188 | { 189 | cv::Rect_ pBox = (*it).predict(); 190 | 191 | bool is_nan = (isnan(pBox.x)) or (isnan(pBox.y)) or (isnan(pBox.width)) or (isnan(pBox.height)); 192 | bool is_bound = (pBox.x > (float)width) or (pBox.y > (float)height) or 193 | (pBox.x + pBox.width < 0) or (pBox.y + pBox.height < 0); 194 | bool is_illegal = (pBox.width <= 0) or (pBox.height <= 0); 195 | bool time_since_update = it->m_time_since_update > max_age; 196 | 197 | TrackerRes trk_box(it->m_classes, it->m_prob, pBox.x, pBox.y, pBox.width, pBox.height, it->m_id); 198 | trk_box.classes = it->m_classes; 199 | trk_box.feature = it->m_feature; 200 | if (!(time_since_update or is_nan or is_bound or is_illegal)) 201 | { 202 | predict_boxes.push_back(trk_box); 203 | it++; 204 | } 205 | else 206 | { 207 | it = kalman_boxes.erase(it); 208 | //cerr << "Box invalid at frame: " << frame_count << endl; 209 | } 210 | } 211 | 212 | std::set unmatchedDetections; 213 | std::set unmatchedTrajectories; 214 | std::vector matchedPairs; 215 | FeatureMatching(predict_boxes, unmatchedDetections, unmatchedTrajectories, matchedPairs); 216 | IOUMatching(predict_boxes, unmatchedDetections, unmatchedTrajectories, matchedPairs); 217 | 218 | for (auto & matchedPair : matchedPairs) { 219 | int trk_id = matchedPair.x; 220 | int det_id = matchedPair.y; 221 | StateType rect_box = { tracker_boxes[det_id].x, tracker_boxes[det_id].y, 222 | tracker_boxes[det_id].w, tracker_boxes[det_id].h }; 223 | kalman_boxes[trk_id].update(rect_box, tracker_boxes[det_id].classes, tracker_boxes[det_id].prob, 224 | tracker_boxes[det_id].feature); 225 | tracker_boxes[det_id].object_id = kalman_boxes[trk_id].m_id; 226 | } 227 | for (auto umd : unmatchedDetections) { 228 | StateType rect_box = { tracker_boxes[umd].x, tracker_boxes[umd].y, 229 | tracker_boxes[umd].w, tracker_boxes[umd].h }; 230 | KalmanTracker tracker = KalmanTracker(rect_box, tracker_boxes[umd].classes, tracker_boxes[umd].prob); 231 | tracker_boxes[umd].object_id = tracker.m_id; 232 | tracker.m_feature = tracker_boxes[umd].feature.clone(); 233 | kalman_boxes.push_back(tracker); 234 | } 235 | } 236 | 237 | void ObjectTracker::DrawResults(cv::Mat &origin_img) { 238 | cv::cvtColor(origin_img, origin_img, cv::COLOR_BGR2RGB); 239 | for(const auto &tracker_box : tracker_boxes) { 240 | char t[256]; 241 | sprintf(t, "%d", tracker_box.object_id); 242 | std::string name = class_labels[tracker_box.classes] + "-" + t; 243 | cv::putText(origin_img, name, cv::Point(tracker_box.x - tracker_box.w / 2, tracker_box.y - tracker_box.h / 2 - 5), 244 | cv::FONT_HERSHEY_COMPLEX, 0.7, id_colors[tracker_box.object_id % 100], 2); 245 | cv::Rect rst(tracker_box.x - tracker_box.w / 2, tracker_box.y - tracker_box.h / 2, tracker_box.w, tracker_box.h); 246 | cv::rectangle(origin_img, rst, id_colors[tracker_box.object_id % 100], 2, cv::LINE_8, 0); 247 | } 248 | } -------------------------------------------------------------------------------- /depends/yaml-cpp/include/yaml-cpp/node/impl.h: -------------------------------------------------------------------------------- 1 | #ifndef NODE_IMPL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 2 | #define NODE_IMPL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 3 | 4 | #if defined(_MSC_VER) || \ 5 | (defined(__GNUC__) && (__GNUC__ == 3 && __GNUC_MINOR__ >= 4) || \ 6 | (__GNUC__ >= 4)) // GCC supports "pragma once" correctly since 3.4 7 | #pragma once 8 | #endif 9 | 10 | #include "yaml-cpp/exceptions.h" 11 | #include "yaml-cpp/node/detail/memory.h" 12 | #include "yaml-cpp/node/detail/node.h" 13 | #include "yaml-cpp/node/iterator.h" 14 | #include "yaml-cpp/node/node.h" 15 | #include 16 | #include 17 | 18 | namespace YAML { 19 | inline Node::Node() 20 | : m_isValid(true), m_invalidKey{}, m_pMemory(nullptr), m_pNode(nullptr) {} 21 | 22 | inline Node::Node(NodeType::value type) 23 | : m_isValid(true), 24 | m_invalidKey{}, 25 | m_pMemory(new detail::memory_holder), 26 | m_pNode(&m_pMemory->create_node()) { 27 | m_pNode->set_type(type); 28 | } 29 | 30 | template 31 | inline Node::Node(const T& rhs) 32 | : m_isValid(true), 33 | m_invalidKey{}, 34 | m_pMemory(new detail::memory_holder), 35 | m_pNode(&m_pMemory->create_node()) { 36 | Assign(rhs); 37 | } 38 | 39 | inline Node::Node(const detail::iterator_value& rhs) 40 | : m_isValid(rhs.m_isValid), 41 | m_invalidKey(rhs.m_invalidKey), 42 | m_pMemory(rhs.m_pMemory), 43 | m_pNode(rhs.m_pNode) {} 44 | 45 | inline Node::Node(const Node& rhs) 46 | : m_isValid(rhs.m_isValid), 47 | m_invalidKey(rhs.m_invalidKey), 48 | m_pMemory(rhs.m_pMemory), 49 | m_pNode(rhs.m_pNode) {} 50 | 51 | inline Node::Node(Zombie) 52 | : m_isValid(false), m_invalidKey{}, m_pMemory{}, m_pNode(nullptr) {} 53 | 54 | inline Node::Node(Zombie, const std::string& key) 55 | : m_isValid(false), m_invalidKey(key), m_pMemory{}, m_pNode(nullptr) {} 56 | 57 | inline Node::Node(detail::node& node, detail::shared_memory_holder pMemory) 58 | : m_isValid(true), m_invalidKey{}, m_pMemory(pMemory), m_pNode(&node) {} 59 | 60 | inline Node::~Node() {} 61 | 62 | inline void Node::EnsureNodeExists() const { 63 | if (!m_isValid) 64 | throw InvalidNode(m_invalidKey); 65 | if (!m_pNode) { 66 | m_pMemory.reset(new detail::memory_holder); 67 | m_pNode = &m_pMemory->create_node(); 68 | m_pNode->set_null(); 69 | } 70 | } 71 | 72 | inline bool Node::IsDefined() const { 73 | if (!m_isValid) { 74 | return false; 75 | } 76 | return m_pNode ? m_pNode->is_defined() : true; 77 | } 78 | 79 | inline Mark Node::Mark() const { 80 | if (!m_isValid) { 81 | throw InvalidNode(m_invalidKey); 82 | } 83 | return m_pNode ? m_pNode->mark() : Mark::null_mark(); 84 | } 85 | 86 | inline NodeType::value Node::Type() const { 87 | if (!m_isValid) 88 | throw InvalidNode(m_invalidKey); 89 | return m_pNode ? m_pNode->type() : NodeType::Null; 90 | } 91 | 92 | // access 93 | 94 | // template helpers 95 | template 96 | struct as_if { 97 | explicit as_if(const Node& node_) : node(node_) {} 98 | const Node& node; 99 | 100 | T operator()(const S& fallback) const { 101 | if (!node.m_pNode) 102 | return fallback; 103 | 104 | T t; 105 | if (convert::decode(node, t)) 106 | return t; 107 | return fallback; 108 | } 109 | }; 110 | 111 | template 112 | struct as_if { 113 | explicit as_if(const Node& node_) : node(node_) {} 114 | const Node& node; 115 | 116 | std::string operator()(const S& fallback) const { 117 | if (node.Type() != NodeType::Scalar) 118 | return fallback; 119 | return node.Scalar(); 120 | } 121 | }; 122 | 123 | template 124 | struct as_if { 125 | explicit as_if(const Node& node_) : node(node_) {} 126 | const Node& node; 127 | 128 | T operator()() const { 129 | if (!node.m_pNode) 130 | throw TypedBadConversion(node.Mark()); 131 | 132 | T t; 133 | if (convert::decode(node, t)) 134 | return t; 135 | throw TypedBadConversion(node.Mark()); 136 | } 137 | }; 138 | 139 | template <> 140 | struct as_if { 141 | explicit as_if(const Node& node_) : node(node_) {} 142 | const Node& node; 143 | 144 | std::string operator()() const { 145 | if (node.Type() != NodeType::Scalar) 146 | throw TypedBadConversion(node.Mark()); 147 | return node.Scalar(); 148 | } 149 | }; 150 | 151 | // access functions 152 | template 153 | inline T Node::as() const { 154 | if (!m_isValid) 155 | throw InvalidNode(m_invalidKey); 156 | return as_if(*this)(); 157 | } 158 | 159 | template 160 | inline T Node::as(const S& fallback) const { 161 | if (!m_isValid) 162 | return fallback; 163 | return as_if(*this)(fallback); 164 | } 165 | 166 | inline const std::string& Node::Scalar() const { 167 | if (!m_isValid) 168 | throw InvalidNode(m_invalidKey); 169 | return m_pNode ? m_pNode->scalar() : detail::node_data::empty_scalar(); 170 | } 171 | 172 | inline const std::string& Node::Tag() const { 173 | if (!m_isValid) 174 | throw InvalidNode(m_invalidKey); 175 | return m_pNode ? m_pNode->tag() : detail::node_data::empty_scalar(); 176 | } 177 | 178 | inline void Node::SetTag(const std::string& tag) { 179 | if (!m_isValid) 180 | throw InvalidNode(m_invalidKey); 181 | EnsureNodeExists(); 182 | m_pNode->set_tag(tag); 183 | } 184 | 185 | inline EmitterStyle::value Node::Style() const { 186 | if (!m_isValid) 187 | throw InvalidNode(m_invalidKey); 188 | return m_pNode ? m_pNode->style() : EmitterStyle::Default; 189 | } 190 | 191 | inline void Node::SetStyle(EmitterStyle::value style) { 192 | if (!m_isValid) 193 | throw InvalidNode(m_invalidKey); 194 | EnsureNodeExists(); 195 | m_pNode->set_style(style); 196 | } 197 | 198 | // assignment 199 | inline bool Node::is(const Node& rhs) const { 200 | if (!m_isValid || !rhs.m_isValid) 201 | throw InvalidNode(m_invalidKey); 202 | if (!m_pNode || !rhs.m_pNode) 203 | return false; 204 | return m_pNode->is(*rhs.m_pNode); 205 | } 206 | 207 | template 208 | inline Node& Node::operator=(const T& rhs) { 209 | if (!m_isValid) 210 | throw InvalidNode(m_invalidKey); 211 | Assign(rhs); 212 | return *this; 213 | } 214 | 215 | inline Node& Node::operator=(const Node& rhs) { 216 | if (!m_isValid || !rhs.m_isValid) 217 | throw InvalidNode(m_invalidKey); 218 | if (is(rhs)) 219 | return *this; 220 | AssignNode(rhs); 221 | return *this; 222 | } 223 | 224 | inline void Node::reset(const YAML::Node& rhs) { 225 | if (!m_isValid || !rhs.m_isValid) 226 | throw InvalidNode(m_invalidKey); 227 | m_pMemory = rhs.m_pMemory; 228 | m_pNode = rhs.m_pNode; 229 | } 230 | 231 | template 232 | inline void Node::Assign(const T& rhs) { 233 | if (!m_isValid) 234 | throw InvalidNode(m_invalidKey); 235 | AssignData(convert::encode(rhs)); 236 | } 237 | 238 | template <> 239 | inline void Node::Assign(const std::string& rhs) { 240 | if (!m_isValid) 241 | throw InvalidNode(m_invalidKey); 242 | EnsureNodeExists(); 243 | m_pNode->set_scalar(rhs); 244 | } 245 | 246 | inline void Node::Assign(const char* rhs) { 247 | if (!m_isValid) 248 | throw InvalidNode(m_invalidKey); 249 | EnsureNodeExists(); 250 | m_pNode->set_scalar(rhs); 251 | } 252 | 253 | inline void Node::Assign(char* rhs) { 254 | if (!m_isValid) 255 | throw InvalidNode(m_invalidKey); 256 | EnsureNodeExists(); 257 | m_pNode->set_scalar(rhs); 258 | } 259 | 260 | inline void Node::AssignData(const Node& rhs) { 261 | if (!m_isValid || !rhs.m_isValid) 262 | throw InvalidNode(m_invalidKey); 263 | EnsureNodeExists(); 264 | rhs.EnsureNodeExists(); 265 | 266 | m_pNode->set_data(*rhs.m_pNode); 267 | m_pMemory->merge(*rhs.m_pMemory); 268 | } 269 | 270 | inline void Node::AssignNode(const Node& rhs) { 271 | if (!m_isValid || !rhs.m_isValid) 272 | throw InvalidNode(m_invalidKey); 273 | rhs.EnsureNodeExists(); 274 | 275 | if (!m_pNode) { 276 | m_pNode = rhs.m_pNode; 277 | m_pMemory = rhs.m_pMemory; 278 | return; 279 | } 280 | 281 | m_pNode->set_ref(*rhs.m_pNode); 282 | m_pMemory->merge(*rhs.m_pMemory); 283 | m_pNode = rhs.m_pNode; 284 | } 285 | 286 | // size/iterator 287 | inline std::size_t Node::size() const { 288 | if (!m_isValid) 289 | throw InvalidNode(m_invalidKey); 290 | return m_pNode ? m_pNode->size() : 0; 291 | } 292 | 293 | inline const_iterator Node::begin() const { 294 | if (!m_isValid) 295 | return const_iterator(); 296 | return m_pNode ? const_iterator(m_pNode->begin(), m_pMemory) 297 | : const_iterator(); 298 | } 299 | 300 | inline iterator Node::begin() { 301 | if (!m_isValid) 302 | return iterator(); 303 | return m_pNode ? iterator(m_pNode->begin(), m_pMemory) : iterator(); 304 | } 305 | 306 | inline const_iterator Node::end() const { 307 | if (!m_isValid) 308 | return const_iterator(); 309 | return m_pNode ? const_iterator(m_pNode->end(), m_pMemory) : const_iterator(); 310 | } 311 | 312 | inline iterator Node::end() { 313 | if (!m_isValid) 314 | return iterator(); 315 | return m_pNode ? iterator(m_pNode->end(), m_pMemory) : iterator(); 316 | } 317 | 318 | // sequence 319 | template 320 | inline void Node::push_back(const T& rhs) { 321 | if (!m_isValid) 322 | throw InvalidNode(m_invalidKey); 323 | push_back(Node(rhs)); 324 | } 325 | 326 | inline void Node::push_back(const Node& rhs) { 327 | if (!m_isValid || !rhs.m_isValid) 328 | throw InvalidNode(m_invalidKey); 329 | EnsureNodeExists(); 330 | rhs.EnsureNodeExists(); 331 | 332 | m_pNode->push_back(*rhs.m_pNode, m_pMemory); 333 | m_pMemory->merge(*rhs.m_pMemory); 334 | } 335 | 336 | // helpers for indexing 337 | namespace detail { 338 | template 339 | struct to_value_t { 340 | explicit to_value_t(const T& t_) : t(t_) {} 341 | const T& t; 342 | typedef const T& return_type; 343 | 344 | const T& operator()() const { return t; } 345 | }; 346 | 347 | template <> 348 | struct to_value_t { 349 | explicit to_value_t(const char* t_) : t(t_) {} 350 | const char* t; 351 | typedef std::string return_type; 352 | 353 | const std::string operator()() const { return t; } 354 | }; 355 | 356 | template <> 357 | struct to_value_t { 358 | explicit to_value_t(char* t_) : t(t_) {} 359 | const char* t; 360 | typedef std::string return_type; 361 | 362 | const std::string operator()() const { return t; } 363 | }; 364 | 365 | template 366 | struct to_value_t { 367 | explicit to_value_t(const char* t_) : t(t_) {} 368 | const char* t; 369 | typedef std::string return_type; 370 | 371 | const std::string operator()() const { return t; } 372 | }; 373 | 374 | // converts C-strings to std::strings so they can be copied 375 | template 376 | inline typename to_value_t::return_type to_value(const T& t) { 377 | return to_value_t(t)(); 378 | } 379 | } // namespace detail 380 | 381 | template 382 | std::string key_to_string(const Key& key) { 383 | return streamable_to_string::value>().impl(key); 384 | } 385 | 386 | // indexing 387 | template 388 | inline const Node Node::operator[](const Key& key) const { 389 | if (!m_isValid) 390 | throw InvalidNode(m_invalidKey); 391 | EnsureNodeExists(); 392 | detail::node* value = static_cast(*m_pNode).get( 393 | detail::to_value(key), m_pMemory); 394 | if (!value) { 395 | return Node(ZombieNode, key_to_string(key)); 396 | } 397 | return Node(*value, m_pMemory); 398 | } 399 | 400 | template 401 | inline Node Node::operator[](const Key& key) { 402 | if (!m_isValid) 403 | throw InvalidNode(m_invalidKey); 404 | EnsureNodeExists(); 405 | detail::node& value = m_pNode->get(detail::to_value(key), m_pMemory); 406 | return Node(value, m_pMemory); 407 | } 408 | 409 | template 410 | inline bool Node::remove(const Key& key) { 411 | if (!m_isValid) 412 | throw InvalidNode(m_invalidKey); 413 | EnsureNodeExists(); 414 | return m_pNode->remove(detail::to_value(key), m_pMemory); 415 | } 416 | 417 | inline const Node Node::operator[](const Node& key) const { 418 | if (!m_isValid || !key.m_isValid) 419 | throw InvalidNode(m_invalidKey); 420 | EnsureNodeExists(); 421 | key.EnsureNodeExists(); 422 | m_pMemory->merge(*key.m_pMemory); 423 | detail::node* value = 424 | static_cast(*m_pNode).get(*key.m_pNode, m_pMemory); 425 | if (!value) { 426 | return Node(ZombieNode, key_to_string(key)); 427 | } 428 | return Node(*value, m_pMemory); 429 | } 430 | 431 | inline Node Node::operator[](const Node& key) { 432 | if (!m_isValid || !key.m_isValid) 433 | throw InvalidNode(m_invalidKey); 434 | EnsureNodeExists(); 435 | key.EnsureNodeExists(); 436 | m_pMemory->merge(*key.m_pMemory); 437 | detail::node& value = m_pNode->get(*key.m_pNode, m_pMemory); 438 | return Node(value, m_pMemory); 439 | } 440 | 441 | inline bool Node::remove(const Node& key) { 442 | if (!m_isValid || !key.m_isValid) 443 | throw InvalidNode(m_invalidKey); 444 | EnsureNodeExists(); 445 | key.EnsureNodeExists(); 446 | return m_pNode->remove(*key.m_pNode, m_pMemory); 447 | } 448 | 449 | // map 450 | template 451 | inline void Node::force_insert(const Key& key, const Value& value) { 452 | if (!m_isValid) 453 | throw InvalidNode(m_invalidKey); 454 | EnsureNodeExists(); 455 | m_pNode->force_insert(detail::to_value(key), detail::to_value(value), 456 | m_pMemory); 457 | } 458 | 459 | // free functions 460 | inline bool operator==(const Node& lhs, const Node& rhs) { return lhs.is(rhs); } 461 | } // namespace YAML 462 | 463 | #endif // NODE_IMPL_H_62B23520_7C8E_11DE_8A39_0800200C9A66 464 | -------------------------------------------------------------------------------- /src/Hungarian.cpp: -------------------------------------------------------------------------------- 1 | /////////////////////////////////////////////////////////////////////////////// 2 | // Hungarian.cpp: Implementation file for Class HungarianAlgorithm. 3 | // 4 | // This is a C++ wrapper with slight modification of a hungarian algorithm implementation by Markus Buehren. 5 | // The original implementation is a few mex-functions for use in MATLAB, found here: 6 | // http://www.mathworks.com/matlabcentral/fileexchange/6543-functions-for-the-rectangular-assignment-problem 7 | // 8 | // Both this code and the orignal code are published under the BSD license. 9 | // by Cong Ma, 2016 10 | // 11 | 12 | #include 13 | #include 14 | #include "Hungarian.h" 15 | 16 | 17 | HungarianAlgorithm::HungarianAlgorithm(){} 18 | HungarianAlgorithm::~HungarianAlgorithm(){} 19 | 20 | 21 | //********************************************************// 22 | // A single function wrapper for solving assignment problem. 23 | //********************************************************// 24 | double HungarianAlgorithm::Solve(std::vector>& DistMatrix, std::vector& Assignment) 25 | { 26 | unsigned int nRows = DistMatrix.size(); 27 | unsigned int nCols = DistMatrix[0].size(); 28 | 29 | auto *distMatrixIn = new double[nRows * nCols]; 30 | int *assignment = new int[nRows]; 31 | double cost = 0.0; 32 | 33 | // Fill in the distMatrixIn. Mind the index is "i + nRows * j". 34 | // Here the cost matrix of size MxN is defined as a double precision array of N*M elements. 35 | // In the solving functions matrices are seen to be saved MATLAB-internally in row-order. 36 | // (i.e. the matrix [1 2; 3 4] will be stored as a vector [1 3 2 4], NOT [1 2 3 4]). 37 | for (unsigned int i = 0; i < nRows; i++) 38 | for (unsigned int j = 0; j < nCols; j++) 39 | distMatrixIn[i + nRows * j] = DistMatrix[i][j]; 40 | 41 | // call solving function 42 | assignmentoptimal(assignment, &cost, distMatrixIn, nRows, nCols); 43 | 44 | Assignment.clear(); 45 | for (unsigned int r = 0; r < nRows; r++) 46 | Assignment.push_back(assignment[r]); 47 | 48 | delete[] distMatrixIn; 49 | delete[] assignment; 50 | return cost; 51 | } 52 | 53 | 54 | //********************************************************// 55 | // Solve optimal solution for assignment problem using Munkres algorithm, also known as Hungarian Algorithm. 56 | //********************************************************// 57 | void HungarianAlgorithm::assignmentoptimal(int *assignment, double *cost, double *distMatrixIn, int nOfRows, int nOfColumns) 58 | { 59 | double *distMatrix, *distMatrixTemp, *distMatrixEnd, *columnEnd, value, minValue; 60 | bool *coveredColumns, *coveredRows, *starMatrix, *newStarMatrix, *primeMatrix; 61 | int nOfElements, minDim, row, col; 62 | 63 | /* initialization */ 64 | *cost = 0; 65 | for (row = 0; row nOfColumns) */ 130 | { 131 | minDim = nOfColumns; 132 | 133 | for (col = 0; col= 0) 213 | *cost += distMatrix[row + nOfRows*col]; 214 | } 215 | } 216 | 217 | /********************************************************/ 218 | void HungarianAlgorithm::step2a(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim) 219 | { 220 | bool *starMatrixTemp, *columnEnd; 221 | int col; 222 | 223 | /* cover every column containing a starred zero */ 224 | for (col = 0; col 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | 29 | using Severity = nvinfer1::ILogger::Severity; 30 | 31 | class LogStreamConsumerBuffer : public std::stringbuf 32 | { 33 | public: 34 | LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) 35 | : mOutput(stream) 36 | , mPrefix(prefix) 37 | , mShouldLog(shouldLog) 38 | { 39 | } 40 | 41 | LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) 42 | : mOutput(other.mOutput) 43 | { 44 | } 45 | 46 | ~LogStreamConsumerBuffer() 47 | { 48 | // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence 49 | // std::streambuf::pptr() gives a pointer to the current position of the output sequence 50 | // if the pointer to the beginning is not equal to the pointer to the current position, 51 | // call putOutput() to log the output to the stream 52 | if (pbase() != pptr()) 53 | { 54 | putOutput(); 55 | } 56 | } 57 | 58 | // synchronizes the stream buffer and returns 0 on success 59 | // synchronizing the stream buffer consists of inserting the buffer contents into the stream, 60 | // resetting the buffer and flushing the stream 61 | virtual int sync() 62 | { 63 | putOutput(); 64 | return 0; 65 | } 66 | 67 | void putOutput() 68 | { 69 | if (mShouldLog) 70 | { 71 | // prepend timestamp 72 | std::time_t timestamp = std::time(nullptr); 73 | tm* tm_local = std::localtime(×tamp); 74 | std::cout << "["; 75 | std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; 76 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; 77 | std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; 78 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; 79 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; 80 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; 81 | // std::stringbuf::str() gets the string contents of the buffer 82 | // insert the buffer contents pre-appended by the appropriate prefix into the stream 83 | mOutput << mPrefix << str(); 84 | // set the buffer to empty 85 | str(""); 86 | // flush the stream 87 | mOutput.flush(); 88 | } 89 | } 90 | 91 | void setShouldLog(bool shouldLog) 92 | { 93 | mShouldLog = shouldLog; 94 | } 95 | 96 | private: 97 | std::ostream& mOutput; 98 | std::string mPrefix; 99 | bool mShouldLog; 100 | }; 101 | 102 | //! 103 | //! \class LogStreamConsumerBase 104 | //! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer 105 | //! 106 | class LogStreamConsumerBase 107 | { 108 | public: 109 | LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) 110 | : mBuffer(stream, prefix, shouldLog) 111 | { 112 | } 113 | 114 | protected: 115 | LogStreamConsumerBuffer mBuffer; 116 | }; 117 | 118 | //! 119 | //! \class LogStreamConsumer 120 | //! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. 121 | //! Order of base classes is LogStreamConsumerBase and then std::ostream. 122 | //! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field 123 | //! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. 124 | //! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. 125 | //! Please do not change the order of the parent classes. 126 | //! 127 | class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream 128 | { 129 | public: 130 | //! \brief Creates a LogStreamConsumer which logs messages with level severity. 131 | //! Reportable severity determines if the messages are severe enough to be logged. 132 | LogStreamConsumer(Severity reportableSeverity, Severity severity) 133 | : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) 134 | , std::ostream(&mBuffer) // links the stream buffer with the stream 135 | , mShouldLog(severity <= reportableSeverity) 136 | , mSeverity(severity) 137 | { 138 | } 139 | 140 | LogStreamConsumer(LogStreamConsumer&& other) 141 | : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) 142 | , std::ostream(&mBuffer) // links the stream buffer with the stream 143 | , mShouldLog(other.mShouldLog) 144 | , mSeverity(other.mSeverity) 145 | { 146 | } 147 | 148 | void setReportableSeverity(Severity reportableSeverity) 149 | { 150 | mShouldLog = mSeverity <= reportableSeverity; 151 | mBuffer.setShouldLog(mShouldLog); 152 | } 153 | 154 | private: 155 | static std::ostream& severityOstream(Severity severity) 156 | { 157 | return severity >= Severity::kINFO ? std::cout : std::cerr; 158 | } 159 | 160 | static std::string severityPrefix(Severity severity) 161 | { 162 | switch (severity) 163 | { 164 | case Severity::kINTERNAL_ERROR: return "[F] "; 165 | case Severity::kERROR: return "[E] "; 166 | case Severity::kWARNING: return "[W] "; 167 | case Severity::kINFO: return "[I] "; 168 | case Severity::kVERBOSE: return "[V] "; 169 | default: assert(0); return ""; 170 | } 171 | } 172 | 173 | bool mShouldLog; 174 | Severity mSeverity; 175 | }; 176 | 177 | //! \class Logger 178 | //! 179 | //! \brief Class which manages logging of TensorRT tools and samples 180 | //! 181 | //! \details This class provides a common interface for TensorRT tools and samples to log information to the console, 182 | //! and supports logging two types of messages: 183 | //! 184 | //! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) 185 | //! - Test pass/fail messages 186 | //! 187 | //! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is 188 | //! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. 189 | //! 190 | //! In the future, this class could be extended to support dumping test results to a file in some standard format 191 | //! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). 192 | //! 193 | //! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger 194 | //! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT 195 | //! library and messages coming from the sample. 196 | //! 197 | //! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the 198 | //! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger 199 | //! object. 200 | 201 | class Logger : public nvinfer1::ILogger 202 | { 203 | public: 204 | Logger(Severity severity = Severity::kWARNING) 205 | : mReportableSeverity(severity) 206 | { 207 | } 208 | 209 | //! 210 | //! \enum TestResult 211 | //! \brief Represents the state of a given test 212 | //! 213 | enum class TestResult 214 | { 215 | kRUNNING, //!< The test is running 216 | kPASSED, //!< The test passed 217 | kFAILED, //!< The test failed 218 | kWAIVED //!< The test was waived 219 | }; 220 | 221 | //! 222 | //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger 223 | //! \return The nvinfer1::ILogger associated with this Logger 224 | //! 225 | //! TODO Once all samples are updated to use this method to register the logger with TensorRT, 226 | //! we can eliminate the inheritance of Logger from ILogger 227 | //! 228 | nvinfer1::ILogger& getTRTLogger() 229 | { 230 | return *this; 231 | } 232 | 233 | //! 234 | //! \brief Implementation of the nvinfer1::ILogger::log() virtual method 235 | //! 236 | //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the 237 | //! inheritance from nvinfer1::ILogger 238 | //! 239 | void log(Severity severity, const char* msg) override 240 | { 241 | LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; 242 | } 243 | 244 | //! 245 | //! \brief Method for controlling the verbosity of logging output 246 | //! 247 | //! \param severity The logger will only emit messages that have severity of this level or higher. 248 | //! 249 | void setReportableSeverity(Severity severity) 250 | { 251 | mReportableSeverity = severity; 252 | } 253 | 254 | //! 255 | //! \brief Opaque handle that holds logging information for a particular test 256 | //! 257 | //! This object is an opaque handle to information used by the Logger to print test results. 258 | //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used 259 | //! with Logger::reportTest{Start,End}(). 260 | //! 261 | class TestAtom 262 | { 263 | public: 264 | TestAtom(TestAtom&&) = default; 265 | 266 | private: 267 | friend class Logger; 268 | 269 | TestAtom(bool started, const std::string& name, const std::string& cmdline) 270 | : mStarted(started) 271 | , mName(name) 272 | , mCmdline(cmdline) 273 | { 274 | } 275 | 276 | bool mStarted; 277 | std::string mName; 278 | std::string mCmdline; 279 | }; 280 | 281 | //! 282 | //! \brief Define a test for logging 283 | //! 284 | //! \param[in] name The name of the test. This should be a string starting with 285 | //! "TensorRT" and containing dot-separated strings containing 286 | //! the characters [A-Za-z0-9_]. 287 | //! For example, "TensorRT.sample_googlenet" 288 | //! \param[in] cmdline The command line used to reproduce the test 289 | // 290 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 291 | //! 292 | static TestAtom defineTest(const std::string& name, const std::string& cmdline) 293 | { 294 | return TestAtom(false, name, cmdline); 295 | } 296 | 297 | //! 298 | //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments 299 | //! as input 300 | //! 301 | //! \param[in] name The name of the test 302 | //! \param[in] argc The number of command-line arguments 303 | //! \param[in] argv The array of command-line arguments (given as C strings) 304 | //! 305 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 306 | static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) 307 | { 308 | auto cmdline = genCmdlineString(argc, argv); 309 | return defineTest(name, cmdline); 310 | } 311 | 312 | //! 313 | //! \brief Report that a test has started. 314 | //! 315 | //! \pre reportTestStart() has not been called yet for the given testAtom 316 | //! 317 | //! \param[in] testAtom The handle to the test that has started 318 | //! 319 | static void reportTestStart(TestAtom& testAtom) 320 | { 321 | reportTestResult(testAtom, TestResult::kRUNNING); 322 | assert(!testAtom.mStarted); 323 | testAtom.mStarted = true; 324 | } 325 | 326 | //! 327 | //! \brief Report that a test has ended. 328 | //! 329 | //! \pre reportTestStart() has been called for the given testAtom 330 | //! 331 | //! \param[in] testAtom The handle to the test that has ended 332 | //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, 333 | //! TestResult::kFAILED, TestResult::kWAIVED 334 | //! 335 | static void reportTestEnd(const TestAtom& testAtom, TestResult result) 336 | { 337 | assert(result != TestResult::kRUNNING); 338 | assert(testAtom.mStarted); 339 | reportTestResult(testAtom, result); 340 | } 341 | 342 | static int reportPass(const TestAtom& testAtom) 343 | { 344 | reportTestEnd(testAtom, TestResult::kPASSED); 345 | return EXIT_SUCCESS; 346 | } 347 | 348 | static int reportFail(const TestAtom& testAtom) 349 | { 350 | reportTestEnd(testAtom, TestResult::kFAILED); 351 | return EXIT_FAILURE; 352 | } 353 | 354 | static int reportWaive(const TestAtom& testAtom) 355 | { 356 | reportTestEnd(testAtom, TestResult::kWAIVED); 357 | return EXIT_SUCCESS; 358 | } 359 | 360 | static int reportTest(const TestAtom& testAtom, bool pass) 361 | { 362 | return pass ? reportPass(testAtom) : reportFail(testAtom); 363 | } 364 | 365 | Severity getReportableSeverity() const 366 | { 367 | return mReportableSeverity; 368 | } 369 | 370 | private: 371 | //! 372 | //! \brief returns an appropriate string for prefixing a log message with the given severity 373 | //! 374 | static const char* severityPrefix(Severity severity) 375 | { 376 | switch (severity) 377 | { 378 | case Severity::kINTERNAL_ERROR: return "[F] "; 379 | case Severity::kERROR: return "[E] "; 380 | case Severity::kWARNING: return "[W] "; 381 | case Severity::kINFO: return "[I] "; 382 | case Severity::kVERBOSE: return "[V] "; 383 | default: assert(0); return ""; 384 | } 385 | } 386 | 387 | //! 388 | //! \brief returns an appropriate string for prefixing a test result message with the given result 389 | //! 390 | static const char* testResultString(TestResult result) 391 | { 392 | switch (result) 393 | { 394 | case TestResult::kRUNNING: return "RUNNING"; 395 | case TestResult::kPASSED: return "PASSED"; 396 | case TestResult::kFAILED: return "FAILED"; 397 | case TestResult::kWAIVED: return "WAIVED"; 398 | default: assert(0); return ""; 399 | } 400 | } 401 | 402 | //! 403 | //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity 404 | //! 405 | static std::ostream& severityOstream(Severity severity) 406 | { 407 | return severity >= Severity::kINFO ? std::cout : std::cerr; 408 | } 409 | 410 | //! 411 | //! \brief method that implements logging test results 412 | //! 413 | static void reportTestResult(const TestAtom& testAtom, TestResult result) 414 | { 415 | severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " 416 | << testAtom.mCmdline << std::endl; 417 | } 418 | 419 | //! 420 | //! \brief generate a command line string from the given (argc, argv) values 421 | //! 422 | static std::string genCmdlineString(int argc, char const* const* argv) 423 | { 424 | std::stringstream ss; 425 | for (int i = 0; i < argc; i++) 426 | { 427 | if (i > 0) 428 | ss << " "; 429 | ss << argv[i]; 430 | } 431 | return ss.str(); 432 | } 433 | 434 | Severity mReportableSeverity; 435 | }; 436 | 437 | namespace 438 | { 439 | 440 | //! 441 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE 442 | //! 443 | //! Example usage: 444 | //! 445 | //! LOG_VERBOSE(logger) << "hello world" << std::endl; 446 | //! 447 | inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) 448 | { 449 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); 450 | } 451 | 452 | //! 453 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO 454 | //! 455 | //! Example usage: 456 | //! 457 | //! LOG_INFO(logger) << "hello world" << std::endl; 458 | //! 459 | inline LogStreamConsumer LOG_INFO(const Logger& logger) 460 | { 461 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); 462 | } 463 | 464 | //! 465 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING 466 | //! 467 | //! Example usage: 468 | //! 469 | //! LOG_WARN(logger) << "hello world" << std::endl; 470 | //! 471 | inline LogStreamConsumer LOG_WARN(const Logger& logger) 472 | { 473 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); 474 | } 475 | 476 | //! 477 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR 478 | //! 479 | //! Example usage: 480 | //! 481 | //! LOG_ERROR(logger) << "hello world" << std::endl; 482 | //! 483 | inline LogStreamConsumer LOG_ERROR(const Logger& logger) 484 | { 485 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); 486 | } 487 | 488 | //! 489 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR 490 | // ("fatal" severity) 491 | //! 492 | //! Example usage: 493 | //! 494 | //! LOG_FATAL(logger) << "hello world" << std::endl; 495 | //! 496 | inline LogStreamConsumer LOG_FATAL(const Logger& logger) 497 | { 498 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); 499 | } 500 | 501 | } // anonymous namespace 502 | 503 | #endif // TENSORRT_LOGGING_H 504 | --------------------------------------------------------------------------------