Skip to content

Commit

Permalink
feat: bytetrack
Browse files Browse the repository at this point in the history
* feat: add optimized bytetrack

* chore: add libeigen
  • Loading branch information
iChizer0 authored Apr 28, 2024
1 parent eda3b56 commit e56f7ab
Show file tree
Hide file tree
Showing 12 changed files with 1,514 additions and 0 deletions.
3 changes: 3 additions & 0 deletions .gitmodules
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
[submodule "third_party/Eigen3"]
path = third_party/Eigen3
url = https://gitlab.com/libeigen/eigen.git
213 changes: 213 additions & 0 deletions third_party/ByteTrack/BYTETracker.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,213 @@
/*
* MIT License
* Copyright (c) 2021 Yifu Zhang
*
* Modified by nullptr, Apr 15, 2024, Seeed Technology Co.,Ltd
*/

#include "BYTETracker.h"

#include <cstdint>
#include <utility>
#include <vector>

using namespace std;

BYTETracker::BYTETracker(int frame_rate, int track_buffer) {
track_thresh = 0.5;
high_thresh = 0.6;
match_thresh = 0.8;

frame_id = 0;
max_time_lost = int(frame_rate / 30.0 * track_buffer);
}

BYTETracker::~BYTETracker() {}

vector<STrack> BYTETracker::update(const vector<Object>& objects) {
////////////////// Step 1: Get detections //////////////////
this->frame_id += 1;

vector<STrack> activated_stracks;
vector<STrack> refind_stracks;
vector<STrack> removed_stracks;
vector<STrack> lost_stracks;
vector<STrack> detections;
vector<STrack> detections_low;

vector<STrack> detections_cp;
vector<STrack> tracked_stracks_swap;
vector<STrack> resa, resb;
vector<STrack> output_stracks;

vector<STrack*> unconfirmed;
vector<STrack*> tracked_stracks;
vector<STrack*> strack_pool;
vector<STrack*> r_tracked_stracks;

for (const auto& obj : objects) {
vector<float> tlwh_;
tlwh_.resize(4);

tlwh_[0] = obj.rect.x;
tlwh_[1] = obj.rect.y;
tlwh_[2] = obj.rect.width;
tlwh_[3] = obj.rect.height;

float score = obj.prob;
if (score >= track_thresh) {
detections.emplace_back(move(tlwh_), score, obj.id);
} else {
detections_low.emplace_back(move(tlwh_), score, obj.id);
}
}

// Add newly detected tracklets to tracked_stracks
for (size_t i = 0; i < this->tracked_stracks.size(); ++i) {
if (!this->tracked_stracks[i].is_activated)
unconfirmed.push_back(&this->tracked_stracks[i]);
else
tracked_stracks.push_back(&this->tracked_stracks[i]);
}

////////////////// Step 2: First association, with IoU //////////////////
strack_pool = joint_stracks(tracked_stracks, this->lost_stracks);
STrack::multi_predict(strack_pool, this->kalman_filter);

vector<vector<float> > dists;
int dist_size = 0, dist_size_size = 0;
dists = iou_distance(strack_pool, detections, dist_size, dist_size_size);

vector<vector<int> > matches;
vector<int> u_track, u_detection;
linear_assignment(dists, dist_size, dist_size_size, match_thresh, matches, u_track, u_detection);

for (int i = 0; i < matches.size(); ++i) {
STrack* track = strack_pool[matches[i][0]];
STrack* det = &detections[matches[i][1]];
if (track->state == TrackState::Tracked) {
track->update(*det, this->frame_id);
activated_stracks.push_back(*track);
} else {
track->re_activate(*det, this->frame_id, false);
refind_stracks.push_back(*track);
}
}

////////////////// Step 3: Second association, using low score dets //////////////////
for (int i = 0; i < u_detection.size(); ++i) {
detections_cp.push_back(move(detections[u_detection[i]]));
}
detections.clear();
detections.assign(detections_low.begin(), detections_low.end());

for (int i = 0; i < u_track.size(); ++i) {
auto idx = u_track[i];
auto st = strack_pool[idx];
if (st->state == TrackState::Tracked) {
r_tracked_stracks.push_back(st);
}
}

dists.clear();
dists = iou_distance(r_tracked_stracks, detections, dist_size, dist_size_size);

matches.clear();
u_track.clear();
u_detection.clear();
linear_assignment(dists, dist_size, dist_size_size, 0.5, matches, u_track, u_detection);

for (int i = 0; i < matches.size(); ++i) {
STrack* track = r_tracked_stracks[matches[i][0]];
STrack* det = &detections[matches[i][1]];
if (track->state == TrackState::Tracked) {
track->update(*det, this->frame_id);
activated_stracks.push_back(*track);
} else {
track->re_activate(*det, this->frame_id, false);
refind_stracks.push_back(*track);
}
}

for (int i = 0; i < u_track.size(); ++i) {
STrack* track = r_tracked_stracks[u_track[i]];
if (track->state != TrackState::Lost) {
track->mark_lost();
lost_stracks.push_back(*track);
}
}

// Deal with unconfirmed tracks, usually tracks with only one beginning frame
detections.clear();
detections.assign(detections_cp.begin(), detections_cp.end());

dists.clear();
dists = iou_distance(unconfirmed, detections, dist_size, dist_size_size);

matches.clear();
vector<int> u_unconfirmed;
u_detection.clear();
linear_assignment(dists, dist_size, dist_size_size, 0.7, matches, u_unconfirmed, u_detection);

for (int i = 0; i < matches.size(); ++i) {
unconfirmed[matches[i][0]]->update(detections[matches[i][1]], this->frame_id);
activated_stracks.push_back(*unconfirmed[matches[i][0]]);
}

for (int i = 0; i < u_unconfirmed.size(); ++i) {
STrack* track = unconfirmed[u_unconfirmed[i]];
track->mark_removed();
removed_stracks.push_back(*track);
}

////////////////// Step 4: Init new stracks //////////////////
for (int i = 0; i < u_detection.size(); ++i) {
STrack* track = &detections[u_detection[i]];
if (track->score < this->high_thresh) continue;
track->activate(this->kalman_filter, this->frame_id);
activated_stracks.push_back(*track);
}

////////////////// Step 5: Update state //////////////////
for (int i = 0; i < this->lost_stracks.size(); ++i) {
if (this->frame_id - this->lost_stracks[i].end_frame() > this->max_time_lost) {
this->lost_stracks[i].mark_removed();
removed_stracks.push_back(this->lost_stracks[i]);
}
}

for (int i = 0; i < this->tracked_stracks.size(); ++i) {
if (this->tracked_stracks[i].state == TrackState::Tracked) {
tracked_stracks_swap.push_back(this->tracked_stracks[i]);
}
}
this->tracked_stracks.clear();
this->tracked_stracks.assign(tracked_stracks_swap.begin(), tracked_stracks_swap.end());

this->tracked_stracks = joint_stracks(this->tracked_stracks, activated_stracks);
this->tracked_stracks = joint_stracks(this->tracked_stracks, refind_stracks);

this->lost_stracks = sub_stracks(this->lost_stracks, this->tracked_stracks);
for (int i = 0; i < lost_stracks.size(); ++i) {
this->lost_stracks.push_back(lost_stracks[i]);
}

this->lost_stracks = sub_stracks(this->lost_stracks, this->removed_stracks);
for (int i = 0; i < removed_stracks.size(); ++i) {
this->removed_stracks.push_back(removed_stracks[i]);
}

remove_duplicate_stracks(resa, resb, this->tracked_stracks, this->lost_stracks);

this->tracked_stracks.clear();
this->tracked_stracks.assign(resa.begin(), resa.end());
this->lost_stracks.clear();
this->lost_stracks.assign(resb.begin(), resb.end());

for (int i = 0; i < this->tracked_stracks.size(); ++i) {
if (this->tracked_stracks[i].is_activated) {
output_stracks.push_back(this->tracked_stracks[i]);
}
}
return output_stracks;
}
75 changes: 75 additions & 0 deletions third_party/ByteTrack/BYTETracker.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/*
* MIT License
* Copyright (c) 2021 Yifu Zhang
*
* Modified by nullptr, Apr 15, 2024, Seeed Technology Co.,Ltd
*/

