mii/usls
1
0
Fork 0
mirror of https://github.com/mii443/usls.git synced 2025-12-03 11:08:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add YOLOE model (#88)

Browse Source
This commit is contained in:
Jamjamjon
2025-05-10 01:47:58 +08:00
committed by GitHub
parent 68a70e6e7b
commit c54775cedd
5 changed files with 152 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/models/picodet/config.rs
View File

@@ -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)
}
}

60
src/models/yolo/config.rs
View File

@@ -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")
}
}

33
src/utils/names.rs
View File

@@ -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",