Add florence2 all tasks

jamjamjon · Sep 21, 2024 · a06b0d8 · a06b0d8
1 parent 329b692
commit a06b0d8
Show file tree

Hide file tree

Showing 13 changed files with 385 additions and 246 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "usls"
-version = "0.0.14"
+version = "0.0.15"
 edition = "2021"
 description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models."
 repository = "https://github.com/jamjamjon/usls"
@@ -22,7 +22,6 @@ dirs = { version = "5.0.1" }
 ureq = { version = "2.9.1", default-features = true, features = [
     "socks-proxy",
 ] }
-walkdir = { version = "2.5.0" }  # TODO: remove
 tokenizers = { version = "0.15.2" }
 rayon = "1.10.0"
 indicatif = "0.17.8"

diff --git a/examples/dataloader/main.rs b/examples/dataloader/main.rs
@@ -18,26 +18,24 @@ fn main() -> anyhow::Result<()> {
 
     // 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:[email protected]:554/h265/ch1/",  // rtsp h264 stream
         "./assets/bus.jpg", // local image
-                            // "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
     )?
     .with_batch(1)
-    .with_progress_bar(true)
-    .with_bound(100)
     .build()?;
 
-    // // build annotator
+    // build annotator
     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(1000));
+        // std::thread::sleep(std::time::Duration::from_millis(100));
         let ys = model.forward(&xs, false)?;
         annotator.annotate(&xs, &ys);
     }

diff --git a/examples/florence2/main.rs b/examples/florence2/main.rs
@@ -134,62 +134,117 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         options_decoder,
         options_decoder_merged,
     )?;
-    // .with_task(Task::Caption(2));
 
     // load images
     let xs = [
+        // DataLoader::try_read("florence2/car.jpg")?, // for testing region-related tasks
         DataLoader::try_read("florence2/car.jpg")?,
+        // DataLoader::try_read("images/db.png")?,
         DataLoader::try_read("assets/bus.jpg")?,
     ];
 