#pragma once

#include <cfloat>
#include <climits>
#include <cstdint>
#include <vector>

#include "STrack.h"

class BYTETracker {
public:
struct Object {
Rect4f rect;
int label;
float prob;
int id;
};

public:
BYTETracker(int frame_rate = 30, int track_buffer = 30);
~BYTETracker();

std::vector<STrack> update(const std::vector<Object>& objects);

private:
std::vector<STrack*> joint_stracks(std::vector<STrack*>& tlista, std::vector<STrack>& tlistb);
std::vector<STrack> joint_stracks(std::vector<STrack>& tlista, std::vector<STrack>& tlistb);

std::vector<STrack> sub_stracks(std::vector<STrack>& tlista, std::vector<STrack>& tlistb);
void remove_duplicate_stracks(std::vector<STrack>& resa,
std::vector<STrack>& resb,
std::vector<STrack>& stracksa,
std::vector<STrack>& stracksb);

void linear_assignment(std::vector<std::vector<float> >& cost_matrix,
int cost_matrix_size,
int cost_matrix_size_size,
float thresh,
std::vector<std::vector<int> >& matches,
std::vector<int>& unmatched_a,
std::vector<int>& unmatched_b);
std::vector<std::vector<float> > iou_distance(std::vector<STrack*>& atracks,
std::vector<STrack>& btracks,
int& dist_size,
int& dist_size_size);
std::vector<std::vector<float> > iou_distance(std::vector<STrack>& atracks, std::vector<STrack>& btracks);
std::vector<std::vector<float> > ious(std::vector<std::vector<float> >& atlbrs,
std::vector<std::vector<float> >& btlbrs);

double lapjv(const std::vector<std::vector<float> >& cost,
std::vector<int>& rowsol,
std::vector<int>& colsol,
bool extend_cost = false,
float cost_limit = LONG_MAX,
bool return_cost = true);

private:
float track_thresh;
float high_thresh;
float match_thresh;
int frame_id;
int max_time_lost;

std::vector<STrack> tracked_stracks;
std::vector<STrack> lost_stracks;
std::vector<STrack> removed_stracks;
byte_kalman::KalmanFilter kalman_filter;
};
21 changes: 21 additions & 0 deletions third_party/ByteTrack/LICENSE
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
MIT License

Copyright (c) 2021 Yifu Zhang

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
Loading

0 comments on commit e56f7ab

Please sign in to comment.