From 8cb853d0dce9dc82bd9832c22a958c6438ecb1f9 Mon Sep 17 00:00:00 2001 From: jamjamjon Date: Tue, 24 Sep 2024 22:59:20 +0800 Subject: [PATCH] Add Viewer for imshow --- Cargo.toml | 1 + examples/dataloader/main.rs | 45 +++++++-- src/core/mod.rs | 5 + src/core/viewer.rs | 185 ++++++++++++++++++++++++++++++++++++ 4 files changed, 227 insertions(+), 9 deletions(-) create mode 100644 src/core/viewer.rs diff --git a/Cargo.toml b/Cargo.toml index 5beb404..50a3e55 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ video-rs = { version = "0.9.0", features = ["ndarray"] } natord = "1.0.9" tracing = "0.1.40" tracing-subscriber = "0.3.18" +minifb = "0.27.0" [features] diff --git a/examples/dataloader/main.rs b/examples/dataloader/main.rs index 40484f2..7ed7eab 100644 --- a/examples/dataloader/main.rs +++ b/examples/dataloader/main.rs @@ -1,4 +1,6 @@ -use usls::{models::YOLO, Annotator, DataLoader, Options, Vision, YOLOTask, YOLOVersion}; +use usls::{ + models::YOLO, Annotator, DataLoader, Key, Options, Viewer, Vision, YOLOTask, YOLOVersion, +}; fn main() -> anyhow::Result<()> { tracing_subscriber::fmt() @@ -16,32 +18,57 @@ fn main() -> anyhow::Result<()> { .with_confs(&[0.2]); let mut model = YOLO::new(options)?; + // build annotator + let annotator = Annotator::new() + .with_bboxes_thickness(4) + .with_saveout("YOLO-DataLoader"); + // build dataloader let dl = DataLoader::new( // "images/bus.jpg", // remote image // "../images", // image folder // "../demo.mp4", // local video - // "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // remote video - // "rtsp://admin:xyz@192.168.2.217:554/h265/ch1/", // rtsp h264 stream - "./assets/bus.jpg", // local image + "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // remote video + // "rtsp://admin:xyz@192.168.2.217:554/h265/ch1/", // rtsp h264 stream + // "./assets/bus.jpg", // local image + // "/home/qweasd/Desktop/SourceVideos/3.mp4", + // "../hall.mp4", + // "../000020489.jpg", )? .with_batch(1) .build()?; - // build annotator - let annotator = Annotator::new() - .with_bboxes_thickness(4) - .with_saveout("YOLO-DataLoader"); + let mut viewer = Viewer::new(); // run for (xs, _) in dl { // std::thread::sleep(std::time::Duration::from_millis(100)); let ys = model.forward(&xs, false)?; - annotator.annotate(&xs, &ys); + // annotator.annotate(&xs, &ys); + + // TODO: option for saving + let images_plotted = annotator.plot(&xs, &ys)?; + + // RgbaImage -> DynamicImage + let frames: Vec = images_plotted + .iter() + .map(|x| image::DynamicImage::from(x.to_owned())) + .collect(); + + // imshow + viewer.imshow(&frames)?; + if !viewer.is_open() || viewer.is_key_pressed(Key::Escape) { + break; + } + + // write video + viewer.write_batch(&frames, 30)?; // fps } // images -> video // DataLoader::is2v("runs/YOLO-DataLoader", &["runs", "is2v"], 24)?; + viewer.finish_write()?; + Ok(()) } diff --git a/src/core/mod.rs b/src/core/mod.rs index 7225f8f..ab5d8d6 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -15,6 +15,7 @@ mod ort_engine; mod task; mod tokenizer_stream; mod ts; +mod viewer; mod vision; mod x; mod xs; @@ -35,6 +36,10 @@ pub use ort_engine::OrtEngine; pub use task::Task; pub use tokenizer_stream::TokenizerStream; pub use ts::Ts; +pub use viewer::Viewer; pub use vision::Vision; pub use x::X; pub use xs::Xs; + +// re-export +pub use minifb::Key; diff --git a/src/core/viewer.rs b/src/core/viewer.rs new file mode 100644 index 0000000..d61ee4d --- /dev/null +++ b/src/core/viewer.rs @@ -0,0 +1,185 @@ +use anyhow::Result; +use image::DynamicImage; +use minifb::{Window, WindowOptions}; +use video_rs::{ + encode::{Encoder, Settings}, + time::Time, +}; + +use crate::{string_now, Dir, Key}; + +pub struct Viewer<'a> { + name: &'a str, + window: Option, + fps_poll: usize, + writer: Option, + position: Time, + saveout: Option, + saveout_base: String, + saveout_subs: Vec, +} + +impl Default for Viewer<'_> { + fn default() -> Self { + Self { + name: "usls-viewer", + window: None, + fps_poll: 100, + writer: None, + position: Time::zero(), + saveout: None, + saveout_subs: vec![], + saveout_base: String::from("runs"), + } + } +} + +impl Viewer<'_> { + pub fn new() -> Self { + Default::default() + } + + /// Create folders for saving annotated results. e.g., `./runs/xxx` + pub fn saveout(&self) -> Result { + let mut subs = vec![self.saveout_base.as_str()]; + if let Some(saveout) = &self.saveout { + // add subs + if !self.saveout_subs.is_empty() { + let xs = self + .saveout_subs + .iter() + .map(|x| x.as_str()) + .collect::>(); + subs.extend(xs); + } + + // add filename + subs.push(saveout); + } + + // mkdir even no filename specified + Dir::Currnet.raw_path_with_subs(&subs) + } + + pub fn is_open(&self) -> bool { + if let Some(window) = &self.window { + window.is_open() + } else { + false + } + } + + pub fn is_key_pressed(&self, key: Key) -> bool { + if let Some(window) = &self.window { + window.is_key_down(key) + } else { + false + } + } + + pub fn is_esc_pressed(&self) -> bool { + self.is_key_pressed(Key::Escape) + } + + pub fn wh(&self) -> Option<(usize, usize)> { + self.window.as_ref().map(|x| x.get_size()) + } + + pub fn imshow(&mut self, xs: &[DynamicImage]) -> Result<()> { + for x in xs.iter() { + let rgb = x.to_rgb8(); + let (w, h) = (rgb.width() as usize, rgb.height() as usize); + + // check if need to reload + let reload_window = match &self.window { + Some(window) => window.get_size() != (w, h), + None => true, + }; + + // reload + if reload_window { + self.window = Window::new( + self.name, + w as _, + h as _, + WindowOptions { + resize: true, + topmost: true, + borderless: false, + ..WindowOptions::default() + }, + ) + .ok() + .map(|mut x| { + x.set_target_fps(self.fps_poll); + x + }); + } + + // update buffer + let mut buffer: Vec = Vec::with_capacity(w * h); + for pixel in rgb.pixels() { + let r = pixel[0]; + let g = pixel[1]; + let b = pixel[2]; + let p = Self::from_u8_rgb(r, g, b); + buffer.push(p); + } + + // Update the window with the image buffer + self.window + .as_mut() + .unwrap() + .update_with_buffer(&buffer, w, h)?; + + // optional: save as videos + } + + Ok(()) + } + + pub fn write(&mut self, frame: &image::DynamicImage, fps: usize) -> Result<()> { + // build writer at the 1st time + let frame = frame.to_rgb8(); + let (w, h) = frame.dimensions(); + if self.writer.is_none() { + let settings = Settings::preset_h264_yuv420p(w as _, h as _, false); + let saveout = self.saveout()?.join(format!("{}.mp4", string_now("-"))); + self.writer = Some(Encoder::new(saveout, settings)?); + } + + // write video + if let Some(writer) = self.writer.as_mut() { + // let raw_data = frame.into_raw(); + let raw_data = frame.to_vec(); + let frame = ndarray::Array3::from_shape_vec((h as usize, w as usize, 3), raw_data)?; + + // encode and update + writer.encode(&frame, self.position)?; + self.position = self + .position + .aligned_with(Time::from_nth_of_a_second(fps)) + .add(); + } + Ok(()) + } + + pub fn write_batch(&mut self, frames: &[image::DynamicImage], fps: usize) -> Result<()> { + for frame in frames.iter() { + self.write(frame, fps)? + } + Ok(()) + } + + pub fn finish_write(&mut self) -> Result<()> { + match &mut self.writer { + Some(writer) => Ok(writer.finish()?), + None => anyhow::bail!("Found no video encoder."), + } + } + + fn from_u8_rgb(r: u8, g: u8, b: u8) -> u32 { + let (r, g, b) = (r as u32, g as u32, b as u32); + (r << 16) | (g << 8) | b + } +}