mii/usls
1
0
Fork 0
mirror of https://github.com/mii443/usls.git synced 2025-08-22 15:45:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

🐍 v0.1.0 (#53)

Browse Source
This commit is contained in:
Jamjamjon
2025-01-12 16:59:57 +08:00
committed by GitHub
parent 4e932c4910
commit 0f2d84b8c5
256 changed files with 12485 additions and 9088 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
examples/yolo/README.md
View File

@ -1,175 +1,65 @@
<h1 align='center'>YOLO-Series</h1>
| Detection | Instance Segmentation | Pose |
| :----------------------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------------------------------: | :-----------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-det.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-seg.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-pose.png' width="300px">` |
| Detection | Instance Segmentation | Pose |
| :---------------: | :------------------------: |:---------------: |
| <img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-det.png' width="300px"> | <img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-seg.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-pose.png' width="300px"> |
| Classification | Obb |
| :------------------------: |:------------------------: |
|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-cls.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-obb-2.png' width="628px">
| Head Detection | Fall Detection | Trash Detection |
| :------------------------: |:------------------------: |:------------------------: |
|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-head.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-falldown.png' width="300px">|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-trash.png' width="300px">
| YOLO-World | Face Parsing | FastSAM |
| :------------------------: |:------------------------: |:------------------------: |
|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-yolov8-world.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-face-parsing.png' width="300px">|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-fastsam.png' width="300px">
| Classification | Obb |
| :----------------------------------------------------------------------------------------------------: | :------------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-cls.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-obb-2.png' width="628px">` |
| Head Detection | Fall Detection | Trash Detection |
| :-----------------------------------------------------------------------------------------------------: | :---------------------------------------------------------------------------------------------------------: | :------------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-head.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-falldown.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-trash.png' width="300px">` |
| YOLO-World | Face Parsing | FastSAM |
| :-------------------------------------------------------------------------------------------------------------: | :-------------------------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-yolov8-world.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-face-parsing.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-fastsam.png' width="300px">` |
## Quick Start
```Shell
# customized
cargo run -r --example yolo -- --task detect --ver v8 --nc 6 --model xxx.onnx # YOLOv8
# Your customized YOLOv8 model
cargo run -r --example yolo -- --task detect --ver v8 --num-classes 6 --model xxx.onnx # YOLOv8
# Classify
cargo run -r --example yolo -- --task classify --ver v5 --scale s --width 224 --height 224 --nc 1000 # YOLOv5
cargo run -r --example yolo -- --task classify --ver v8 --scale n --width 224 --height 224 --nc 1000 # YOLOv8
cargo run -r --example yolo -- --task classify --ver v11 --scale n --width 224 --height 224 --nc 1000 # YOLOv11
cargo run -r --example yolo -- --task classify --ver 5 --scale s --image-width 224 --image-height 224 --num-classes 1000 --use-imagenet-1k-classes # YOLOv5
cargo run -r --example yolo -- --task classify --ver 8 --scale n --image-width 224 --image-height 224 # YOLOv8
cargo run -r --example yolo -- --task classify --ver 11 --scale n --image-width 224 --image-height 224 # YOLOv11
# Detect
cargo run -r --example yolo -- --task detect --ver v5 --scale n # YOLOv5
cargo run -r --example yolo -- --task detect --ver v6 --scale n # YOLOv6
cargo run -r --example yolo -- --task detect --ver v7 --scale t # YOLOv7
cargo run -r --example yolo -- --task detect --ver v8 --scale n # YOLOv8
cargo run -r --example yolo -- --task detect --ver v9 --scale t # YOLOv9
cargo run -r --example yolo -- --task detect --ver v10 --scale n # YOLOv10
cargo run -r --example yolo -- --task detect --ver v11 --scale n # YOLOv11
cargo run -r --example yolo -- --task detect --ver rtdetr --scale l # RTDETR
cargo run -r --example yolo -- --task detect --ver v8 --model yolo/v8-s-world-v2-shoes.onnx # YOLOv8-world
cargo run -r --example yolo -- --task detect --ver 5 --scale n --use-coco-80-classes # YOLOv5
cargo run -r --example yolo -- --task detect --ver 6 --scale n --use-coco-80-classes # YOLOv6
cargo run -r --example yolo -- --task detect --ver 7 --scale t --use-coco-80-classes # YOLOv7
cargo run -r --example yolo -- --task detect --ver 8 --scale n --use-coco-80-classes # YOLOv8
cargo run -r --example yolo -- --task detect --ver 9 --scale t --use-coco-80-classes # YOLOv9
cargo run -r --example yolo -- --task detect --ver 10 --scale n --use-coco-80-classes # YOLOv10
cargo run -r --example yolo -- --task detect --ver 11 --scale n --use-coco-80-classes # YOLOv11
cargo run -r --example yolo -- --task detect --ver 8 --model v8-s-world-v2-shoes.onnx # YOLOv8-world
# Pose
cargo run -r --example yolo -- --task pose --ver v8 --scale n # YOLOv8-Pose
cargo run -r --example yolo -- --task pose --ver v11 --scale n # YOLOv11-Pose
cargo run -r --example yolo -- --task pose --ver 8 --scale n # YOLOv8-Pose
cargo run -r --example yolo -- --task pose --ver 11 --scale n # YOLOv11-Pose
# Segment
cargo run -r --example yolo -- --task segment --ver v5 --scale n # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver v8 --scale n # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver v11 --scale n # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver v8 --model yolo/FastSAM-s-dyn-f16.onnx # FastSAM
cargo run -r --example yolo -- --task segment --ver 5 --scale n # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver 8 --scale n # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver 11 --scale n # YOLOv8-Segment
# Obb
cargo run -r --example yolo -- --ver v8 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv8-Obb
cargo run -r --example yolo -- --ver v11 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv11-Obb
cargo run -r --example yolo -- --ver 8 --task obb --scale n --image-width 1024 --image-height 1024 --source images/dota.png # YOLOv8-Obb
cargo run -r --example yolo -- --ver 11 --task obb --scale n --image-width 1024 --image-height 1024 --source images/dota.png # YOLOv11-Obb
```
**`cargo run -r --example yolo -- --help` for more options**
## YOLOs configs with `Options`
<details open>
<summary>Use official YOLO Models</summary>
```Rust
let options = Options::default()
.with_yolo_version(YOLOVersion::V5) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
.with_yolo_task(YOLOTask::Classify) // YOLOTask: Classify, Detect, Pose, Segment, Obb
.with_model("xxxx.onnx")?;
```
</details>
<details open>
<summary>Cutomized your own YOLO model</summary>
```Rust
// This config is for YOLOv8-Segment
use usls::{AnchorsPosition, BoxType, ClssType, YOLOPreds};
let options = Options::default()
.with_yolo_preds(
YOLOPreds {
bbox: Some(BoxType::Cxcywh),
clss: ClssType::Clss,
coefs: Some(true),
anchors: Some(AnchorsPosition::After),
..Default::default()
}
)
// .with_nc(80)
// .with_names(&COCO_CLASS_NAMES_80)
.with_model("xxxx.onnx")?;
```
</details>
## Other YOLOv8 Solution Models
| Model | Weights | Datasets|
|:---------------------: | :--------------------------: | :-------------------------------: |
| Face-Landmark Detection | [yolov8-face-dyn-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-n-face-dyn-f16.onnx) | |
| Head Detection | [yolov8-head-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-head-f16.onnx) | |
| Fall Detection | [yolov8-falldown-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-falldown-f16.onnx) | |
| Trash Detection | [yolov8-plastic-bag-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-plastic-bag-f16.onnx) | |
| FaceParsing | [yolov8-face-parsing-dyn](https://github.com/jamjamjon/assets/releases/download/yolo/v8-face-parsing-dyn.onnx) | [CelebAMask-HQ](https://github.com/switchablenorms/CelebAMask-HQ/tree/master/face_parsing)<br />[[Processed YOLO labels]](https://github.com/jamjamjon/assets/releases/download/yolo/CelebAMask-HQ-YOLO-Labels.zip)[[Python Script]](../../scripts/CelebAMask-HQ-To-YOLO-Labels.py) |
## Export ONNX Models
<details close>
<summary>YOLOv5</summary>
[Here](https://docs.ultralytics.com/yolov5/tutorials/model_export/)
</details>
<details close>
<summary>YOLOv6</summary>
[Here](https://github.com/meituan/YOLOv6/tree/main/deploy/ONNX)
</details>
<details close>
<summary>YOLOv7</summary>
[Here](https://github.com/WongKinYiu/yolov7?tab=readme-ov-file#export)
</details>
<details close>
<summary>YOLOv8, YOLOv11</summary>
```Shell
pip install -U ultralytics
# export onnx model with dynamic shapes
yolo export model=yolov8m.pt format=onnx simplify dynamic
yolo export model=yolov8m-cls.pt format=onnx simplify dynamic
yolo export model=yolov8m-pose.pt format=onnx simplify dynamic
yolo export model=yolov8m-seg.pt format=onnx simplify dynamic
yolo export model=yolov8m-obb.pt format=onnx simplify dynamic
# export onnx model with fixed shapes
yolo export model=yolov8m.pt format=onnx simplify
yolo export model=yolov8m-cls.pt format=onnx simplify
yolo export model=yolov8m-pose.pt format=onnx simplify
yolo export model=yolov8m-seg.pt format=onnx simplify
yolo export model=yolov8m-obb.pt format=onnx simplify
```
</details>
<details close>
<summary>YOLOv9</summary>
[Here](https://github.com/WongKinYiu/yolov9/blob/main/export.py)
</details>
<details close>
<summary>YOLOv10</summary>
[Here](https://github.com/THU-MIG/yolov10#export)
</details>
| Model | Weights |
| :---------------------: | :------------------------------------------------------: |
| Face-Landmark Detection | [yolov8-n-face](https://github.com/jamjamjon/assets/releases/download/yolo/v8-n-face-fp16.onnx) |
| Head Detection | [yolov8-head](https://github.com/jamjamjon/assets/releases/download/yolo/v8-head-fp16.onnx) |
| Fall Detection | [yolov8-falldown](https://github.com/jamjamjon/assets/releases/download/yolo/v8-falldown-fp16.onnx) |
| Trash Detection | [yolov8-plastic-bag](https://github.com/jamjamjon/assets/releases/download/yolo/v8-plastic-bag-fp16.onnx) |
| FaceParsing | [yolov8-face-parsing-seg](https://github.com/jamjamjon/assets/releases/download/yolo/v8-face-parsing.onnx) |

362
examples/yolo/main.rs
View File

@ -1,171 +1,213 @@
use anyhow::Result;
use clap::Parser;
use usls::{
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOScale, YOLOTask,
YOLOVersion, COCO_SKELETONS_16,
models::YOLO, Annotator, DataLoader, Options, COCO_CLASS_NAMES_80, COCO_SKELETONS_16,
IMAGENET_NAMES_1K,
};
#[derive(Parser, Clone)]
#[command(author, version, about, long_about = None)]
pub struct Args {
/// Path to the model
#[arg(long)]
pub model: Option<String>,
#[derive(argh::FromArgs, Debug)]
/// Example
struct Args {
/// model file
#[argh(option)]
model: Option<String>,
/// Input source path
#[arg(long, default_value_t = String::from("./assets/bus.jpg"))]
pub source: String,
/// source
#[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
source: String,
/// YOLO Task
#[arg(long, value_enum, default_value_t = YOLOTask::Detect)]
pub task: YOLOTask,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// YOLO Version
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)]
pub ver: YOLOVersion,
/// task
#[argh(option, default = "String::from(\"det\")")]
task: String,
/// YOLO Scale
#[arg(long, value_enum, default_value_t = YOLOScale::N)]
pub scale: YOLOScale,
/// version
#[argh(option, default = "8.0")]
ver: f32,
/// Batch size
#[arg(long, default_value_t = 1)]
pub batch_size: usize,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// Minimum input width
#[arg(long, default_value_t = 224)]
pub width_min: isize,
/// scale
#[argh(option, default = "String::from(\"n\")")]
scale: String,
/// Input width
#[arg(long, default_value_t = 640)]
pub width: isize,
/// trt_fp16
#[argh(option, default = "true")]
trt_fp16: bool,
/// Maximum input width
#[arg(long, default_value_t = 1024)]
pub width_max: isize,
/// find_contours
#[argh(option, default = "true")]
find_contours: bool,
/// Minimum input height
#[arg(long, default_value_t = 224)]
pub height_min: isize,
/// batch_size
#[argh(option, default = "1")]
batch_size: usize,
/// Input height
#[arg(long, default_value_t = 640)]
pub height: isize,
/// min_batch_size
#[argh(option, default = "1")]
min_batch_size: usize,
/// Maximum input height
#[arg(long, default_value_t = 1024)]
pub height_max: isize,
/// max_batch_size
#[argh(option, default = "4")]
max_batch_size: usize,
/// Number of classes
#[arg(long, default_value_t = 80)]
pub nc: usize,
/// min_image_width
#[argh(option, default = "224")]
min_image_width: isize,
/// Class confidence
#[arg(long)]
pub confs: Vec<f32>,
/// image_width
#[argh(option, default = "640")]
image_width: isize,
/// Enable TensorRT support
#[arg(long)]
pub trt: bool,
/// max_image_width
#[argh(option, default = "1280")]
max_image_width: isize,
/// Enable CUDA support
#[arg(long)]
pub cuda: bool,
/// min_image_height
#[argh(option, default = "224")]
min_image_height: isize,
/// Enable CoreML support
#[arg(long)]
pub coreml: bool,
/// image_height
#[argh(option, default = "640")]
image_height: isize,
/// Use TensorRT half precision
#[arg(long)]
pub half: bool,
/// max_image_height
#[argh(option, default = "1280")]
max_image_height: isize,
/// Device ID to use
#[arg(long, default_value_t = 0)]
pub device_id: usize,
/// num_classes
#[argh(option)]
num_classes: Option<usize>,
/// Enable performance profiling
#[arg(long)]
pub profile: bool,
/// num_keypoints
#[argh(option)]
num_keypoints: Option<usize>,
/// Disable contour drawing
#[arg(long)]
pub no_contours: bool,
/// use_coco_80_classes
#[argh(switch)]
use_coco_80_classes: bool,
/// Show result
#[arg(long)]
pub view: bool,
/// use_imagenet_1k_classes
#[argh(switch)]
use_imagenet_1k_classes: bool,
/// Do not save output
#[arg(long)]
pub nosave: bool,
/// confs
#[argh(option)]
confs: Vec<f32>,
/// keypoint_confs
#[argh(option)]
keypoint_confs: Vec<f32>,
/// exclude_classes
#[argh(option)]
exclude_classes: Vec<usize>,
/// retain_classes
#[argh(option)]
retain_classes: Vec<usize>,
/// class_names
#[argh(option)]
class_names: Vec<String>,
/// keypoint_names
#[argh(option)]
keypoint_names: Vec<String>,
}
fn main() -> Result<()> {
let args = Args::parse();
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// model path
let path = match &args.model {
None => format!(
"yolo/{}-{}-{}.onnx",
args.ver.name(),
args.scale.name(),
args.task.name()
),
Some(x) => x.to_string(),
};
let args: Args = argh::from_env();
// saveout
let saveout = match &args.model {
None => format!(
"{}-{}-{}",
args.ver.name(),
args.scale.name(),
args.task.name()
),
Some(x) => {
let p = std::path::PathBuf::from(&x);
p.file_stem().unwrap().to_str().unwrap().to_string()
}
};
// device
let device = if args.cuda {
Device::Cuda(args.device_id)
} else if args.trt {
Device::Trt(args.device_id)
} else if args.coreml {
Device::CoreML(args.device_id)
} else {
Device::Cpu(args.device_id)
};
// build options
let options = Options::new()
.with_model(&path)?
.with_yolo_version(args.ver)
.with_yolo_task(args.task)
.with_device(device)
.with_trt_fp16(args.half)
.with_ixx(0, 0, (1, args.batch_size as _, 4).into())
.with_ixx(0, 2, (args.height_min, args.height, args.height_max).into())
.with_ixx(0, 3, (args.width_min, args.width, args.width_max).into())
.with_confs(if args.confs.is_empty() {
let mut options = Options::yolo()
.with_model_file(&args.model.unwrap_or_default())
.with_model_task(args.task.as_str().try_into()?)
.with_model_version(args.ver.into())
.with_model_scale(args.scale.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_trt_fp16(args.trt_fp16)
.with_model_ixx(
0,
0,
(args.min_batch_size, args.batch_size, args.max_batch_size).into(),
)
.with_model_ixx(
0,
2,
(
args.min_image_height,
args.image_height,
args.max_image_height,
)
.into(),
)
.with_model_ixx(
0,
3,
(args.min_image_width, args.image_width, args.max_image_width).into(),
)
.with_class_confs(if args.confs.is_empty() {
&[0.2, 0.15]
} else {
&args.confs
})
.with_nc(args.nc)
// .with_names(&COCO_CLASS_NAMES_80)
// .with_names2(&COCO_KEYPOINTS_17)
.with_find_contours(!args.no_contours) // find contours or not
.exclude_classes(&[0])
// .retain_classes(&[0, 5])
.with_profile(args.profile);
.with_keypoint_confs(if args.keypoint_confs.is_empty() {
&[0.5]
} else {
&args.keypoint_confs
})
.with_find_contours(args.find_contours)
.retain_classes(&args.retain_classes)
.exclude_classes(&args.exclude_classes);
if args.use_coco_80_classes {
options = options.with_class_names(&COCO_CLASS_NAMES_80);
}
if args.use_imagenet_1k_classes {
options = options.with_class_names(&IMAGENET_NAMES_1K);
}
if let Some(nc) = args.num_classes {
options = options.with_nc(nc);
}
if let Some(nk) = args.num_keypoints {
options = options.with_nk(nk);
}
if !args.class_names.is_empty() {
options = options.with_class_names(
&args
.class_names
.iter()
.map(|x| x.as_str())
.collect::<Vec<_>>(),
);
}
if !args.keypoint_names.is_empty() {
options = options.with_keypoint_names(
&args
.keypoint_names
.iter()
.map(|x| x.as_str())
.collect::<Vec<_>>(),
);
}
// build model
let mut model = YOLO::new(options)?;
let mut model = YOLO::try_from(options.commit()?)?;
// build dataloader
let dl = DataLoader::new(&args.source)?
@ -175,56 +217,28 @@ fn main() -> Result<()> {
// build annotator
let annotator = Annotator::default()
.with_skeletons(&COCO_SKELETONS_16)
.without_masks(true) // No masks plotting when doing segment task.
.without_masks(true)
.with_bboxes_thickness(3)
.with_keypoints_name(false) // Enable keypoints names
.with_saveout_subs(&["YOLO"])
.with_saveout(&saveout);
// build viewer
let mut viewer = if args.view {
Some(Viewer::new().with_delay(5).with_scale(1.).resizable(true))
} else {
None
};
.with_saveout(model.spec());
// run & annotate
for (xs, _paths) in dl {
// let ys = model.run(&xs)?; // way one
let ys = model.forward(&xs, args.profile)?; // way two
let images_plotted = annotator.plot(&xs, &ys, !args.nosave)?;
let ys = model.forward(&xs)?;
// extract bboxes
// for y in ys.iter() {
// if let Some(bboxes) = y.bboxes() {
// println!("[Bboxes]: Found {} objects", bboxes.len());
// for (i, bbox) in bboxes.iter().enumerate() {
// println!("{}: {:?}", i, bbox)
// }
// }
// }
// show image
match &mut viewer {
Some(viewer) => viewer.imshow(&images_plotted)?,
None => continue,
}
// check out window and key event
match &mut viewer {
Some(viewer) => {
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) {
break;
}
}
None => continue,
}
// write video
if !args.nosave {
match &mut viewer {
Some(viewer) => viewer.write_batch(&images_plotted)?,
None => continue,
}
}
// plot
annotator.annotate(&xs, &ys);
}
// finish video write
if !args.nosave {
if let Some(viewer) = &mut viewer {
viewer.finish_write()?;
}
}
model.summary();
Ok(())
}