-    // run with a batch of tasks
-    let ys = model.run(
+    // region-related tasks
+    let quantizer = usls::Quantizer::default();
+    // let coords = [449., 270., 556., 372.];  // wheel
+    let coords = [31., 156., 581., 373.]; // car
+    let (width_car, height_car) = (xs[0].width(), xs[0].height());
+    let quantized_coords = quantizer.quantize(&coords, (width_car as _, height_car as _));
+
+    // run with tasks
+    let ys = model.run_with_tasks(
         &xs,
         &[
             // w/ inputs
             Task::Caption(0),
             Task::Caption(1),
             Task::Caption(2),
             Task::Ocr,
+            Task::OcrWithRegion,
             Task::RegionProposal,
             Task::ObjectDetection,
             Task::DenseRegionCaption,
-            // // Task::OcrWithRegion, // TODO
             // w/o inputs
-            // Task::OpenSetDetection("A green car".into()),
-            // Task::CaptionToPhraseGrounding("A green car".into()),
-            // Task::ReferringExpressionSegmentation("A green car".into()),
-            // Task::RegionToSegmentation(702, 575, 866, 772),
-            // Task::RegionToCategory(52, 332, 932, 774),
-            // Task::RegionToDescription(52, 332, 932, 774),
-            // Task::RegionToOcr(100, 100, 300, 300),
+            Task::OpenSetDetection("a vehicle".into()),
+            Task::CaptionToPhraseGrounding(
+                "A vehicle with two wheels parked in front of a building.".into(),
+            ),
+            Task::ReferringExpressionSegmentation("a vehicle".into()),
+            Task::RegionToSegmentation(
+                quantized_coords[0],
+                quantized_coords[1],
+                quantized_coords[2],
+                quantized_coords[3],
+            ),
+            Task::RegionToCategory(
+                quantized_coords[0],
+                quantized_coords[1],
+                quantized_coords[2],
+                quantized_coords[3],
+            ),
+            Task::RegionToDescription(
+                quantized_coords[0],
+                quantized_coords[1],
+                quantized_coords[2],
+                quantized_coords[3],
+            ),
         ],
     )?;
 
     // annotator
+    let annotator = Annotator::new()
+        .without_bboxes_conf(true)
+        .with_bboxes_thickness(3)
+        .with_saveout_subs(&["Florence2"]);
     for (task, ys_) in ys.iter() {
         match task {
-            Task::Caption(_) | Task::Ocr => println!("Task: {:?}\n{:?}\n", task, ys_),
+            Task::Caption(_)
+            | Task::Ocr
+            | Task::RegionToCategory(..)
+            | Task::RegionToDescription(..) => {
+                println!("Task: {:?}\n{:?}\n", task, ys_)
+            }
             Task::DenseRegionCaption => {
-                let annotator = Annotator::default()
-                    .without_bboxes_conf(true)
-                    .with_bboxes_thickness(4)
-                    .with_saveout("Florence2-DenseRegionCaption");
+                let annotator = annotator.clone().with_saveout("Dense-Region-Caption");
                 annotator.annotate(&xs, ys_);
             }
             Task::RegionProposal => {
-                let annotator = Annotator::default()
-                    .without_bboxes_conf(true)
-                    .without_bboxes_name(true)
-                    .with_bboxes_thickness(4)
-                    .with_saveout("Florence2-RegionProposal");
+                let annotator = annotator
+                    .clone()
+                    .without_bboxes_name(false)
+                    .with_saveout("Region-Proposal");
+
                 annotator.annotate(&xs, ys_);
             }
             Task::ObjectDetection => {
-                let annotator = Annotator::default()
-                    .without_bboxes_conf(true)
-                    .with_bboxes_thickness(4)
-                    .with_saveout("Florence2-ObjectDetection");
+                let annotator = annotator.clone().with_saveout("Object-Detection");
+                annotator.annotate(&xs, ys_);
+            }
+            Task::OpenSetDetection(_) => {
+                let annotator = annotator.clone().with_saveout("Open-Set-Detection");
+                annotator.annotate(&xs, ys_);
+            }
+            Task::CaptionToPhraseGrounding(_) => {
+                let annotator = annotator
+                    .clone()
+                    .with_saveout("Caption-To-Phrase-Grounding");
+                annotator.annotate(&xs, ys_);
+            }
+            Task::ReferringExpressionSegmentation(_) => {
+                let annotator = annotator
+                    .clone()
+                    .with_polygons_alpha(200)
+                    .with_saveout("Referring-Expression-Segmentation");
+                annotator.annotate(&xs, ys_);
+            }
+            Task::RegionToSegmentation(..) => {
+                let annotator = annotator
+                    .clone()
+                    .with_polygons_alpha(200)
+                    .with_saveout("Region-To-Segmentation");
+                annotator.annotate(&xs, ys_);
+            }
+            Task::OcrWithRegion => {
+                let annotator = annotator.clone().with_saveout("Ocr-With-Region");
                 annotator.annotate(&xs, ys_);
             }
 

diff --git a/src/core/annotator.rs b/src/core/annotator.rs
@@ -2,19 +2,21 @@ use crate::{
     colormap256, string_now, Bbox, Dir, Hub, Keypoint, Mask, Mbr, Polygon, Prob, CHECK_MARK,
     CROSS_MARK, Y,
 };
-use ab_glyph::{FontVec, PxScale};
+use ab_glyph::{FontArc, PxScale};
 use anyhow::Result;
 use image::{DynamicImage, GenericImage, Rgba, RgbaImage};
 use imageproc::map::map_colors;
 
 /// Annotator for struct `Y`
-// #[derive(Debug)]
+#[derive(Clone)]
 pub struct Annotator {
-    font: FontVec,
+    // TODO: Add lifetime
+    font: FontArc,
     _scale: f32, // Cope with ab_glyph & imageproc=0.24.0
     scale_dy: f32,
     saveout_base: String,
     saveout: Option<String>,
+    saveout_subs: Vec<String>,
     decimal_places: usize,
 
     // About mbrs
@@ -72,6 +74,7 @@ impl Default for Annotator {
             scale_dy: 28.,
             polygons_alpha: 179,
             saveout: None,
+            saveout_subs: vec![],
             saveout_base: String::from("runs"),
             decimal_places: 4,
             without_bboxes: false,
@@ -323,18 +326,35 @@ impl Annotator {
         self
     }
 
+    pub fn with_saveout_subs(mut self, xs: &[&str]) -> Self {
+        self.saveout_subs = xs.iter().map(|x| x.to_string()).collect::<Vec<String>>();
+        self
+    }
+
     pub fn with_font(mut self, path: &str) -> Result<Self> {
         self.font = Self::load_font(Some(path))?;
         Ok(self)
     }
 
     /// Create folders for saving annotated results. e.g., `./runs/xxx`
     pub fn saveout(&self) -> Result<std::path::PathBuf> {
-        let subs = match &self.saveout {
-            Some(x) => vec![self.saveout_base.as_str(), x.as_str()],
-            None => vec![self.saveout_base.as_str()],
-        };
+        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)
     }
 
@@ -345,7 +365,7 @@ 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 span = tracing::span!(tracing::Level::INFO, "Annotator-plot");
         let _guard = span.enter();
 
         let mut vs: Vec<RgbaImage> = Vec::new();
@@ -749,13 +769,13 @@ impl Annotator {
     }
 
     /// Load custom font
-    fn load_font(path: Option<&str>) -> Result<FontVec> {
+    fn load_font(path: Option<&str>) -> Result<FontArc> {
         let path_font = match path {
             None => Hub::new()?.fetch("fonts/Arial.ttf")?.commit()?,
             Some(p) => p.into(),
         };
         let buffer = std::fs::read(path_font)?;
-        Ok(FontVec::try_from_vec(buffer.to_owned())?)
+        Ok(FontArc::try_from_vec(buffer.to_owned())?)
     }
 
     /// Pick color from pallette
@@ -765,6 +785,7 @@ impl Annotator {
 
     /// Color pallette
     fn color_palette() -> [(u8, u8, u8, u8); 20] {
+        // TODO: more colors
         [
             (0, 255, 127, 255),   // spring green
             (255, 105, 180, 255), // hot pink

diff --git a/src/core/dataloader.rs b/src/core/dataloader.rs
@@ -35,7 +35,7 @@ impl Iterator for DataLoaderIterator {
                 None => {
                     progress_bar.set_prefix("    Iterated");
                     progress_bar.set_style(
-                        indicatif::ProgressStyle::with_template(crate::PROGRESS_BAR_STYLE_GREEN)
+                        indicatif::ProgressStyle::with_template(crate::PROGRESS_BAR_STYLE_FINISH_2)
                             .unwrap(),
                     );
                     progress_bar.finish();
@@ -56,7 +56,7 @@ impl IntoIterator for DataLoader {
                 self.nf / self.batch_size as u64,
                 "   Iterating",
                 Some(&format!("{:?}", self.media_type)),
-                crate::PROGRESS_BAR_STYLE_CYAN,
+                crate::PROGRESS_BAR_STYLE_CYAN_2,
             )
             .ok()
         } else {
@@ -365,7 +365,12 @@ impl DataLoader {
         let saveout = Dir::Currnet
             .raw_path_with_subs(subs)?
             .join(format!("{}.mp4", string_now("-")));
-        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})")?;
+        let pb = build_progress_bar(
+            paths.len() as u64,
+            "  Converting",
+            Some(&format!("{:?}", MediaType::Video(Location::Local))),
+            crate::PROGRESS_BAR_STYLE_CYAN_2,
+        )?;
 
         // loop
         for path in paths {
@@ -397,10 +402,10 @@ impl DataLoader {
         }
 
         // update
-        pb.set_prefix("  Downloaded");
         pb.set_prefix("   Converted");
+        pb.set_message(saveout.to_str().unwrap_or_default().to_string());
         pb.set_style(ProgressStyle::with_template(
-            "{prefix:.green.bold} {msg} in {elapsed}",
+            crate::PROGRESS_BAR_STYLE_FINISH_4,
         )?);
         pb.finish();