Skip to content

Commit

Permalink
Update
Browse files Browse the repository at this point in the history
  • Loading branch information
@jamjamjon
jamjamjon committed Jul 9, 2024
1 parent b155d44 commit b5f031a
Show file tree
Hide file tree
Showing 9 changed files with 871 additions and 321 deletions.
37 changes: 23 additions & 14 deletions examples/yolo/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
use anyhow::Result;
use clap::Parser;

use usls::{
coco, models::YOLO, Annotator, DataLoader, Options, Vision, YOLOFormat, YOLOTask, YOLOVersion,
};
use usls::{coco, models::YOLO, Annotator, DataLoader, Options, Vision, YOLOTask, YOLOVersion};

#[derive(Parser, Clone)]
#[command(author, version, about, long_about = None)]
Expand All @@ -16,9 +15,8 @@ pub struct Args {
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)]
pub version: YOLOVersion,

#[arg(long, value_enum, default_value_t = YOLOFormat::NCxcywhClssA)]
pub format: YOLOFormat,

// #[arg(long, value_enum, default_value_t = YOLOFormat::NCxcywhClssA)]
// pub format: YOLOFormat,
#[arg(long, default_value_t = 224)]
pub width_min: isize,

Expand All @@ -37,6 +35,9 @@ pub struct Args {
#[arg(long, default_value_t = 800)]
pub height_max: isize,

#[arg(long, default_value_t = 80)]
pub nc: usize,

#[arg(long)]
pub trt: bool,

Expand All @@ -59,7 +60,7 @@ pub struct Args {
pub plot: bool,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
fn main() -> Result<()> {
let args = Args::parse();

// build options
Expand All @@ -68,25 +69,33 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
// version & task
let options = match args.version {
YOLOVersion::V5 => match args.task {
YOLOTask::Classify => options.with_model("../models/yolov5s-cls.onnx")?,
YOLOTask::Detect => options.with_model("../models/yolov5s.onnx")?,
YOLOTask::Segment => options.with_model("../models/yolov5s.onnx")?,
t => todo!("{t:?} is unsupported for {:?}", args.version),
YOLOTask::Classify => options.with_model("yolov5n-cls-dyn.onnx")?,
YOLOTask::Detect => options.with_model("yolov5n-dyn.onnx")?,
YOLOTask::Segment => options.with_model("yolov5n-seg-dyn.onnx")?,
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.version),
},
YOLOVersion::V6 => match args.task {
YOLOTask::Detect => options.with_model("yolov6n-dyn.onnx")?.with_nc(args.nc),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.version),
},
YOLOVersion::V7 => match args.task {
YOLOTask::Detect => options.with_model("yolov7-tiny-dyn.onnx")?.with_nc(args.nc),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.version),
},
YOLOVersion::V8 => match args.task {
YOLOTask::Classify => options.with_model("yolov8m-cls-dyn-cls.onnx")?,
YOLOTask::Classify => options.with_model("yolov8m-cls-dyn.onnx")?,
YOLOTask::Detect => options.with_model("yolov8m-dyn.onnx")?,
YOLOTask::Segment => options.with_model("yolov8m-seg-dyn.onnx")?,
YOLOTask::Pose => options.with_model("yolov8m-pose-dyn.onnx")?,
YOLOTask::Obb => options.with_model("yolov8m-obb-dyn.onnx")?,
},
YOLOVersion::V9 => match args.task {
YOLOTask::Detect => options.with_model("yolov9-c-dyn-f16.onnx")?,
t => todo!("{t:?} is unsupported for {:?}", args.version),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.version),
},
YOLOVersion::V10 => match args.task {
YOLOTask::Detect => options.with_model("yolov10n-dyn.onnx")?,
t => todo!("{t:?} is unsupported for {:?}", args.version),
t => anyhow::bail!("Task: {t:?} is unsupported for {:?}", args.version),
},
}
.with_yolo_version(args.version)
Expand Down
30 changes: 8 additions & 22 deletions src/core/options.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use anyhow::Result;

use crate::{
auto_load,
models::{YOLOFormat, YOLOTask, YOLOVersion},
models::{YOLOPreds, YOLOTask, YOLOVersion},
Device, MinOptMax,
};

Expand Down Expand Up @@ -51,8 +51,7 @@ pub struct Options {
pub nm: Option<usize>,
pub confs: Vec<f32>,
pub kconfs: Vec<f32>,
pub iou: f32,
pub apply_nms: bool,
pub iou: Option<f32>,
pub tokenizer: Option<String>,
pub vocab: Option<String>,
pub names: Option<Vec<String>>, // names
Expand All @@ -61,10 +60,9 @@ pub struct Options {
pub min_width: Option<f32>,
pub min_height: Option<f32>,
pub unclip_ratio: f32, // DB
pub apply_probs_softmax: bool,
pub yolo_task: Option<YOLOTask>,
pub yolo_version: Option<YOLOVersion>,
pub yolo_format: Option<YOLOFormat>,
pub yolo_preds: Option<YOLOPreds>,
}

impl Default for Options {
Expand Down Expand Up @@ -106,8 +104,7 @@ impl Default for Options {
nm: None,
confs: vec![0.4f32],
kconfs: vec![0.5f32],
iou: 0.45f32,
apply_nms: true,
iou: None,
tokenizer: None,
vocab: None,
names: None,
Expand All @@ -118,8 +115,7 @@ impl Default for Options {
unclip_ratio: 1.5,
yolo_task: None,
yolo_version: None,
apply_probs_softmax: false,
yolo_format: None,
yolo_preds: None,
}
}
}
Expand Down Expand Up @@ -170,11 +166,6 @@ impl Options {
self
}

pub fn apply_probs_softmax(mut self, x: bool) -> Self {
self.apply_probs_softmax = x;
self
}

pub fn with_profile(mut self, profile: bool) -> Self {
self.profile = profile;
self
Expand Down Expand Up @@ -220,13 +211,8 @@ impl Options {
self
}

pub fn with_yolo_format(mut self, x: YOLOFormat) -> Self {
self.yolo_format = Some(x);
self
}

pub fn with_nms(mut self, apply_nms: bool) -> Self {
self.apply_nms = apply_nms;
pub fn with_yolo_preds(mut self, x: YOLOPreds) -> Self {
self.yolo_preds = Some(x);
self
}

Expand All @@ -241,7 +227,7 @@ impl Options {
}

pub fn with_iou(mut self, x: f32) -> Self {
self.iou = x;
self.iou = Some(x);
self
}

Expand Down
7 changes: 5 additions & 2 deletions src/models/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ mod rtdetr;
mod rtmo;
mod svtr;
mod yolo;
mod yolo_format;
mod yolo_;
mod yolop;

pub use blip::Blip;
Expand All @@ -23,5 +23,8 @@ pub use rtdetr::RTDETR;
pub use rtmo::RTMO;
pub use svtr::SVTR;
pub use yolo::YOLO;
pub use yolo_format::{BoxType, YOLOFormat, YOLOTask, YOLOVersion};
pub use yolo_::*;
// {
// AnchorsPosition, BoxType, ClssType, KptsType, YOLOFormat, YOLOPreds, YOLOTask, YOLOVersion,
// };
pub use yolop::YOLOPv2;
Loading

0 comments on commit b5f031a

Please sign in to comment.