Update Viewer

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "usls"
-version = "0.0.15"
+version = "0.0.16"
 edition = "2021"
 description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models."
 repository = "https://github.com/jamjamjon/usls"

examples/dataloader/main.rs

@@ -28,47 +28,39 @@ fn main() -> anyhow::Result<()> {
         // "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:[email protected]:554/h265/ch1/",  // rtsp h264 stream
-                                                                                             // "./assets/bus.jpg", // local image
-                                                                                             // "/home/qweasd/Desktop/SourceVideos/3.mp4",
-                                                                                             // "../hall.mp4",
-                                                                                             // "../000020489.jpg",
+        // "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // remote video
+        // "rtsp://admin:[email protected]:554/h265/ch1/",  // rtsp h264 stream
+        // "./assets/bus.jpg", // local image
+        "../SourceVideos/3.mp4",
     )?
     .with_batch(1)
     .build()?;
 
-    let mut viewer = Viewer::new();
+    let mut viewer = Viewer::new().with_delay(100).with_scale(1.).resizable(true);
 
-    // run
+    // iteration
     for (xs, _) in dl {
-        // std::thread::sleep(std::time::Duration::from_millis(100));
-        let ys = model.forward(&xs, false)?;
-        // annotator.annotate(&xs, &ys);
+        // inference & annotate
+        let ys = model.run(&xs)?;
+        let images_plotted = annotator.plot(&xs, &ys, false)?;
 
-        // TODO: option for saving
-        let images_plotted = annotator.plot(&xs, &ys)?;
+        // show image
+        viewer.imshow(&images_plotted)?;
 
-        // RgbaImage -> DynamicImage
-        let frames: Vec<image::DynamicImage> = images_plotted
-            .iter()
-            .map(|x| image::DynamicImage::from(x.to_owned()))
-            .collect();
+        // write video
+        viewer.write_batch(&images_plotted)?;
 
-        // imshow
-        viewer.imshow(&frames)?;
+        // check out window and key event
         if !viewer.is_open() || viewer.is_key_pressed(Key::Escape) {
             break;
         }
-
-        // write video
-        viewer.write_batch(&frames, 30)?; // fps
     }
 
+    // finish video write
+    viewer.finish_write()?;
+
     // images -> video
     // DataLoader::is2v("runs/YOLO-DataLoader", &["runs", "is2v"], 24)?;
 
-    viewer.finish_write()?;
-
     Ok(())
 }

src/core/annotator.rs

@@ -358,17 +358,17 @@ impl Annotator {
         Dir::Currnet.raw_path_with_subs(&subs)
     }
 
-    /// Annotate images, and no return
+    /// Annotate images, save, and no return
     pub fn annotate(&self, imgs: &[DynamicImage], ys: &[Y]) {
-        let _ = self.plot(imgs, ys);
+        let _ = self.plot(imgs, ys, true);
     }
 
