BYTETracker.h 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. #pragma once
  2. #include "STrack.h"
  3. #include "BYTETracker.h"
  4. #include "nvdstracker.h"
  5. #include <map>
  6. #include "DataType.h"
  7. #include <iostream>
  8. class NvObject {
  9. public:
  10. float rect[4];
  11. int label;
  12. float prob;
  13. NvMOTObjToTrack *associatedObjectIn;
  14. };
  15. class BYTETracker {
  16. public:
  17. BYTETracker(int frame_rate = 30, int track_buffer = 30);
  18. ~BYTETracker();
  19. vector<STrack> update(const vector<NvObject> &nvObjects);
  20. private:
  21. vector<STrack *> joint_stracks(vector<STrack *> &tlista, vector<STrack> &tlistb);
  22. vector<STrack> joint_stracks(vector<STrack> &tlista, vector<STrack> &tlistb);
  23. vector<STrack> sub_stracks(vector<STrack> &tlista, vector<STrack> &tlistb);
  24. void remove_duplicate_stracks(vector<STrack> &resa, vector<STrack> &resb, vector<STrack> &stracksa,
  25. vector<STrack> &stracksb);
  26. void linear_assignment(vector<vector<float> > &cost_matrix, int cost_matrix_size, int cost_matrix_size_size,
  27. float thresh,
  28. vector<vector<int> > &matches, vector<int> &unmatched_a, vector<int> &unmatched_b);
  29. vector<vector<float>>iou_distance(vector<STrack *> &atracks, vector<STrack> &btracks, int &dist_size, int &dist_size_size);
  30. vector<vector<float>> iou_distance(vector<STrack> &atracks, vector<STrack> &btracks);
  31. vector<vector<float>> ious(vector<vector<float> > &atlbrs, vector<vector<float> > &btlbrs);
  32. double lapjv(const vector<vector<float> > &cost, vector<int> &rowsol, vector<int> &colsol,
  33. bool extend_cost = false, float cost_limit = LONG_MAX, bool return_cost = true);
  34. private:
  35. float track_thresh;
  36. float high_thresh;
  37. float match_thresh;
  38. int frame_id;
  39. int max_time_lost;
  40. vector<STrack> tracked_stracks;
  41. vector<STrack> lost_stracks;
  42. vector<STrack> removed_stracks;
  43. byte_kalman::KalmanFilter kalman_filter;
  44. };