Skip to content

Commit

Permalink
Add YOLOv11
Browse files Browse the repository at this point in the history
  • Loading branch information
jamjamjon committed Sep 30, 2024
1 parent 2cb9e57 commit f4d505c
Show file tree
Hide file tree
Showing 4 changed files with 211 additions and 144 deletions.
46 changes: 19 additions & 27 deletions examples/yolo/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,45 +24,37 @@
## Quick Start
```Shell

# customized
cargo run -r --example yolo -- --task detect --ver v8 --nc 6 --model xxx.onnx # YOLOv8

# Classify
cargo run -r --example yolo -- --task classify --ver v5 # YOLOv5
cargo run -r --example yolo -- --task classify --ver v8 # YOLOv8
cargo run -r --example yolo -- --task classify --ver v5 --scale n --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

# Detect
cargo run -r --example yolo -- --task detect --ver v5 # YOLOv5
cargo run -r --example yolo -- --task detect --ver v6 # YOLOv6
cargo run -r --example yolo -- --task detect --ver v7 # YOLOv7
cargo run -r --example yolo -- --task detect --ver v8 # YOLOv8
cargo run -r --example yolo -- --task detect --ver v9 # YOLOv9
cargo run -r --example yolo -- --task detect --ver v10 # YOLOv10
cargo run -r --example yolo -- --task detect --ver rtdetr # YOLOv8-RTDETR
cargo run -r --example yolo -- --task detect --ver v5 --scale n --nc 80 # YOLOv5
cargo run -r --example yolo -- --task detect --ver v6 --scale n --nc 80 # YOLOv6
cargo run -r --example yolo -- --task detect --ver v7 --scale t --nc 80 # YOLOv7
cargo run -r --example yolo -- --task detect --ver v8 --scale n --nc 80 # YOLOv8
cargo run -r --example yolo -- --task detect --ver v9 --scale t --nc 80 # YOLOv9
cargo run -r --example yolo -- --task detect --ver v10 --scale n --nc 80 # YOLOv10
cargo run -r --example yolo -- --task detect --ver rtdetr --scale l --nc 80 # YOLOv8-RTDETR
cargo run -r --example yolo -- --task detect --ver v8 --model yolov8s-world-v2-shoes.onnx # YOLOv8-world

# Pose
cargo run -r --example yolo -- --task pose --ver v8 # YOLOv8-Pose
cargo run -r --example yolo -- --task pose --ver v8 --scale s # YOLOv8-Pose

# Segment
cargo run -r --example yolo -- --task segment --ver v5 # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver v8 # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver v5 --scale n --nc 80 # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver v8 --scale n --nc 80 # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver v8 --model FastSAM-s-dyn-f16.onnx # FastSAM

# Obb
cargo run -r --example yolo -- --task obb --ver v8 # YOLOv8-Obb
cargo run -r --example yolo -- --ver v8 --task obb --scale s --source images/dota.png # YOLOv8-Obb
cargo run -r --example yolo -- --ver v11 --task obb --scale s --source images/dota.png # YOLOv11-Obb
```

<details close>
<summary>other options</summary>

`--source` to specify the input images
`--model` to specify the ONNX model
`--width --height` to specify the input resolution
`--nc` to specify the number of model's classes
`--plot` to annotate with inference results
`--profile` to profile
`--cuda --trt --coreml --device_id` to select device
`--half` to use float16 when using TensorRT EP

</details>
**`cargo run -r --example yolo -- --help` for more options**


## YOLOs configs with `Options`
Expand Down Expand Up @@ -140,7 +132,7 @@ let options = Options::default()
</details>

<details close>
<summary>YOLOv8</summary>
<summary>YOLOv8, YOLOv11</summary>

