Add YOLOE model (#88)

examples/yoloe/README.md

@@ -0,0 +1,6 @@
+## Quick Start
+
+```shell
+cargo run -r --example yoloe 
+```
+

examples/yoloe/main.rs

@@ -0,0 +1,59 @@
+use anyhow::Result;
+use usls::{models::YOLO, Annotator, DataLoader, Options, Style};
+
+#[derive(argh::FromArgs)]
+/// Example
+struct Args {
+    /// dtype
+    #[argh(option, default = "String::from(\"fp16\")")]
+    dtype: String,
+
+    /// device
+    #[argh(option, default = "String::from(\"cpu:0\")")]
+    device: String,
+}
+
+fn main() -> Result<()> {
+    tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
+        .init();
+
+    let args: Args = argh::from_env();
+
+    // options
+    let options = Options::yoloe_v8s_seg_pf()
+        // yoloe_v8m_seg_pf()
+        // yoloe_v8l_seg_pf()
+        // yoloe_11s_seg_pf()
+        // yoloe_11m_seg_pf()
+        // yoloe_11l_seg_pf()
+        .with_model_dtype(args.dtype.as_str().try_into()?)
+        .with_model_device(args.device.as_str().try_into()?)
+        .commit()?;
+    let mut model = YOLO::new(options)?;
+
+    // load
+    let xs = DataLoader::try_read_n(&["./assets/bus.jpg"])?;
+
+    // run
+    let ys = model.forward(&xs)?;
+    println!("{:?}", ys);
+
+    // annotate
+    let annotator = Annotator::default()
+        .with_hbb_style(Style::hbb().with_draw_fill(true))
+        .with_mask_style(Style::mask().with_draw_mask_polygon_largest(true));
+
+    for (x, y) in xs.iter().zip(ys.iter()) {
+        annotator.annotate(x, y)?.save(format!(
+            "{}.jpg",
+            usls::Dir::Current
+                .base_dir_with_subs(&["runs", model.spec()])?
+                .join(usls::timestamp(None))
+                .display(),
+        ))?;
+    }
+
+    Ok(())
+}

src/models/picodet/config.rs

@@ -1,4 +1,7 @@
-use crate::{ResizeMode, NAMES_COCO_80};
+use crate::{
+    ResizeMode, NAMES_COCO_80, NAMES_PICODET_LAYOUT_17, NAMES_PICODET_LAYOUT_3,
+    NAMES_PICODET_LAYOUT_5,
+};
 
 /// Model configuration for `PicoDet`
 impl crate::Options {
@@ -26,36 +29,18 @@ impl crate::Options {
     pub fn picodet_layout_1x() -> Self {
         Self::picodet()
             .with_model_file("layout-1x.onnx")
-            .with_class_names(&["Text", "Title", "List", "Table", "Figure"])
+            .with_class_names(&NAMES_PICODET_LAYOUT_5)
     }
 
     pub fn picodet_l_layout_3cls() -> Self {
         Self::picodet()
             .with_model_file("l-layout-3cls.onnx")
-            .with_class_names(&["image", "table", "seal"])
+            .with_class_names(&NAMES_PICODET_LAYOUT_3)
     }
 
     pub fn picodet_l_layout_17cls() -> Self {
         Self::picodet()
             .with_model_file("l-layout-17cls.onnx")
-            .with_class_names(&[
-                "paragraph_title",
-                "image",
-                "text",
-                "number",
-                "abstract",
-                "content",
-                "figure_title",
-                "formula",
-                "table",
-                "table_title",
-                "reference",
-                "doc_title",
-                "footnote",
-                "header",
-                "algorithm",
-                "footer",
-                "seal",
-            ])
+            .with_class_names(&NAMES_PICODET_LAYOUT_17)
     }
 }

src/models/yolo/config.rs

@@ -1,4 +1,7 @@
-use crate::{models::YOLOPredsFormat, Options, ResizeMode, Scale, Task, NAMES_COCO_KEYPOINTS_17};
+use crate::{
+    models::YOLOPredsFormat, Options, ResizeMode, Scale, Task, NAMES_COCO_KEYPOINTS_17,
+    NAMES_YOLO_DOCLAYOUT_10,
+};
 
 impl Options {
     pub fn yolo() -> Self {
@@ -19,18 +22,7 @@ impl Options {
             .with_model_ixx(0, 2, (640, 1024, 1024).into())
             .with_model_ixx(0, 3, (640, 1024, 1024).into())
             .with_class_confs(&[0.4])
-            .with_class_names(&[
-                "title",
-                "plain text",
-                "abandon",
-                "figure",
-                "figure_caption",
-                "table",
-                "table_caption",
-                "table_footnote",
-                "isolate_formula",
-                "formula_caption",
-            ])
+            .with_class_names(&NAMES_YOLO_DOCLAYOUT_10)
     }
 
     pub fn yolo_classify() -> Self {
@@ -197,4 +189,46 @@ impl Options {
             .with_model_version(11.into())
             .with_model_scale(Scale::X)
     }
+
+    pub fn yoloe_v8s_seg_pf() -> Self {
+        Self::yolo()
+            .with_model_version(8.into())
+            .with_model_scale(Scale::S)
+            .with_model_file("yoloe-v8s-seg-pf.onnx")
+    }
+
+    pub fn yoloe_v8m_seg_pf() -> Self {
+        Self::yolo()
+            .with_model_version(8.into())
+            .with_model_scale(Scale::M)
+            .with_model_file("yoloe-v8m-seg-pf.onnx")
+    }
+
+    pub fn yoloe_v8l_seg_pf() -> Self {
+        Self::yolo()
+            .with_model_version(8.into())
+            .with_model_scale(Scale::L)
+            .with_model_file("yoloe-v8l-seg-pf.onnx")
+    }
+
+    pub fn yoloe_11s_seg_pf() -> Self {
+        Self::yolo()
+            .with_model_version(11.into())
+            .with_model_scale(Scale::S)
+            .with_model_file("yoloe-11s-seg-pf.onnx")
+    }
+
+    pub fn yoloe_11m_seg_pf() -> Self {
+        Self::yolo()
+            .with_model_version(11.into())
+            .with_model_scale(Scale::M)
+            .with_model_file("yoloe-v8m-seg-pf.onnx")
+    }
+
+    pub fn yoloe_11l_seg_pf() -> Self {
+        Self::yolo()
+            .with_model_version(11.into())
+            .with_model_scale(Scale::L)
+            .with_model_file("yoloe-11l-seg-pf.onnx")
+    }
 }

src/utils/names.rs

@@ -1,3 +1,36 @@
+pub const NAMES_YOLO_DOCLAYOUT_10: [&str; 10] = [
+    "title",
+    "plain text",
+    "abandon",
+    "figure",
+    "figure_caption",
+    "table",
+    "table_caption",
+    "table_footnote",
+    "isolate_formula",
+    "formula_caption",
+];
+pub const NAMES_PICODET_LAYOUT_17: [&str; 17] = [
+    "paragraph_title",
+    "image",
+    "text",
+    "number",
+    "abstract",
+    "content",
+    "figure_title",
+    "formula",
+    "table",
+    "table_title",
+    "reference",
+    "doc_title",
+    "footnote",
+    "header",
+    "algorithm",
+    "footer",
+    "seal",
+];
+pub const NAMES_PICODET_LAYOUT_3: [&str; 3] = ["image", "table", "seal"];
+pub const NAMES_PICODET_LAYOUT_5: [&str; 5] = ["Text", "Title", "List", "Table", "Figure"];
 pub const NAMES_COCO_KEYPOINTS_17: [&str; 17] = [
     "nose",
     "left_eye",