#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct BoundingBox {
pub x1: f32,
pub y1: f32,
pub x2: f32,
pub y2: f32,
pub score: f32,
pub class_id: Option<usize>,
}
impl BoundingBox {
pub fn new(x1: f32, y1: f32, x2: f32, y2: f32, score: f32, class_id: Option<usize>) -> Self {
Self {
x1,
y1,
x2,
y2,
score,
class_id,
}
}
pub fn area(&self) -> f32 {
let w = (self.x2 - self.x1).max(0.0);
let h = (self.y2 - self.y1).max(0.0);
w * h
}
pub fn iou(&self, other: &BoundingBox) -> f32 {
let inter_x1 = self.x1.max(other.x1);
let inter_y1 = self.y1.max(other.y1);
let inter_x2 = self.x2.min(other.x2);
let inter_y2 = self.y2.min(other.y2);
if inter_x2 <= inter_x1 || inter_y2 <= inter_y1 {
return 0.0;
}
let inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1);
let union_area = self.area() + other.area() - inter_area;
if union_area <= 0.0 {
0.0
} else {
inter_area / union_area
}
}
pub fn to_xywh(&self) -> (f32, f32, f32, f32) {
let cx = (self.x1 + self.x2) * 0.5;
let cy = (self.y1 + self.y2) * 0.5;
let w = self.x2 - self.x1;
let h = self.y2 - self.y1;
(cx, cy, w, h)
}
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TrackState {
Tentative,
Confirmed,
Lost,
Deleted,
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct Track {
pub track_id: u64,
pub state: TrackState,
pub bbox: BoundingBox,
pub age: usize,
pub hits: usize,
pub time_since_update: usize,
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct SortConfig {
pub max_age: usize,
pub min_hits: usize,
pub iou_threshold: f32,
}
impl Default for SortConfig {
fn default() -> Self {
Self {
max_age: 3,
min_hits: 3,
iou_threshold: 0.3,
}
}
}
impl SortConfig {
pub fn new(max_age: usize, min_hits: usize, iou_threshold: f32) -> Self {
Self {
max_age,
min_hits,
iou_threshold,
}
}
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct ByteTrackConfig {
pub high_thresh: f32,
pub low_thresh: f32,
pub match_thresh: f32,
pub max_age: usize,
pub min_hits: usize,
}
impl Default for ByteTrackConfig {
fn default() -> Self {
Self {
high_thresh: 0.6,
low_thresh: 0.1,
match_thresh: 0.8,
max_age: 30,
min_hits: 3,
}
}
}
impl ByteTrackConfig {
pub fn new(
high_thresh: f32,
low_thresh: f32,
match_thresh: f32,
max_age: usize,
min_hits: usize,
) -> Self {
Self {
high_thresh,
low_thresh,
match_thresh,
max_age,
min_hits,
}
}
}
pub struct TrackerResult {
pub tracks: Vec<Track>,
pub frame_id: usize,
}