115 lines
5.7 KiB
C++
115 lines
5.7 KiB
C++
#pragma once
|
|
#include <chrono>
|
|
#include <stdexcept>
|
|
#include <stdlib.h>
|
|
#include <memory>
|
|
#include <vector>
|
|
#include <sstream>
|
|
#include "EigenKalmanFilter.h"
|
|
#include "EigenKalmanBBoxTrack.h"
|
|
#include "EigenBaseTrack.h"
|
|
#include "EigenBoundingBoxIoUMatching.h"
|
|
#include "EigenLinearAssignment.h"
|
|
#include "EigenBoundingBoxTrackUtils.h"
|
|
|
|
namespace ByteTrackEigen {
|
|
/**
|
|
* @brief BYTETracker class for tracking objects in video frames using Kalman Filters.
|
|
*
|
|
* This class implements the BYTETrack algorithm, which combines the use of Kalman Filters
|
|
* and Hungarian algorithm for object tracking in video streams.
|
|
*
|
|
* Key functionalities include processing frame detections, updating track states, and
|
|
* maintaining lists of active, lost, and removed tracks.
|
|
*/
|
|
class BYTETracker {
|
|
public:
|
|
/**
|
|
* @brief Constructor for BYTETracker.
|
|
*
|
|
* Initializes the tracker with specific thresholds and parameters for tracking.
|
|
*
|
|
* @param track_thresh Threshold for track confidence.
|
|
* @param track_buffer Size of the buffer to store track history.
|
|
* @param match_thresh Threshold for matching detections to tracks.
|
|
* @param frame_rate Frame rate of the video being processed.
|
|
*/
|
|
BYTETracker(float track_thresh = 0.25, int track_buffer = 30, float match_thresh = 0.8, int frame_rate = 30);
|
|
void update_parameters(int frameRate = 30, int trackBuffer = 30, double trackThreshold = 0.5, double highThreshold = 0.6, double matchThresold = 0.8, bool autoFrameRate = false);
|
|
float getEstimatedFps() const;
|
|
std::vector<KalmanBBoxTrack> update(const Eigen::MatrixXf& output_results, const std::vector<std::string> obj_ids);
|
|
|
|
private:
|
|
// Internal constants and thresholds
|
|
const float BASE_FRAME_RATE = 30.0;
|
|
const float MIN_KEEP_THRESH = 0.1f;
|
|
const float LOWER_CONFIDENCE_MATCHING_THRESHOLD = 0.5;
|
|
const float ACTIVATION_MATCHING_THRESHOLD = 0.7;
|
|
|
|
void estimateFrameRate();
|
|
|
|
// Member variables for tracking settings and state
|
|
float track_thresh; // Threshold for determining track validity.
|
|
float match_thresh; // Threshold for matching detections to existing tracks.
|
|
int frame_id; // Current frame identifier.
|
|
float det_thresh; // Detection threshold for filtering weak detections.
|
|
int buffer_size; // Size of the history buffer for tracks.
|
|
int max_time_lost; // Maximum time a track can be lost before removal.
|
|
int track_buffer_; // Stored track buffer for recalculating max_time_lost.
|
|
|
|
// Frame rate auto-estimation
|
|
bool auto_frame_rate_;
|
|
float estimated_fps_;
|
|
float time_scale_factor_;
|
|
size_t fps_sample_count_;
|
|
std::chrono::steady_clock::time_point last_update_time_;
|
|
bool has_last_update_time_;
|
|
|
|
// Kalman filter and track lists
|
|
KalmanFilter kalman_filter; // Kalman filter for state estimation.
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>> tracked_tracks; // Active tracks.
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>> lost_tracks; // Tracks that are currently lost.
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>> removed_tracks; // Tracks that are removed.
|
|
|
|
LinearAssignment linear_assignment; // For solving assignment problems.
|
|
|
|
// Internal methods for processing detections and tracks
|
|
std::vector<KalmanBBoxTrack>extract_kalman_bbox_tracks(const Eigen::MatrixXf dets, const Eigen::VectorXf scores_keep, const Eigen::VectorXf clss_ids, const std::vector<std::string> obj_ids );
|
|
Eigen::MatrixXf select_matrix_rows_by_indices(const Eigen::MatrixXf matrix, const std::vector<int> indices);
|
|
std::vector<std::string> select_vector_by_indices(const std::vector<std::string> obj_ids, const std::vector<int> indices);
|
|
|
|
std::pair<std::vector<KalmanBBoxTrack>, std::vector<KalmanBBoxTrack>> filter_and_partition_detections(const Eigen::MatrixXf& output_results, const std::vector<std::string> obj_ids);
|
|
std::pair<std::vector<std::shared_ptr<KalmanBBoxTrack>>, std::vector<std::shared_ptr<KalmanBBoxTrack>>> partition_tracks_by_activation();
|
|
std::tuple<std::vector<std::pair<int, int>>, std::set<int>, std::set<int>> assign_tracks_to_detections(
|
|
const std::vector<std::shared_ptr<KalmanBBoxTrack>> tracks,
|
|
const std::vector<KalmanBBoxTrack> detections,
|
|
double thresh
|
|
);
|
|
void update_tracks_from_detections(
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& tracks,
|
|
const std::vector<KalmanBBoxTrack> detections,
|
|
const std::vector<std::pair<int, int>> track_detection_pair_indices,
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& reacquired_tracked_tracks,
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& activated_tracks
|
|
);
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>> extract_active_tracks(
|
|
const std::vector<std::shared_ptr<KalmanBBoxTrack>>& tracks,
|
|
std::set<int> unpaired_track_indices
|
|
);
|
|
void flag_unpaired_tracks_as_lost(
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& currently_tracked_tracks,
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& lost_tracks,
|
|
std::set<int> unpaired_track_indices
|
|
);
|
|
void prune_and_merge_tracked_tracks(
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& reacquired_tracked_tracks,
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& activated_tracks
|
|
);
|
|
void handle_lost_and_removed_tracks(
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& removed_tracks,
|
|
std::vector<std::shared_ptr<KalmanBBoxTrack>>& lost_tracks
|
|
);
|
|
};
|
|
|
|
}
|