```Shell
pip install -U ultralytics
Expand Down
225 changes: 117 additions & 108 deletions examples/yolo/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,188 +2,160 @@ use anyhow::Result;
use clap::Parser;

use usls::{
models::YOLO, Annotator, DataLoader, Options, Vision, YOLOTask, YOLOVersion, COCO_KEYPOINTS_17,
COCO_SKELETONS_16,
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOScale, YOLOTask,
YOLOVersion, COCO_KEYPOINTS_17, COCO_SKELETONS_16,
};

#[derive(Parser, Clone)]
#[command(author, version, about, long_about = None)]
pub struct Args {
/// Path to the model
#[arg(long)]
pub model: Option<String>,

/// Input source path
#[arg(long, default_value_t = String::from("./assets/bus.jpg"))]
pub source: String,

/// YOLO Task
#[arg(long, value_enum, default_value_t = YOLOTask::Detect)]
pub task: YOLOTask,

/// YOLO Version
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)]
pub ver: YOLOVersion,

/// YOLO Scale
#[arg(long, value_enum, default_value_t = YOLOScale::N)]
pub scale: YOLOScale,

/// Batch size
#[arg(long, default_value_t = 1)]
pub batch_size: usize,

/// Minimum input width
#[arg(long, default_value_t = 224)]
pub width_min: isize,

/// Input width
#[arg(long, default_value_t = 640)]
pub width: isize,

/// Maximum input width
#[arg(long, default_value_t = 800)]
pub width_max: isize,

/// Minimum input height
#[arg(long, default_value_t = 224)]
pub height_min: isize,

/// Input height
#[arg(long, default_value_t = 640)]
pub height: isize,

/// Maximum input height
#[arg(long, default_value_t = 800)]
pub height_max: isize,

/// Number of classes
#[arg(long, default_value_t = 80)]
pub nc: usize,

/// Class confidence
#[arg(long)]
pub confs: Vec<f32>,

/// Enable TensorRT support
#[arg(long)]
pub trt: bool,

/// Enable CUDA support
#[arg(long)]
pub cuda: bool,

/// Enable CoreML support
#[arg(long)]
pub half: bool,
pub coreml: bool,

/// Use TensorRT half precision
#[arg(long)]
pub coreml: bool,
pub half: bool,

/// Device ID to use
#[arg(long, default_value_t = 0)]
pub device_id: usize,

/// Enable performance profiling
#[arg(long)]
pub profile: bool,

/// Disable contour drawing
#[arg(long)]
pub no_plot: bool,
pub no_contours: bool,

/// Show result
#[arg(long)]
pub no_contours: bool,
pub view: bool,

/// Do not save output
#[arg(long)]
pub nosave: bool,
}

fn main() -> Result<()> {
let args = Args::parse();

// build options
let options = Options::default();

// version & task
let (options, saveout) = match args.ver {
YOLOVersion::V5 => match args.task {
YOLOTask::Classify => (
options.with_model(&args.model.unwrap_or("yolo/v5-n-cls-dyn.onnx".to_string()))?,
"YOLOv5-Classify",
),
YOLOTask::Detect => (
options.with_model(&args.model.unwrap_or("yolo/v5-n-dyn.onnx".to_string()))?,
"YOLOv5-Detect",
),
YOLOTask::Segment => (
options.with_model(&args.model.unwrap_or("yolo/v5-n-seg-dyn.onnx".to_string()))?,
"YOLOv5-Segment",
),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.ver),
},
YOLOVersion::V6 => match args.task {
YOLOTask::Detect => (
options
.with_model(&args.model.unwrap_or("yolo/v6-n-dyn.onnx".to_string()))?
.with_nc(args.nc),
"YOLOv6-Detect",
),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.ver),
},
YOLOVersion::V7 => match args.task {
YOLOTask::Detect => (
options
.with_model(&args.model.unwrap_or("yolo/v7-tiny-dyn.onnx".to_string()))?
.with_nc(args.nc),
"YOLOv7-Detect",
),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.ver),
},
YOLOVersion::V8 => match args.task {
YOLOTask::Classify => (
options.with_model(&args.model.unwrap_or("yolo/v8-m-cls-dyn.onnx".to_string()))?,
"YOLOv8-Classify",
),
YOLOTask::Detect => (
options.with_model(&args.model.unwrap_or("yolo/v8-m-dyn.onnx".to_string()))?,
"YOLOv8-Detect",
),
YOLOTask::Segment => (
options.with_model(&args.model.unwrap_or("yolo/v8-m-seg-dyn.onnx".to_string()))?,
"YOLOv8-Segment",
),
YOLOTask::Pose => (
options.with_model(&args.model.unwrap_or("yolo/v8-m-pose-dyn.onnx".to_string()))?,
"YOLOv8-Pose",
),
YOLOTask::Obb => (
options.with_model(&args.model.unwrap_or("yolo/v8-m-obb-dyn.onnx".to_string()))?,
"YOLOv8-Obb",
),
},
YOLOVersion::V9 => match args.task {
YOLOTask::Detect => (
options.with_model(&args.model.unwrap_or("yolo/v9-c-dyn-f16.onnx".to_string()))?,
"YOLOv9-Detect",
),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.ver),
},
YOLOVersion::V10 => match args.task {
YOLOTask::Detect => (
options.with_model(&args.model.unwrap_or("yolo/v10-n.onnx".to_string()))?,
"YOLOv10-Detect",
),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.ver),
},
YOLOVersion::RTDETR => match args.task {
YOLOTask::Detect => (
options.with_model(&args.model.unwrap_or("yolo/rtdetr-l-f16.onnx".to_string()))?,
"RTDETR",
),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.ver),
},
};

let options = options
.with_yolo_version(args.ver)
.with_yolo_task(args.task);
// path
let path = args.model.unwrap_or({
format!(
"yolo/{}-{}-{}.onnx",
args.ver.name(),
args.scale.name(),
args.task.name()
)
});

// saveout
let saveout = format!(
"{}-{}-{}",
args.ver.name(),
args.scale.name(),
args.task.name()
);

// device
let options = if args.cuda {
options.with_cuda(args.device_id)
let device = if args.cuda {
Device::Cuda(args.device_id)
} else if args.trt {
let options = options.with_trt(args.device_id);
if args.half {
options.with_trt_fp16(true)
} else {
options
}
Device::Trt(args.device_id)
} else if args.coreml {
options.with_coreml(args.device_id)
Device::CoreML(args.device_id)
} else {
options.with_cpu()
Device::Cpu(args.device_id)
};
let options = options

// 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(&[0.2, 0.15]) // class_0: 0.4, others: 0.15
.with_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
.with_profile(args.profile);

// build model
let mut model = YOLO::new(options)?;

// build dataloader
Expand All @@ -194,16 +166,53 @@ fn main() -> Result<()> {
// build annotator
let annotator = Annotator::default()
.with_skeletons(&COCO_SKELETONS_16)
.with_bboxes_thickness(4)
.without_masks(true) // No masks plotting when doing segment task.
.with_saveout(saveout);
.with_bboxes_thickness(3)
.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
};

// run & annotate
for (xs, _paths) in dl {
// let ys = model.run(&xs)?; // way one
let ys = model.forward(&xs, args.profile)?; // way two
if !args.no_plot {
annotator.annotate(&xs, &ys);
let images_plotted = annotator.plot(&xs, &ys, !args.nosave)?;

// 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,
}
}
}

// finish video write
if !args.nosave {
if let Some(viewer) = &mut viewer {
viewer.finish_write()?;
}
}

Expand Down
Loading

0 comments on commit f4d505c

Please sign in to comment.