diff --git a/examples/yoloe/README.md b/examples/yoloe/README.md new file mode 100644 index 0000000..110c5a4 --- /dev/null +++ b/examples/yoloe/README.md @@ -0,0 +1,6 @@ +## Quick Start + +```shell +cargo run -r --example yoloe +``` + diff --git a/examples/yoloe/main.rs b/examples/yoloe/main.rs new file mode 100644 index 0000000..841945c --- /dev/null +++ b/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(()) +} diff --git a/src/models/picodet/config.rs b/src/models/picodet/config.rs index 8a28f8f..b3988a5 100644 --- a/src/models/picodet/config.rs +++ b/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) } } diff --git a/src/models/yolo/config.rs b/src/models/yolo/config.rs index 133f3f3..8034827 100644 --- a/src/models/yolo/config.rs +++ b/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") + } } diff --git a/src/utils/names.rs b/src/utils/names.rs index c814351..644973e 100644 --- a/src/utils/names.rs +++ b/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",