diff --git a/Cargo.toml b/Cargo.toml index 90dd433..799bd06 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,8 @@ serde_json = "1.0" tempfile = "3.12.0" video-rs = { version = "0.9.0", features = ["ndarray"] } natord = "1.0.9" +tracing = "0.1.40" +tracing-subscriber = "0.3.18" [features] diff --git a/examples/dl/main.rs b/examples/dl/main.rs index b5b1a18..3da7ef3 100644 --- a/examples/dl/main.rs +++ b/examples/dl/main.rs @@ -1,6 +1,10 @@ use usls::{models::YOLO, Annotator, DataLoader, Options, Vision, YOLOTask, YOLOVersion}; fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .init(); + let options = Options::new() .with_cuda(0) .with_model("yolo/v8-m-dyn.onnx")? @@ -17,27 +21,31 @@ fn main() -> anyhow::Result<()> { // "rtsp://admin:zfsoft888@192.168.2.217:554/h265/ch1/", // "rtsp://admin:KCNULU@192.168.2.193:554/h264/ch1/", // "../hall.mp4", - "./assets/bus.jpg", + // "./assets/bus.jpg", // "images/car.jpg", // "../set-negs", // "/home/qweasd/Desktop/coco/val2017/images/test", - // "/home/qweasd/Desktop/SourceVideos/3.mp4", + "/home/qweasd/Desktop/SourceVideos/3.mp4", // "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", )? .with_batch(1) + .with_progress_bar(true) .build()?; // // build annotator - let annotator = Annotator::new().with_saveout("YOLO-DataLoader"); + let annotator = Annotator::new() + .with_bboxes_thickness(4) + .with_saveout("YOLO-DataLoader"); // run for (xs, _) in dl { + // std::thread::sleep(std::time::Duration::from_millis(10000)); let ys = model.forward(&xs, false)?; annotator.annotate(&xs, &ys); } // images -> video - DataLoader::is2v("runs/YOLO-DataLoader", &["runs", "is2v"], 24)?; + // DataLoader::is2v("runs/YOLO-DataLoader", &["runs", "is2v"], 24)?; Ok(()) } diff --git a/src/core/annotator.rs b/src/core/annotator.rs index 99dd774..75a667c 100644 --- a/src/core/annotator.rs +++ b/src/core/annotator.rs @@ -10,7 +10,6 @@ use imageproc::map::map_colors; /// Annotator for struct `Y` // #[derive(Debug)] pub struct Annotator { - verbose: bool, font: FontVec, _scale: f32, // Cope with ab_glyph & imageproc=0.24.0 scale_dy: f32, @@ -65,7 +64,6 @@ pub struct Annotator { impl Default for Annotator { fn default() -> Self { Self { - verbose: true, font: match Self::load_font(None) { Ok(x) => x, Err(err) => panic!("Failed to load font: {}", err), @@ -347,6 +345,9 @@ impl Annotator { /// Plot images and return plotted images(RGBA8) pub fn plot(&self, imgs: &[DynamicImage], ys: &[Y]) -> Result> { + let span = tracing::span!(tracing::Level::INFO, "YOLO-new"); + let _guard = span.enter(); + let mut vs: Vec = Vec::new(); // annotate @@ -396,11 +397,9 @@ impl Annotator { // save let saveout = self.saveout()?.join(format!("{}.png", string_now("-"))); match img_rgba.save(&saveout) { - Err(err) => println!("{} Saving failed: {:?}", CROSS_MARK, err), + Err(err) => tracing::error!("{} Saving failed: {:?}", CROSS_MARK, err), Ok(_) => { - if self.verbose { - println!("{} Annotated image saved to: {:?}", CHECK_MARK, saveout); - } + tracing::info!("{} Annotated image saved to: {:?}", CHECK_MARK, saveout); } } diff --git a/src/core/dataloader.rs b/src/core/dataloader.rs index c0dc4f1..9233c89 100644 --- a/src/core/dataloader.rs +++ b/src/core/dataloader.rs @@ -10,88 +10,133 @@ use video_rs::{ Decoder, Url, }; -use crate::{string_now, Dir, Hub, Location, MediaType, CHECK_MARK, CROSS_MARK}; +use crate::{ + build_progress_bar, string_now, Dir, Hub, Location, MediaType, CHECK_MARK, CROSS_MARK, +}; type TempReturnType = (Vec, Vec); -impl IntoIterator for DataLoader { - type Item = TempReturnType; - type IntoIter = DataLoaderIterator; - - fn into_iter(self) -> Self::IntoIter { - DataLoaderIterator { - receiver: self.receiver, - } - } -} - pub struct DataLoaderIterator { receiver: mpsc::Receiver, + progress_bar: Option, } impl Iterator for DataLoaderIterator { type Item = TempReturnType; fn next(&mut self) -> Option { - self.receiver.recv().ok() + match &self.progress_bar { + None => self.receiver.recv().ok(), + Some(progress_bar) => { + progress_bar.inc(1); + match self.receiver.recv().ok() { + Some(item) => { + // progress_bar.inc(1); + Some(item) + } + None => { + progress_bar.set_prefix(" Iterated"); + progress_bar.finish(); + None + } + } + } + } + } +} + +impl IntoIterator for DataLoader { + type Item = TempReturnType; + type IntoIter = DataLoaderIterator; + + fn into_iter(self) -> Self::IntoIter { + let progress_bar = if self.with_pb { + build_progress_bar( + self.nf / self.batch_size as u64, + " Iterating", + Some(&format!("{:?}", self.media_type)), + "{prefix:.green.bold} {msg} {human_pos}/{human_len} |{bar}| {elapsed_precise}", + ) + .ok() + } else { + None + }; + + DataLoaderIterator { + receiver: self.receiver, + progress_bar, + } } } /// Load images, video, stream pub struct DataLoader { pub paths: Option>, - pub media_type: Option, + pub media_type: MediaType, pub batch_size: usize, sender: Option>, receiver: mpsc::Receiver, pub decoder: Option, + nf: u64, // MAX means live stream + with_pb: bool, } impl Default for DataLoader { fn default() -> Self { Self { paths: None, - media_type: Some(MediaType::Unknown), + media_type: MediaType::Unknown, batch_size: 1, sender: None, receiver: mpsc::channel().1, decoder: None, + nf: 0, + with_pb: true, } } } impl DataLoader { pub fn new(source: &str) -> Result { + let span = tracing::span!(tracing::Level::INFO, "DataLoader-new"); + let _guard = span.enter(); + + // Number of frames or stream + let mut nf = 0; + // paths & media_type let source_path = Path::new(source); let (paths, media_type) = match source_path.exists() { false => { // remote + nf = 1; ( Some(VecDeque::from([source_path.to_path_buf()])), - Some(MediaType::from_url(source)), + MediaType::from_url(source), ) } true => { // local if source_path.is_file() { + nf = 1; ( Some(VecDeque::from([source_path.to_path_buf()])), - Some(MediaType::from_path(source_path)), + MediaType::from_path(source_path), ) } else if source_path.is_dir() { let paths_sorted = Self::load_from_folder(source_path)?; + nf = paths_sorted.len() as _; ( Some(VecDeque::from(paths_sorted)), - Some(MediaType::Image(Location::Local)), + MediaType::Image(Location::Local), ) } else { - (None, Some(MediaType::Unknown)) + (None, MediaType::Unknown) } } }; - if let Some(MediaType::Unknown) = media_type { + if let MediaType::Unknown = media_type { anyhow::bail!("Could not locate the source path: {:?}", source_path); } @@ -100,20 +145,21 @@ impl DataLoader { // decoder let decoder = match &media_type { - Some(MediaType::Video(Location::Local)) => Some(Decoder::new(source_path)?), - Some(MediaType::Video(Location::Remote)) | Some(MediaType::Stream) => { + MediaType::Video(Location::Local) => Some(Decoder::new(source_path)?), + MediaType::Video(Location::Remote) | MediaType::Stream => { let location: video_rs::location::Location = source.parse::()?.into(); Some(Decoder::new(location)?) } _ => None, }; + // get frames + if let Some(decoder) = &decoder { + nf = decoder.frames().unwrap_or(u64::MAX); + } + // summary - println!( - "{CHECK_MARK} Found {:?} x{}", - media_type.as_ref().unwrap_or(&MediaType::Unknown), - paths.as_ref().map_or(0, |p| p.len()) - ); + tracing::info!("{} Found {:?} x{}", CHECK_MARK, media_type, nf,); Ok(DataLoader { paths, @@ -122,6 +168,8 @@ impl DataLoader { receiver, batch_size: 1, decoder, + nf, + with_pb: true, }) } @@ -129,7 +177,8 @@ impl DataLoader { let sender = self.sender.take().expect("Sender should be available"); let batch_size = self.batch_size; let data = self.paths.take().unwrap_or_default(); - let media_type = self.media_type.take().unwrap_or(MediaType::Unknown); + // let media_type = self.media_type.take().unwrap_or(MediaType::Unknown); + let media_type = self.media_type.clone(); let decoder = self.decoder.take(); // Spawn the producer thread @@ -147,6 +196,8 @@ impl DataLoader { media_type: MediaType, mut decoder: Option, ) { + let span = tracing::span!(tracing::Level::INFO, "DataLoader-producer-thread"); + let _guard = span.enter(); let mut yis: Vec = Vec::with_capacity(batch_size); let mut yps: Vec = Vec::with_capacity(batch_size); @@ -155,7 +206,7 @@ impl DataLoader { while let Some(path) = data.pop_front() { match Self::try_read(&path) { Err(err) => { - println!("{} {:?} | {:?}", CROSS_MARK, path, err); + tracing::warn!("{} {:?} | {:?}", CROSS_MARK, path, err); continue; } Ok(img) => { @@ -211,7 +262,7 @@ impl DataLoader { // Deal with remaining data if !yis.is_empty() && sender.send((yis, yps)).is_err() { - println!("Receiver dropped, stopping production"); + tracing::info!("Receiver dropped, stopping production"); } } @@ -220,6 +271,11 @@ impl DataLoader { self } + pub fn with_progress_bar(mut self, x: bool) -> Self { + self.with_pb = x; + self + } + pub fn load_from_folder>(path: P) -> Result> { let mut paths: Vec = std::fs::read_dir(path)? .filter_map(|entry| entry.ok()) @@ -299,17 +355,7 @@ impl DataLoader { let saveout = Dir::Currnet .raw_path_with_subs(subs)? .join(format!("{}.mp4", string_now("-"))); - - // pb - let pb = ProgressBar::new(paths.len() as u64); - pb.set_style( - ProgressStyle::with_template( - "{prefix:.cyan.bold} {msg} |{bar}| ({percent_precise}%, {human_pos}/{human_len}, {per_sec})", - )? - .progress_chars("ā–ˆā–ˆ "), - ); - pb.set_prefix(" Converting"); - pb.set_message(saveout.to_str().unwrap_or_default().to_string()); + let pb = build_progress_bar(paths.len() as u64, " Converting", saveout.to_str(),"{prefix:.cyan.bold} {msg} |{bar}| ({percent_precise}%, {human_pos}/{human_len}, {per_sec})")?; // loop for path in paths { diff --git a/src/core/hub.rs b/src/core/hub.rs index 6f14835..f7b49e5 100644 --- a/src/core/hub.rs +++ b/src/core/hub.rs @@ -286,21 +286,24 @@ impl Hub { } pub fn connect_remote(&mut self) -> Result> { + let span = tracing::span!(tracing::Level::INFO, "OrtEngine-run"); + let _guard = span.enter(); + let should_download = if !self.cache.exists() { - // println!("No cache found, fetching data from GitHub"); + tracing::info!("No cache found, fetching data from GitHub"); true } else { match std::fs::metadata(&self.cache)?.modified() { Err(_) => { - // println!("Cannot get file modified time, fetching new data from GitHub"); + tracing::info!("Cannot get file modified time, fetching new data from GitHub"); true } Ok(modified_time) => { if std::time::SystemTime::now().duration_since(modified_time)? < self.ttl { - // println!("Using cached data"); + tracing::info!("Using cached data"); false } else { - // println!("Cache expired, fetching new data from GitHub"); + tracing::info!("Cache expired, fetching new data from GitHub"); true } } diff --git a/src/core/ort_engine.rs b/src/core/ort_engine.rs index d76cadf..46ee553 100644 --- a/src/core/ort_engine.rs +++ b/src/core/ort_engine.rs @@ -8,7 +8,8 @@ use prost::Message; use std::collections::HashSet; use crate::{ - human_bytes, onnx, Device, Dir, MinOptMax, Ops, Options, Ts, Xs, CHECK_MARK, CROSS_MARK, X, + build_progress_bar, human_bytes, onnx, Device, Dir, MinOptMax, Ops, Options, Ts, Xs, + CHECK_MARK, CROSS_MARK, X, }; /// Ort Tensor Attrs: name, data_type, dims @@ -37,6 +38,9 @@ pub struct OrtEngine { impl OrtEngine { pub fn new(config: &Options) -> Result { + let span = tracing::span!(tracing::Level::INFO, "OrtEngine-new"); + let _guard = span.enter(); + // onnx graph let model_proto = Self::load_onnx(&config.onnx_path)?; let graph = match &model_proto.graph { @@ -150,13 +154,13 @@ impl OrtEngine { } Device::Cuda(device_id) => { Self::build_cuda(&builder, device_id).unwrap_or_else(|err| { + tracing::warn!("{err}, Using cpu"); device = Device::Cpu(0); - println!("{err}"); }) } Device::CoreML(_) => Self::build_coreml(&builder).unwrap_or_else(|err| { + tracing::warn!("{err}, Using cpu"); device = Device::Cpu(0); - println!("{err}"); }), Device::Cpu(_) => { Self::build_cpu(&builder)?; @@ -169,7 +173,7 @@ impl OrtEngine { .commit_from_file(&config.onnx_path)?; // summary - println!( + tracing::info!( "{CHECK_MARK} Backend: ONNXRuntime | Opset: {} | Device: {:?} | Params: {}", model_proto.opset_import[0].version, device, @@ -200,6 +204,9 @@ impl OrtEngine { fp16_enable: bool, engine_cache_enable: bool, ) -> Result<()> { + let span = tracing::span!(tracing::Level::INFO, "OrtEngine-new"); + let _guard = span.enter(); + // auto generate shapes let mut spec_min = String::new(); let mut spec_opt = String::new(); @@ -240,7 +247,7 @@ impl OrtEngine { .with_profile_opt_shapes(spec_opt) .with_profile_max_shapes(spec_max); if trt.is_available()? && trt.register(builder).is_ok() { - println!("\nšŸ¢ Initial model serialization with TensorRT may require a wait...\n"); + tracing::info!("šŸ¢ Initial model serialization with TensorRT may require a wait...\n"); Ok(()) } else { anyhow::bail!("{CROSS_MARK} TensorRT initialization failed") @@ -276,6 +283,15 @@ impl OrtEngine { pub fn dry_run(&mut self) -> Result<()> { if self.num_dry_run > 0 { + // pb + let pb = build_progress_bar( + self.num_dry_run as u64, + " DryRun", + None, + "{prefix:.cyan.bold} {msg:.dim} {human_pos}/{human_len} |{bar}| {elapsed_precise}", + )?; + + // dummy inputs let mut xs = Vec::new(); for i in self.inputs_minoptmax.iter() { let mut x: Vec = Vec::new(); @@ -286,16 +302,23 @@ impl OrtEngine { xs.push(X::from(x)); } let xs = Xs::from(xs); + + // run for _ in 0..self.num_dry_run { + pb.inc(1); self.run(xs.clone())?; } self.ts.clear(); - println!("{CHECK_MARK} Dryrun x{}", self.num_dry_run); + + pb.finish(); } Ok(()) } pub fn run(&mut self, xs: Xs) -> Result { + let span = tracing::span!(tracing::Level::INFO, "OrtEngine-run"); + let _guard = span.enter(); + // inputs dtype alignment let mut xs_ = Vec::new(); let t_pre = std::time::Instant::now(); @@ -366,7 +389,7 @@ impl OrtEngine { if self.profile { let len = 10usize; let n = 4usize; - println!( + tracing::info!( "[Profile] {:>len$.n$?} ({:>len$.n$?} avg) [alignment: {:>len$.n$?} ({:>len$.n$?} avg) | inference: {:>len$.n$?} ({:>len$.n$?} avg) | to_f32: {:>len$.n$?} ({:>len$.n$?} avg)]", t_pre + t_run + t_post, self.ts.avg(), diff --git a/src/core/vision.rs b/src/core/vision.rs index f7fe6ed..a4c8931 100644 --- a/src/core/vision.rs +++ b/src/core/vision.rs @@ -25,6 +25,9 @@ pub trait Vision: Sized { /// Executes the full pipeline. fn forward(&mut self, xs: &[Self::Input], profile: bool) -> anyhow::Result> { + let span = tracing::span!(tracing::Level::INFO, "DataLoader-new"); + let _guard = span.enter(); + let t_pre = std::time::Instant::now(); let ys = self.preprocess(xs)?; let t_pre = t_pre.elapsed(); @@ -38,7 +41,9 @@ pub trait Vision: Sized { let t_post = t_post.elapsed(); if profile { - println!("> Preprocess: {t_pre:?} | Execution: {t_exe:?} | Postprocess: {t_post:?}"); + tracing::info!( + "> Preprocess: {t_pre:?} | Execution: {t_exe:?} | Postprocess: {t_post:?}" + ); } Ok(ys) diff --git a/src/models/yolo.rs b/src/models/yolo.rs index e5da132..7e7c9e6 100644 --- a/src/models/yolo.rs +++ b/src/models/yolo.rs @@ -32,6 +32,9 @@ impl Vision for YOLO { type Input = DynamicImage; fn new(options: Options) -> Result { + let span = tracing::span!(tracing::Level::INFO, "YOLO-new"); + let _guard = span.enter(); + let mut engine = OrtEngine::new(&options)?; let (batch, height, width) = ( engine.batch().to_owned(), @@ -49,7 +52,7 @@ impl Vision for YOLO { "segment" => Some(YOLOTask::Segment), "obb" => Some(YOLOTask::Obb), s => { - println!("YOLO Task: {s:?} is unsupported"); + tracing::error!("YOLO Task: {s:?} is unsupported"); None } })); @@ -135,7 +138,7 @@ impl Vision for YOLO { let iou = options.iou.unwrap_or(0.45); // Summary - println!("YOLO Task: {:?}, Version: {:?}", task, version); + tracing::info!("YOLO Task: {:?}, Version: {:?}", task, version); engine.dry_run()?; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 41de922..19479d4 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,5 +1,6 @@ #![allow(dead_code)] +use indicatif::{ProgressBar, ProgressStyle}; use rand::{distributions::Alphanumeric, thread_rng, Rng}; pub mod colormap256; @@ -55,3 +56,17 @@ pub(crate) fn string_now(delimiter: &str) -> String { ); t_now.format(&fmt).to_string() } + +pub fn build_progress_bar( + n: u64, + prefix: &str, + msg: Option<&str>, + style_temp: &str, +) -> anyhow::Result { + let pb = ProgressBar::new(n); + pb.set_style(ProgressStyle::with_template(style_temp)?.progress_chars("ā–ˆā–ˆ ")); + pb.set_prefix(prefix.to_string()); + pb.set_message(msg.unwrap_or_default().to_string()); + + Ok(pb) +}