This commit is contained in:
jamjamjon
2025-05-19 23:25:48 +08:00
parent 5394c9ba27
commit 26de63d239
36 changed files with 4946 additions and 329 deletions

View File

@ -27,27 +27,29 @@ cargo run -r --example yolo -- --task detect --ver v8 --num-classes 6 --model xx
# Classify
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
cargo run -r --example yolo -- --task classify --ver 8 --scale n --image-width 224 --image-height 224 --use-imagenet-1k-classes # YOLOv8
cargo run -r --example yolo -- --task classify --ver 11 --scale n --image-width 224 --image-height 224 # YOLO11
# Detect
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
cargo run -r --example yolo -- --task detect --ver 5 --scale n --use-coco-80-classes --dtype fp16 # YOLOv5
cargo run -r --example yolo -- --task detect --ver 6 --scale n --use-coco-80-classes --dtype fp16 # YOLOv6
cargo run -r --example yolo -- --task detect --ver 7 --scale t --use-coco-80-classes --dtype fp16 # YOLOv7
cargo run -r --example yolo -- --task detect --ver 8 --scale n --use-coco-80-classes --dtype fp16 # YOLOv8
cargo run -r --example yolo -- --task detect --ver 9 --scale t --use-coco-80-classes --dtype fp16 # YOLOv9
cargo run -r --example yolo -- --task detect --ver 10 --scale n --use-coco-80-classes --dtype fp16 # YOLOv10
cargo run -r --example yolo -- --task detect --ver 11 --scale n --use-coco-80-classes --dtype fp16 # YOLO11
cargo run -r --example yolo -- --task detect --ver 12 --scale n --use-coco-80-classes --dtype fp16 # YOLOv12
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 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 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
cargo run -r --example yolo -- --task segment --ver 5 --scale n --use-coco-80-classes --dtype fp16 # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver 8 --scale n --use-coco-80-classes --dtype fp16 # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver 9 --scale c --use-coco-80-classes --dtype fp16 # YOLOv9-Segment
cargo run -r --example yolo -- --task segment --ver 11 --scale n --use-coco-80-classes --dtype fp16 # YOLO11-Segment
# Obb
cargo run -r --example yolo -- --ver 8 --task obb --scale n --image-width 1024 --image-height 1024 --source images/dota.png # YOLOv8-Obb

View File

@ -5,21 +5,21 @@ use usls::{
};
#[derive(argh::FromArgs, Debug)]
/// Example
/// YOLO Example
struct Args {
/// model file
/// model file(.onnx)
#[argh(option)]
model: Option<String>,
/// source
/// source: image, image folder, video stream
#[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
source: String,
/// dtype
/// model dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// task
/// task: det, seg, pose, classify, obb
#[argh(option, default = "String::from(\"det\")")]
task: String,
@ -27,101 +27,101 @@ struct Args {
#[argh(option, default = "8.0")]
ver: f32,
/// device
/// device: cuda, cpu, mps
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// scale
/// scale: n, s, m, l, x
#[argh(option, default = "String::from(\"n\")")]
scale: String,
/// trt_fp16
/// enable TensorRT FP16
#[argh(option, default = "true")]
trt_fp16: bool,
/// batch_size
/// batch size
#[argh(option, default = "1")]
batch_size: usize,
/// min_batch_size
/// bin batch size: For TensorRT
#[argh(option, default = "1")]
min_batch_size: usize,
/// max_batch_size
/// max Batch size: For TensorRT
#[argh(option, default = "4")]
max_batch_size: usize,
/// min_image_width
/// min image width: For TensorRT
#[argh(option, default = "224")]
min_image_width: isize,
/// image_width
/// image width: For TensorRT
#[argh(option, default = "640")]
image_width: isize,
/// max_image_width
/// max image width: For TensorRT
#[argh(option, default = "1280")]
max_image_width: isize,
/// min_image_height
/// min image height: For TensorRT
#[argh(option, default = "224")]
min_image_height: isize,
/// image_height
/// image height: For TensorRT
#[argh(option, default = "640")]
image_height: isize,
/// max_image_height
/// max image height: For TensorRT
#[argh(option, default = "1280")]
max_image_height: isize,
/// num_classes
/// num classes
#[argh(option)]
num_classes: Option<usize>,
/// num_keypoints
/// num keypoints
#[argh(option)]
num_keypoints: Option<usize>,
/// use_coco_80_classes
#[argh(switch)]
use_coco_80_classes: bool,
/// use_coco_17_keypoints_classes
#[argh(switch)]
use_coco_17_keypoints_classes: bool,
/// use_imagenet_1k_classes
#[argh(switch)]
use_imagenet_1k_classes: 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
/// class names
#[argh(option)]
class_names: Vec<String>,
/// keypoint_names
/// keypoint names
#[argh(option)]
keypoint_names: Vec<String>,
/// topk
/// top-k
#[argh(option, default = "5")]
topk: usize,
/// use COCO 80 classes
#[argh(switch)]
use_coco_80_classes: bool,
/// use COCO 17 keypoints classes
#[argh(switch)]
use_coco_17_keypoints_classes: bool,
/// use ImageNet 1K classes
#[argh(switch)]
use_imagenet_1k_classes: bool,
/// confidences
#[argh(option)]
confs: Vec<f32>,
/// keypoint nonfidences
#[argh(option)]
keypoint_confs: Vec<f32>,
/// exclude nlasses
#[argh(option)]
exclude_classes: Vec<usize>,
/// retain classes
#[argh(option)]
retain_classes: Vec<usize>,
}
fn main() -> Result<()> {
@ -129,9 +129,7 @@ fn main() -> Result<()> {
.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();
let mut config = ModelConfig::yolo()
.with_model_file(&args.model.unwrap_or_default())
.with_task(args.task.as_str().try_into()?)
@ -139,7 +137,6 @@ fn main() -> Result<()> {
.with_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_trt_fp16(args.trt_fp16)
.with_model_ixx(
0,
@ -174,27 +171,21 @@ fn main() -> Result<()> {
.with_topk(args.topk)
.retain_classes(&args.retain_classes)
.exclude_classes(&args.exclude_classes);
if args.use_coco_80_classes {
config = config.with_class_names(&NAMES_COCO_80);
}
if args.use_coco_17_keypoints_classes {
config = config.with_keypoint_names(&NAMES_COCO_KEYPOINTS_17);
}
if args.use_imagenet_1k_classes {
config = config.with_class_names(&NAMES_IMAGENET_1K);
}
if let Some(nc) = args.num_classes {
config = config.with_nc(nc);
}
if let Some(nk) = args.num_keypoints {
config = config.with_nk(nk);
}
if !args.class_names.is_empty() {
config = config.with_class_names(
&args
@ -204,7 +195,6 @@ fn main() -> Result<()> {
.collect::<Vec<_>>(),
);
}
if !args.keypoint_names.is_empty() {
config = config.with_keypoint_names(
&args
@ -216,7 +206,7 @@ fn main() -> Result<()> {
}
// build model
let mut model = YOLO::try_from(config.auto_yolo_model_file().commit()?)?;
let mut model = YOLO::new(config.commit()?)?;
// build dataloader
let dl = DataLoader::new(&args.source)?
@ -256,6 +246,7 @@ fn main() -> Result<()> {
}
}
// summary
model.summary();
Ok(())