feat(yoloe): add visual prompt support and fix font loading bug (#152)

Author: jamjamjon

Cargo.toml

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

README.md

@@ -93,7 +93,7 @@ usls = "latest-version"
 | [SAM-HQ](https://github.com/SysCV/sam-hq) | Segment Anything | [demo](examples/sam) |
 | [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM) | Instance Segmentation | [demo](examples/yolo) |
 | [YOLO-World](https://github.com/AILab-CVC/YOLO-World) | Open-Set Detection With Language | [demo](examples/yolo) |
-| [YOLOE](https://github.com/THU-MIG/yoloe) | Open-Set Detection And Segmentation | [demo-prompt-free](examples/yoloe)<br />[demo-text-prompt](examples/yoloe-text-prompt) |
+| [YOLOE](https://github.com/THU-MIG/yoloe) | Open-Set Detection And Segmentation | [demo-prompt-free](examples/yoloe-prompt-free)<br />[demo-prompt(visual & textual)](examples/yoloe-prompt) |
 | [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO) | Open-Set Detection With Language | [demo](examples/grounding-dino) |
 | [CLIP](https://github.com/openai/CLIP) | Vision-Language Embedding | [demo](examples/clip) |
 | [jina-clip-v1](https://huggingface.co/jinaai/jina-clip-v1) | Vision-Language Embedding | [demo](examples/clip) |

examples/yoloe-prompt-free/README.md

@@ -0,0 +1,6 @@
+## Quick Start
+
+```shell
+cargo run -r -F cuda --example yoloe-prompt-free -- --device cuda
+```
+

examples/yoloe/main.rs examples/yoloe-prompt-free/main.rsexamples/yoloe/main.rs renamed to examples/yoloe-prompt-free/main.rs

@@ -49,7 +49,7 @@ fn main() -> Result<()> {
         annotator.annotate(x, y)?.save(format!(
             "{}.jpg",
             usls::Dir::Current
-                .base_dir_with_subs(&["runs", model.spec()])?
+                .base_dir_with_subs(&["runs", "YOLOE-prompt-free", model.spec()])?
                 .join(usls::timestamp(None))
                 .display(),
         ))?;

examples/yoloe-prompt/README.md

@@ -0,0 +1,10 @@
+## Quick Start
+
+```shell
+# Textual prompt (CPU)
+cargo run -r --example yoloe-prompt -- --source ./assets/bus.jpg --visual false
+
+# Visual prompt (TensorRT)
+cargo run -r --example yoloe-prompt -F tensorrt -- --source ./assets/bus.jpg --visual true --device tensorrt
+```
+

examples/yoloe-text-prompt/main.rs examples/yoloe-prompt/main.rsexamples/yoloe-text-prompt/main.rs renamed to examples/yoloe-prompt/main.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use usls::{models::YOLO, Annotator, Config, DataLoader, Style};
+use usls::{models::YOLO, Annotator, Config, DataLoader, Hbb, Style};
 
 #[derive(argh::FromArgs)]
 /// Example
@@ -9,7 +9,7 @@ struct Args {
     source: String,
 
     /// dtype
-    #[argh(option, default = "String::from(\"fp16\")")]
+    #[argh(option, default = "String::from(\"fp32\")")]
     dtype: String,
 
     /// device
@@ -33,6 +33,10 @@ struct Args {
     /// batch size
     #[argh(option, default = "1")]
     batch_size: usize,
+
+    /// visual or textual
+    #[argh(option, default = "true")]
+    visual: bool,
 }
 
 fn main() -> Result<()> {
@@ -43,14 +47,32 @@ fn main() -> Result<()> {
     let args: Args = argh::from_env();
 
     // config
-    let config = Config::yoloe_v8m_seg_tp()
-        .with_batch_size_all_min_opt_max(1, args.batch_size, 8)
-        .with_model_dtype(args.dtype.as_str().parse()?)
-        .with_textual_dtype("fp16".parse()?) // Use FP32 when TensorRT is enabled
-        .with_device_all(args.device.as_str().parse()?)
-        .commit()?;
+    let config = if args.visual {
+        Config::yoloe_11m_seg_vp()
+    } else {
+        Config::yoloe_v8m_seg_tp().with_textual_encoder_dtype("fp16".parse()?) // Use FP32 when TensorRT is enabled
+    }
+    .with_batch_size_all_min_opt_max(1, args.batch_size, 8)
+    .with_model_dtype(args.dtype.as_str().parse()?)
+    .with_device_all(args.device.as_str().parse()?)
+    .with_class_confs(&[0.25])
+    .commit()?;
     let mut model = YOLO::new(config)?;
 
+    // encode visual or textual
+    let embedding = if args.visual {
+        let prompt_image = DataLoader::try_read_one("./assets/bus.jpg")?;
+        model.encode_visual_prompt(
+            prompt_image,
+            &[
+                Hbb::from_xyxy(221.52, 405.8, 344.98, 857.54).with_name("person"),
+                // Hbb::from_xyxy(120., 425., 160., 445.).with_name("glasses"), // TODO
+            ],
+        )?
+    } else {
+        model.encode_class_names(&args.labels.iter().map(|x| x.as_str()).collect::<Vec<_>>())?
+    };
+
     // build dataloader
     let dl = DataLoader::new(&args.source)?
         .with_batch(model.batch() as _)
@@ -61,14 +83,9 @@ fn main() -> Result<()> {
         .with_hbb_style(Style::hbb().with_draw_fill(true))
         .with_mask_style(Style::mask().with_draw_mask_polygon_largest(true));
 
-    // encode text prompts
-    let text_embeddings =
-        model.encode_class_names(&args.labels.iter().map(|x| x.as_str()).collect::<Vec<_>>())?;
-
     // run & annotate
     for xs in &dl {
-        // infer with text embeddings
-        let ys = model.forward_with_te(&xs, &text_embeddings)?;
+        let ys = model.forward_with_embedding(&xs, &embedding)?;
         println!("ys: {:?}", ys);
 
         for (x, y) in xs.iter().zip(ys.iter()) {
@@ -78,7 +95,7 @@ fn main() -> Result<()> {
             annotator.annotate(x, y)?.save(format!(
                 "{}.jpg",
                 usls::Dir::Current
-                    .base_dir_with_subs(&["runs", model.spec()])?
+                    .base_dir_with_subs(&["runs", "YOLOE-prompt", model.spec()])?
                     .join(usls::timestamp(None))
                     .display(),
             ))?;

examples/yoloe-text-prompt/README.md

@@ -42,6 +42,7 @@ mod svtr;
 mod swin2sr;
 mod trocr;
 mod yolo;
+mod yoloe;
 mod yolop;
 
 pub use blip::*;

examples/yoloe/README.md

@@ -1,7 +1,6 @@
 use crate::{
     models::YOLOPredsFormat, Config, ResizeMode, Scale, Task, NAMES_COCO_80,
-    NAMES_COCO_KEYPOINTS_17, NAMES_DOTA_V1_15, NAMES_IMAGENET_1K, NAMES_YOLOE_4585,
-    NAMES_YOLO_DOCLAYOUT_10,
+    NAMES_COCO_KEYPOINTS_17, NAMES_DOTA_V1_15, NAMES_IMAGENET_1K, NAMES_YOLO_DOCLAYOUT_10,
 };
 
 impl Config {
@@ -90,64 +89,6 @@ impl Config {
             .with_model_file("doclayout-docstructbench.onnx") // TODO: batch_size > 1
     }
 
-    /// Creates a base YOLOE configuration with 4585 classes.
-    ///
-    /// Configures the model for instance segmentation with a large class vocabulary.
-    pub fn yoloe() -> Self {
-        Self::yolo()
-            .with_task(Task::InstanceSegmentation)
-            .with_class_names(&NAMES_YOLOE_4585)
-    }
-    /// Creates a configuration for YOLOE-v8s segmentation model.
-    /// Uses the small variant of YOLOv8 architecture.
-    pub fn yoloe_v8s_seg_pf() -> Self {
-        Self::yoloe()
-            .with_version(8.into())
-            .with_scale(Scale::S)
-            .with_model_file("yoloe-v8s-seg-pf.onnx")
-    }
-
-    /// Creates a configuration for YOLOE-v8m segmentation model.
-    /// Uses the medium variant of YOLOv8 architecture.
-    pub fn yoloe_v8m_seg_pf() -> Self {
-        Self::yoloe()
-            .with_version(8.into())
-            .with_scale(Scale::M)
-            .with_model_file("yoloe-v8m-seg-pf.onnx")
-    }
-
-    /// Creates a configuration for YOLOE-v8l segmentation model.
-    /// Uses the large variant of YOLOv8 architecture.
-    pub fn yoloe_v8l_seg_pf() -> Self {
-        Self::yoloe()
-            .with_version(8.into())
-            .with_scale(Scale::L)
-            .with_model_file("yoloe-v8l-seg-pf.onnx")
-    }
-
-    /// Creates a configuration for YOLOE-11s segmentation model.
-    /// Uses the small variant of YOLOv11 architecture.
-    pub fn yoloe_11s_seg_pf() -> Self {
-        Self::yoloe()
-            .with_version(11.into())
-            .with_scale(Scale::S)
-            .with_model_file("yoloe-11s-seg-pf.onnx")
-    }
-
-    pub fn yoloe_11m_seg_pf() -> Self {
-        Self::yoloe()
-            .with_version(11.into())
-            .with_scale(Scale::M)
-            .with_model_file("yoloe-11m-seg-pf.onnx")
-    }
-
-    pub fn yoloe_11l_seg_pf() -> Self {
-        Self::yoloe()
-            .with_version(11.into())
-            .with_scale(Scale::L)
-            .with_model_file("yoloe-11l-seg-pf.onnx")
-    }
-
     pub fn fastsam_s() -> Self {
         Self::yolo_segment()
             .with_class_names(&["object"])
@@ -177,60 +118,4 @@ impl Config {
             .with_scale(Scale::X)
             .with_model_file("rtdetr-x.onnx")
     }
-
-    fn yoloe_seg_tp() -> Self {
-        Self::yolo()
-            .with_batch_size_all(1)
-            .with_nc(80)
-            .with_model_ixx(1, 1, (1, 80, 300).into()) // max_text_classes
-            .with_task(Task::InstanceSegmentation)
-            .with_textual_file("mobileclip/blt-textual.onnx")
-            .with_model_max_length(77)
-            .with_textual_ixx(0, 1, 77.into())
-            .with_tokenizer_file("clip/tokenizer.json")
-            .with_tokenizer_config_file("clip/tokenizer_config.json")
-            .with_special_tokens_map_file("clip/special_tokens_map.json")
-    }
-
-    pub fn yoloe_v8s_seg_tp() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::S)
-            .with_model_file("yoloe-v8s-seg-tp.onnx")
-    }
-
-    pub fn yoloe_v8m_seg_tp() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::M)
-            .with_model_file("yoloe-v8m-seg-tp.onnx")
-    }
-
-    pub fn yoloe_v8l_seg_tp() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::L)
-            .with_model_file("yoloe-v8l-seg-tp.onnx")
-    }
-
-    pub fn yoloe_11s_seg_tp() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::S)
-            .with_model_file("yoloe-11s-seg-tp.onnx")
-    }
-
-    pub fn yoloe_11m_seg_tp() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::M)
-            .with_model_file("yoloe-11m-seg-tp.onnx")
-    }
-
-    pub fn yoloe_11l_seg_tp() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::L)
-            .with_model_file("yoloe-11l-seg-tp.onnx")
-    }
 }