Skip to content

Commit

Permalink
Add tracing
Browse files Browse the repository at this point in the history
  • Loading branch information
@jamjamjon
jamjamjon committed Sep 14, 2024
1 parent 03f8bd5 commit af8e8a3
Show file tree
Hide file tree
Showing 9 changed files with 169 additions and 65 deletions.
2 changes: 2 additions & 0 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
16 changes: 12 additions & 4 deletions examples/dl/main.rs
View file
Original file line number Diff line number Diff line change
@@ -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")?
Expand All @@ -17,27 +21,31 @@ fn main() -> anyhow::Result<()> {
// "rtsp://admin:[email protected]:554/h265/ch1/",
// "rtsp://admin:[email protected]: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(())
}
11 changes: 5 additions & 6 deletions src/core/annotator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -347,6 +345,9 @@ impl Annotator {

/// Plot images and return plotted images(RGBA8)
pub fn plot(&self, imgs: &[DynamicImage], ys: &[Y]) -> Result<Vec<RgbaImage>> {
let span = tracing::span!(tracing::Level::INFO, "YOLO-new");
let _guard = span.enter();

let mut vs: Vec<RgbaImage> = Vec::new();

// annotate
Expand Down Expand Up @@ -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);
}
}

Expand Down
128 changes: 87 additions & 41 deletions src/core/dataloader.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<DynamicImage>, Vec<PathBuf>);

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<TempReturnType>,
progress_bar: Option<ProgressBar>,
}

impl Iterator for DataLoaderIterator {
type Item = TempReturnType;

fn next(&mut self) -> Option<Self::Item> {
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<VecDeque<PathBuf>>,
pub media_type: Option<MediaType>,
pub media_type: MediaType,
pub batch_size: usize,
sender: Option<mpsc::Sender<TempReturnType>>,
receiver: mpsc::Receiver<TempReturnType>,
pub decoder: Option<video_rs::decode::Decoder>,
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<Self> {
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);
}

Expand All @@ -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::<Url>()?.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,
Expand All @@ -122,14 +168,17 @@ impl DataLoader {
receiver,
batch_size: 1,
decoder,
nf,
with_pb: true,
})
}

pub fn build(mut self) -> Result<Self> {
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
Expand All @@ -147,6 +196,8 @@ impl DataLoader {
media_type: MediaType,
mut decoder: Option<video_rs::decode::Decoder>,
) {
let span = tracing::span!(tracing::Level::INFO, "DataLoader-producer-thread");
let _guard = span.enter();
let mut yis: Vec<DynamicImage> = Vec::with_capacity(batch_size);
let mut yps: Vec<PathBuf> = Vec::with_capacity(batch_size);

Expand All @@ -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) => {
Expand Down Expand Up @@ -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");
}
}

Expand All @@ -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<P: AsRef<std::path::Path>>(path: P) -> Result<Vec<std::path::PathBuf>> {
let mut paths: Vec<PathBuf> = std::fs::read_dir(path)?
.filter_map(|entry| entry.ok())
Expand Down Expand Up @@ -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 {
Expand Down
11 changes: 7 additions & 4 deletions src/core/hub.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -286,21 +286,24 @@ impl Hub {
}

pub fn connect_remote(&mut self) -> Result<Vec<Release>> {
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
}
}
Expand Down
Loading

0 comments on commit af8e8a3

Please sign in to comment.