update

Signed-off-by: wep21 <daisuke.nishimatsu1021@gmail.com>

Author: wep21

MODULE.bazel

@@ -29,7 +29,7 @@ register_toolchains("@rust_toolchains//:all")
 crate = use_extension("@rules_rust//crate_universe:extensions.bzl", "crate")
 crate.spec(
     package = "image",
-    version = "0.24",
+    version = "0.25.0",
 )
 crate.spec(
     features = ["ndarray"],
@@ -64,6 +64,26 @@ crate.spec(
     package = "rand",
     version = "0.8.5",
 )
+crate.spec(
+    package = "ab_glyph",
+    version = "0.2.29",
+)
+crate.spec(
+    package = "imageproc",
+    version = "0.25.0",
+)
+crate.annotation(
+    build_script_env = dict(
+        CARGO_PKG_AUTHORS = "",
+        CARGO_PKG_DESCRIPTION = "",
+        CARGO_PKG_HOMEPAGE = "",
+        CARGO_PKG_LICENSE = "",
+        CARGO_PKG_REPOSITORY = "",
+        RUSTDOC = "",
+    ),
+    crate = "rav1e",
+    repositories = ["crates"],
+)
 crate.from_specs()
 use_repo(crate, "crates")
 

detector/src/rtmo.rs

@@ -1,27 +1,27 @@
 #[cxx::bridge(namespace = "rtmo")]
 pub mod ffi {
-    #[derive(Debug)]
+    #[derive(Debug, Clone)]
     struct Point {
         x: f32,
         y: f32,
     }
 
-    #[derive(Debug)]
+    #[derive(Debug, Clone)]
     struct Bbox {
         tl: Point,
         br: Point,
         score: f32,
         class_index: i32,
     }
 
-    #[derive(Debug)]
+    #[derive(Debug, Clone)]
     struct Keypoint {
         x: f32,
         y: f32,
         score: f32,
     }
 
-    #[derive(Debug)]
+    #[derive(Debug, Clone)]
     struct PoseResult {
         keypoints: Vec<Keypoint>,
         bbox: Bbox,

examples/BUILD.bazel

@@ -19,8 +19,11 @@ rust_binary(
     ],
     deps = [
         "//detector:rtmo",
+        "//tracker:tracker",
         "@crates//:image",
+        "@crates//:imageproc",
         "@crates//:minifb",
+        "@crates//:ab_glyph",
         "@crates//:video-rs",
         "@cxx.rs//:cxx",
     ],

examples/src/video_demo.rs

@@ -1,8 +1,11 @@
+use ab_glyph::FontArc;
 use cxx::let_cxx_string;
 use cxx::CxxVector;
 use image::{ImageBuffer, Rgb};
+use imageproc::drawing::draw_text_mut;
 use minifb::{Key, Window, WindowOptions};
 use rtmo::rtmo::ffi::{make_rtmo, PoseResult};
+use tracker::byte_tracker::{ByteTracker, PoseResultWithTrackID};
 use video_rs::decode::Decoder;
 use video_rs::location::Location;
 
@@ -23,6 +26,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         (height as u32).try_into().unwrap(),
     );
     let mut detector = binding.pin_mut();
+    let mut tracker = ByteTracker::new(12, 30, 0.5, 0.6, 0.8);
     let mut pose_results = CxxVector::<PoseResult>::new();
 
     // Create a window for displaying frames
@@ -33,6 +37,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         WindowOptions::default(),
     )?;
 
+    let font_data: &[u8] = include_bytes!("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf");
+    let font = FontArc::try_from_slice(font_data).unwrap();
+    let scale = 36.0f32;
+
     for (_, frame) in decoder
         .decode_iter()
         .take_while(Result::is_ok)
@@ -49,20 +57,34 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         let inference_time = start_time.elapsed();
         println!("Inference status: {}", status);
         println!("Inference time: {:?}", inference_time);
+        let pose_results_with_track_id: Vec<PoseResultWithTrackID> = pose_results
+            .iter()
+            .map(|pose_result| PoseResultWithTrackID {
+                track_id: None,
+                pose: pose_result.clone(),
+            })
+            .collect();
+
+        let track_results = tracker
+            .update(&pose_results_with_track_id)
+            .unwrap_or_else(|e| {
+                eprintln!("Error updating tracker: {}", e);
+                vec![]
+            });
 
         // Draw inference results on the frame
         let mut annotated_image = ImageBuffer::from_fn(width as u32, height as u32, |x, y| {
             let pixel = rgb_image.get_pixel(x, y);
             *pixel
         });
 
-        annotated_image = pose_results
+        annotated_image = track_results
             .iter()
-            .fold(annotated_image, |mut img, pose_result| {
-                let x_min = pose_result.bbox.tl.x as u32;
-                let y_min = pose_result.bbox.tl.y as u32;
-                let x_max = std::cmp::min(pose_result.bbox.br.x as u32, width - 1);
-                let y_max = std::cmp::min(pose_result.bbox.br.y as u32, height - 1);
+            .fold(annotated_image, |mut img, track_result| {
+                let x_min = track_result.pose.bbox.tl.x as u32;
+                let y_min = track_result.pose.bbox.tl.y as u32;
+                let x_max = std::cmp::min(track_result.pose.bbox.br.x as u32, width - 1);
+                let y_max = std::cmp::min(track_result.pose.bbox.br.y as u32, height - 1);
 
                 // Draw bounding box
                 for x in x_min..=x_max {
@@ -75,7 +97,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 }
 
                 // Draw keypoints
-                for keypoint in &pose_result.keypoints {
+                for keypoint in &track_result.pose.keypoints {
                     let x = keypoint.x as u32;
                     let y = keypoint.y as u32;
                     if x < width && y < height {
@@ -102,6 +124,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                             }
                         }
                     }
+
+                    // Draw track ID
+                    if let Some(track_id) = track_result.track_id {
+                        draw_text_mut(
+                            &mut img,
+                            Rgb([255, 255, 0]),
+                            x_min as i32,
+                            std::cmp::max(y_min as i32 - scale as i32, 0),
+                            scale,
+                            &font,
+                            &format!("{}", track_id),
+                        );
+                    }
                 }
 
                 img

tracker/BUILD.bazel

@@ -5,7 +5,9 @@ rust_library(
     srcs = glob([
         "src/*.rs",
     ]),
+    visibility = ["//visibility:public"],
     deps = [
+        "//detector:rtmo",
         "@crates//:nalgebra",
         "@crates//:num",
         "@crates//:thiserror",
@@ -19,6 +21,7 @@ rust_test(
         "src/*.rs",
     ]),
     deps = [
+        "//detector:rtmo",
         "@crates//:nalgebra",
         "@crates//:nearly_eq",
         "@crates//:num",

tracker/src/byte_tracker.rs

@@ -1,7 +1,6 @@
 use crate::{
     error::ByteTrackError,
     lapjv::lapjv,
-    object::Object,
     rect::Rect,
     strack::{STrack, STrackState},
 };
@@ -10,16 +9,94 @@ use std::{collections::HashMap, vec};
  * ByteTracker
  * ---------------------------------------------------------------------------- */
 
+#[derive(Debug, Clone)]
+pub struct PoseResultWithTrackID {
+    pub pose: rtmo::rtmo::ffi::PoseResult,
+    pub track_id: Option<usize>,
+}
+
+impl PoseResultWithTrackID {
+    pub fn new(pose: rtmo::rtmo::ffi::PoseResult, track_id: Option<usize>) -> Self {
+        Self { pose, track_id }
+    }
+    pub fn get_rect(&self) -> Rect<f32> {
+        Rect::<f32>::new(
+            self.pose.bbox.tl.x,
+            self.pose.bbox.tl.y,
+            self.pose.bbox.br.x - self.pose.bbox.tl.x,
+            self.pose.bbox.br.y - self.pose.bbox.tl.y,
+        )
+    }
+    pub fn get_prob(&self) -> f32 {
+        self.pose.bbox.score as f32
+    }
+    pub fn get_pose(&self) -> Vec<rtmo::rtmo::ffi::Keypoint> {
+        self.pose.keypoints.clone()
+    }
+}
+
+impl From<STrack> for PoseResultWithTrackID {
+    fn from(strack: STrack) -> Self {
+        let rect = strack.get_rect();
+        let pose = strack.get_pose();
+        let score = strack.get_score();
+        let track_id = strack.get_track_id();
+        PoseResultWithTrackID {
+            pose: rtmo::rtmo::ffi::PoseResult {
+                keypoints: pose,
+                bbox: rtmo::rtmo::ffi::Bbox {
+                    tl: rtmo::rtmo::ffi::Point {
+                        x: rect.x(),
+                        y: rect.y(),
+                    },
+                    br: rtmo::rtmo::ffi::Point {
+                        x: rect.x() + rect.width(),
+                        y: rect.y() + rect.height(),
+                    },
+                    score: score,
+                    class_index: 0, // Assuming class_index is not used
+                },
+            },
+            track_id: Some(track_id),
+        }
+    }
+}
+
+impl From<&STrack> for PoseResultWithTrackID {
+    fn from(strack: &STrack) -> Self {
+        let rect = strack.get_rect();
+        let pose = strack.get_pose();
+        let score = strack.get_score();
+        let track_id = strack.get_track_id();
+        PoseResultWithTrackID {
+            pose: rtmo::rtmo::ffi::PoseResult {
+                keypoints: pose,
+                bbox: rtmo::rtmo::ffi::Bbox {
+                    tl: rtmo::rtmo::ffi::Point {
+                        x: rect.x(),
+                        y: rect.y(),
+                    },
+                    br: rtmo::rtmo::ffi::Point {
+                        x: rect.x() + rect.width(),
+                        y: rect.y() + rect.height(),
+                    },
+                    score: score,
+                    class_index: 0, // Assuming class_index is not used
+                },
+            },
+            track_id: Some(track_id),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct ByteTracker {
     track_thresh: f32,
     high_thresh: f32,
     match_thresh: f32,
     max_time_lost: usize,
-
     frame_id: usize,
     track_id_count: usize,
-
     tracked_stracks: Vec<STrack>,
     lost_stracks: Vec<STrack>,
     removed_stracks: Vec<STrack>,
@@ -38,17 +115,18 @@ impl ByteTracker {
             high_thresh,
             match_thresh,
             max_time_lost: (track_buffer as f32 * frame_rate as f32 / 30.0) as usize,
-
             frame_id: 0,
             track_id_count: 0,
-
             tracked_stracks: Vec::new(),
             lost_stracks: Vec::new(),
             removed_stracks: Vec::new(),
         }
     }
 
-    pub fn update(&mut self, objects: &Vec<Object>) -> Result<Vec<Object>, ByteTrackError> {
+    pub fn update(
+        &mut self,
+        poses: &Vec<PoseResultWithTrackID>,
+    ) -> Result<Vec<PoseResultWithTrackID>, ByteTrackError> {
         self.frame_id += 1;
 
         /* ------------------ Step 1: Get detections ------------------------- */
@@ -57,9 +135,9 @@ impl ByteTracker {
         let mut det_stracks = Vec::new();
         let mut det_low_stracks = Vec::new();
 
-        for obj in objects {
-            let strack = STrack::new(obj.get_rect(), obj.get_prob());
-            if obj.get_prob() >= self.track_thresh {
+        for pose in poses {
+            let strack = STrack::new(pose.get_rect(), pose.get_pose(), pose.get_prob());
+            if pose.get_prob() >= self.track_thresh {
                 det_stracks.push(strack);
             } else {
                 det_low_stracks.push(strack);