64 lines
1.7 KiB
C++
64 lines
1.7 KiB
C++
#pragma once
|
|
#include "NCNNBytekalmanFilter.h"
|
|
#include <unordered_map>
|
|
namespace ByteTrackNCNN {
|
|
enum TrackState { New = 0, Tracked, Lost, Removed };
|
|
class STrack
|
|
{
|
|
public:
|
|
STrack(std::vector<float> tlwh_, float _score,
|
|
int _class_id, float _left, float _top,
|
|
float _right, float _bottom, std::string _object_id);
|
|
~STrack();
|
|
|
|
std::vector<float> static tlbr_to_tlwh(std::vector<float>& tlbr);
|
|
void static multi_predict(std::vector<STrack*>& stracks, ByteTrackNCNN::ByteKalmanFilter& kalman_filter);
|
|
void static_tlwh();
|
|
void static_tlbr();
|
|
std::vector<float> tlwh_to_xyah(std::vector<float> tlwh_tmp);
|
|
std::vector<float> to_xyah();
|
|
void mark_lost();
|
|
void mark_removed();
|
|
int next_id();
|
|
int end_frame();
|
|
void reset_count();
|
|
|
|
void activate(ByteTrackNCNN::ByteKalmanFilter& kalman_filter, int frame_id);
|
|
void activate(ByteTrackNCNN::ByteKalmanFilter& kalman_filter, int frame_id, int track_id_count);
|
|
void re_activate(STrack& new_track, int frame_id, bool new_id = false);
|
|
void update(STrack& new_track, int frame_id);
|
|
|
|
public:
|
|
bool is_activated;
|
|
int track_id;
|
|
int state;
|
|
int count;
|
|
|
|
std::vector<float> _tlwh;
|
|
std::vector<float> tlwh;
|
|
std::vector<float> tlbr;
|
|
float left; // left, top, right, bottom (original bounding of detected object)
|
|
float top;
|
|
float right;
|
|
float bottom;
|
|
int frame_id;
|
|
int tracklet_len;
|
|
int start_frame;
|
|
std::string object_id;
|
|
|
|
KAL_MEAN mean;
|
|
KAL_COVA covariance;
|
|
float score;
|
|
int class_id;
|
|
|
|
private:
|
|
ByteTrackNCNN::ByteKalmanFilter kalman_filter;
|
|
std::unordered_map<int, float> class_id_scores_;
|
|
int detection_count_;
|
|
bool class_id_locked_;
|
|
static const int CLASS_ID_LOCK_FRAMES = 10;
|
|
void voteClassId(int new_class_id, float score);
|
|
};
|
|
}
|
|
|