-    /// Plot images and return plotted images(RGBA8)
-    pub fn plot(&self, imgs: &[DynamicImage], ys: &[Y]) -> Result<Vec<RgbaImage>> {
+    /// Plot images and return plotted images
+    pub fn plot(&self, imgs: &[DynamicImage], ys: &[Y], save: bool) -> Result<Vec<DynamicImage>> {
         let span = tracing::span!(tracing::Level::INFO, "Annotator-plot");
         let _guard = span.enter();
 
-        let mut vs: Vec<RgbaImage> = Vec::new();
+        let mut vs: Vec<DynamicImage> = Vec::new();
 
         // annotate
         for (img, y) in imgs.iter().zip(ys.iter()) {
@@ -414,16 +414,19 @@ impl Annotator {
                 self.plot_probs(&mut img_rgba, xs);
             }
 
-            // save
-            let saveout = self.saveout()?.join(format!("{}.png", string_now("-")));
-            match img_rgba.save(&saveout) {
-                Err(err) => tracing::error!("{} Saving failed: {:?}", CROSS_MARK, err),
-                Ok(_) => {
-                    tracing::info!("{} Annotated image saved to: {:?}", CHECK_MARK, saveout);
+            // save or not
+            if save {
+                let saveout = self.saveout()?.join(format!("{}.png", string_now("-")));
+                match img_rgba.save(&saveout) {
+                    Err(err) => tracing::error!("{} Saving failed: {:?}", CROSS_MARK, err),
+                    Ok(_) => {
+                        tracing::info!("{} Annotated image saved to: {:?}", CHECK_MARK, saveout);
+                    }
                 }
             }
 
-            vs.push(img_rgba);
+            // RgbaImage -> DynamicImage
+            vs.push(image::DynamicImage::from(img_rgba));
         }
         Ok(vs)
     }

src/core/dir.rs

@@ -14,6 +14,10 @@ pub enum Dir {
 }
 
 impl Dir {
+    pub fn saveout(subs: &[&str]) -> anyhow::Result<std::path::PathBuf> {
+        Self::Currnet.raw_path_with_subs(subs)
+    }
+
     /// Retrieves the base path for the specified directory type, optionally appending the `usls` subdirectory.
     ///
     /// # Arguments

src/core/viewer.rs

@@ -11,25 +11,25 @@ use crate::{string_now, Dir, Key};
 pub struct Viewer<'a> {
     name: &'a str,
     window: Option<Window>,
+    window_scale: f32,
+    window_resizable: bool,
     fps_poll: usize,
+    fps: usize,
     writer: Option<Encoder>,
     position: Time,
-    saveout: Option<String>,
-    saveout_base: String,
-    saveout_subs: Vec<String>,
 }
 
 impl Default for Viewer<'_> {
     fn default() -> Self {
         Self {
             name: "usls-viewer",
             window: None,
+            window_scale: 0.5,
+            window_resizable: true,
             fps_poll: 100,
+            fps: 25,
             writer: None,
             position: Time::zero(),
-            saveout: None,
-            saveout_subs: vec![],
-            saveout_base: String::from("runs"),
         }
     }
 }
@@ -39,73 +39,38 @@ impl Viewer<'_> {
         Default::default()
     }
 
-    /// Create folders for saving annotated results. e.g., `./runs/xxx`
-    pub fn saveout(&self) -> Result<std::path::PathBuf> {
-        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::<Vec<&str>>();
-                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);
+            let (w_scale, h_scale) = (
+                (w as f32 * self.window_scale) as usize,
+                (h as f32 * self.window_scale) as usize,
+            );
 
-            // check if need to reload
-            let reload_window = match &self.window {
-                Some(window) => window.get_size() != (w, h),
+            // should reload?
+            let should_reload = match &self.window {
                 None => true,
+                Some(window) => {
+                    if self.window_resizable {
+                        false
+                    } else {
+                        window.get_size() != (w_scale, h_scale)
+                    }
+                }
             };
 
-            // reload
-            if reload_window {
+            // create window
+            if should_reload {
                 self.window = Window::new(
                     self.name,
-                    w as _,
-                    h as _,
+                    w_scale,
+                    h_scale,
                     WindowOptions {
                         resize: true,
                         topmost: true,
                         borderless: false,
+                        scale: minifb::Scale::X1,
                         ..WindowOptions::default()
                     },
                 )
@@ -116,7 +81,7 @@ impl Viewer<'_> {
                 });
             }
 
-            // update buffer
+            // build buffer
             let mut buffer: Vec<u32> = Vec::with_capacity(w * h);
             for pixel in rgb.pixels() {
                 let r = pixel[0];
@@ -126,56 +91,101 @@ impl Viewer<'_> {
                 buffer.push(p);
             }
 
-            // Update the window with the image buffer
+            // update 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<()> {
+    pub fn write(&mut self, frame: &image::DynamicImage) -> 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("-")));
+            let saveout = Dir::saveout(&["runs"])?.join(format!("{}.mp4", string_now("-")));
+            tracing::info!("Video will be save to: {:?}", saveout);
             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))
+                .aligned_with(Time::from_nth_of_a_second(self.fps))
                 .add();
         }
         Ok(())
     }
 
-    pub fn write_batch(&mut self, frames: &[image::DynamicImage], fps: usize) -> Result<()> {
+    pub fn write_batch(&mut self, frames: &[image::DynamicImage]) -> Result<()> {
         for frame in frames.iter() {
-            self.write(frame, fps)?
+            self.write(frame)?
         }
         Ok(())
     }
 
     pub fn finish_write(&mut self) -> Result<()> {
         match &mut self.writer {
-            Some(writer) => Ok(writer.finish()?),
-            None => anyhow::bail!("Found no video encoder."),
+            Some(writer) => writer.finish()?,
+            None => {
+                tracing::info!("Found no video writer. No need to release.");
+            }
         }
+        Ok(())
+    }
+
+    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 resizable(mut self, x: bool) -> Self {
+        self.window_resizable = x;
+        self
+    }
+
+    pub fn with_scale(mut self, x: f32) -> Self {
+        self.window_scale = x;
+        self
+    }
+
+    pub fn with_fps(mut self, x: usize) -> Self {
+        self.fps = x;
+        self
+    }
+
+    pub fn with_delay(mut self, x: usize) -> Self {
+        self.fps_poll = 1000 / x; // 1000 / 60 = 16.7
+        self
+    }
+
+    pub fn wh(&self) -> Option<(usize, usize)> {
+        self.window.as_ref().map(|x| x.get_size())
     }
 
     fn from_u8_rgb(r: u8, g: u8, b: u8) -> u32 {