diff --git a/examples/yolov5/main.rs b/examples/yolov5/main.rs index 60eae5b..d1c3c8c 100644 --- a/examples/yolov5/main.rs +++ b/examples/yolov5/main.rs @@ -1,15 +1,14 @@ use usls::{ - models::{YOLOTask, YOLO}, + models::{YOLOTask, YOLOVersion, YOLO}, Annotator, DataLoader, Options, }; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_conf_independent(true) - .with_anchors_first(true) + .with_yolo_version(YOLOVersion::V5) + .with_model("../models/yolov5s-seg.onnx")? .with_yolo_task(YOLOTask::Segment) - .with_model("yolov5s-seg.onnx")? // .with_trt(0) // .with_fp16(true) .with_i00((1, 1, 4).into()) diff --git a/examples/yolov8/main.rs b/examples/yolov8/main.rs index 91c527f..171b712 100644 --- a/examples/yolov8/main.rs +++ b/examples/yolov8/main.rs @@ -3,11 +3,11 @@ use usls::{coco, models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - // .with_model("yolov8m-dyn.onnx")? + .with_model("yolov8m-dyn.onnx")? // .with_model("yolov8m-dyn-f16.onnx")? // .with_model("yolov8m-pose-dyn.onnx")? // .with_model("yolov8m-cls-dyn.onnx")? - .with_model("yolov8m-seg-dyn.onnx")? + // .with_model("yolov8m-seg-dyn.onnx")? // .with_model("yolov8m-obb-dyn.onnx")? // .with_model("yolov8m-oiv7-dyn.onnx")? // .with_trt(0) @@ -19,7 +19,6 @@ fn main() -> Result<(), Box> { .with_i03((224, 640, 800).into()) .with_confs(&[0.4, 0.15]) // class 0: 0.4, others: 0.15 .with_names2(&coco::KEYPOINTS_NAMES_17) - // .with_dry_run(10) .with_profile(false); let mut model = YOLO::new(options)?; diff --git a/examples/yolov9/main.rs b/examples/yolov9/main.rs index c82842e..efb172c 100644 --- a/examples/yolov9/main.rs +++ b/examples/yolov9/main.rs @@ -1,9 +1,13 @@ -use usls::{models::YOLO, Annotator, DataLoader, Options}; +use usls::{ + models::{YOLOVersion, YOLO}, + Annotator, DataLoader, Options, +}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("yolov9-c-dyn-f16.onnx")? + .with_model("../models/yolov9-c.onnx")? + .with_yolo_version(YOLOVersion::V9) .with_i00((1, 1, 4).into()) .with_i02((416, 640, 800).into()) .with_i03((416, 640, 800).into()) diff --git a/src/core/options.rs b/src/core/options.rs index 9faa92e..38cf86c 100644 --- a/src/core/options.rs +++ b/src/core/options.rs @@ -54,7 +54,7 @@ pub struct Options { pub tokenizer: Option, pub vocab: Option, pub names: Option>, // names - pub names2: Option>, // names2: could be keypoints names + pub names2: Option>, // names2 pub names3: Option>, // names3 pub min_width: Option, pub min_height: Option, diff --git a/src/models/yolo.rs b/src/models/yolo.rs index 9721ddb..77211d4 100644 --- a/src/models/yolo.rs +++ b/src/models/yolo.rs @@ -20,8 +20,11 @@ pub enum YOLOTask { #[derive(Debug, Copy, Clone, ValueEnum)] pub enum YOLOVersion { + V5, V8, + V9, V10, + Customized, } #[derive(Debug)] @@ -54,24 +57,43 @@ impl YOLO { engine.height().to_owned(), engine.width().to_owned(), ); - let task = match options.yolo_task { Some(task) => task, - None => match engine - .try_fetch("task") - .unwrap_or("detect".to_string()) - .as_str() - { - "classify" => YOLOTask::Classify, - "detect" => YOLOTask::Detect, - "pose" => YOLOTask::Pose, - "segment" => YOLOTask::Segment, - "obb" => YOLOTask::Obb, - x => todo!("Not supported: {x:?} "), + None => match engine.try_fetch("task") { + None => { + println!("No clear YOLO task specified, using default: Detect"); + YOLOTask::Detect + } + Some(x) => match x.as_str() { + "classify" => YOLOTask::Classify, + "detect" => YOLOTask::Detect, + "pose" => YOLOTask::Pose, + "segment" => YOLOTask::Segment, + "obb" => YOLOTask::Obb, + x => todo!("YOLO Task: {x:?} is not supported"), + }, }, }; + let version = match options.yolo_version { + None => { + println!("No clear YOLO version specified, using default: YOLOv8"); + YOLOVersion::V8 + } + Some(x) => x, + }; - let version = options.yolo_version.unwrap_or(YOLOVersion::V8); + // output format + let (anchors_first, conf_independent, apply_nms, apply_probs_softmax) = match version { + YOLOVersion::V5 => (true, true, true, true), + YOLOVersion::V8 | YOLOVersion::V9 => (false, false, true, false), + YOLOVersion::V10 => (true, false, false, false), + YOLOVersion::Customized => ( + options.anchors_first, + options.conf_independent, + options.apply_nms, + options.apply_probs_softmax, + ), + }; // try from custom class names, and then model metadata let mut names = options.names.or(Self::fetch_names(&engine)); @@ -122,7 +144,6 @@ impl YOLO { confs, kconfs, iou: options.iou, - apply_nms: options.apply_nms, nc, nk, nm, @@ -133,9 +154,10 @@ impl YOLO { version, names, names_kpt, - anchors_first: options.anchors_first, - conf_independent: options.conf_independent, - apply_probs_softmax: options.apply_probs_softmax, + anchors_first, + conf_independent, + apply_nms, + apply_probs_softmax, }) } @@ -154,11 +176,7 @@ impl YOLO { }; let xs_ = ops::normalize(xs_, 0., 255.); let ys = self.engine.run(&[xs_])?; - - match self.version { - YOLOVersion::V10 => self.postprocess_v10(ys, xs), - _ => self.postprocess(ys, xs), - } + self.postprocess(ys, xs) } pub fn postprocess(&self, xs: Vec>, xs0: &[DynamicImage]) -> Result> { @@ -178,7 +196,6 @@ impl YOLO { } else { preds.into_owned() }; - ys.push( Y::default().with_probs( Prob::default() @@ -194,7 +211,6 @@ impl YOLO { for pred in preds.axis_iter(if self.anchors_first { Axis(0) } else { Axis(1) }) { // xywhclsr - let xywh = pred.slice(s![0..CXYWH_OFFSET]); let clss = pred.slice(s![CXYWH_OFFSET..CXYWH_OFFSET + self.nc]); let radians = pred[pred.len() - 1]; let (id, &confidence) = clss @@ -206,7 +222,7 @@ impl YOLO { continue; } - // re-scale + let xywh = pred.slice(s![0..CXYWH_OFFSET]); let cx = xywh[0] / ratio; let cy = xywh[1] / ratio; let w = xywh[2] / ratio; @@ -242,41 +258,70 @@ impl YOLO { .axis_iter(if self.anchors_first { Axis(0) } else { Axis(1) }) .enumerate() { - let bbox = pred.slice(s![0..CXYWH_OFFSET]); - let (conf_, clss) = if self.conf_independent { - ( - pred[CXYWH_OFFSET], - pred.slice(s![CXYWH_OFFSET + 1..CXYWH_OFFSET + self.nc + 1]), - ) - } else { - (1.0, pred.slice(s![CXYWH_OFFSET..CXYWH_OFFSET + self.nc])) - }; - let (id, &confidence) = clss - .into_iter() - .enumerate() - .max_by(|a, b| a.1.total_cmp(b.1)) - .unwrap(); - let confidence = confidence * conf_; - if confidence < self.confs[id] { - continue; + match self.version { + YOLOVersion::V10 => { + let class_id = pred[CXYWH_OFFSET + 1] as usize; + let confidence = pred[CXYWH_OFFSET]; + if confidence < self.confs[class_id] { + continue; + } + let bbox = pred.slice(s![0..CXYWH_OFFSET]); + let x = bbox[0] / ratio; + let y = bbox[1] / ratio; + let x2 = bbox[2] / ratio; + let y2 = bbox[3] / ratio; + let w = x2 - x; + let h = y2 - y; + let y_bbox = Bbox::default() + .with_xywh(x, y, w, h) + .with_confidence(confidence) + .with_id(class_id as isize) + .with_id_born(i as isize) + .with_name( + self.names.as_ref().map(|names| names[class_id].to_owned()), + ); + y_bboxes.push(y_bbox); + } + _ => { + let (conf_, clss) = if self.conf_independent { + ( + pred[CXYWH_OFFSET], + pred.slice( + s![CXYWH_OFFSET + 1..CXYWH_OFFSET + self.nc + 1], + ), + ) + } else { + (1.0, pred.slice(s![CXYWH_OFFSET..CXYWH_OFFSET + self.nc])) + }; + let (id, &confidence) = clss + .into_iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .unwrap(); + let confidence = confidence * conf_; + if confidence < self.confs[id] { + continue; + } + let bbox = pred.slice(s![0..CXYWH_OFFSET]); + let cx = bbox[0] / ratio; + let cy = bbox[1] / ratio; + let w = bbox[2] / ratio; + let h = bbox[3] / ratio; + let x = cx - w / 2.; + let y = cy - h / 2.; + let x = x.max(0.0).min(image_width); + let y = y.max(0.0).min(image_height); + let y_bbox = Bbox::default() + .with_xywh(x, y, w, h) + .with_confidence(confidence) + .with_id(id as isize) + .with_id_born(i as isize) + .with_name( + self.names.as_ref().map(|names| names[id].to_owned()), + ); + y_bboxes.push(y_bbox); + } } - - // re-scale - let cx = bbox[0] / ratio; - let cy = bbox[1] / ratio; - let w = bbox[2] / ratio; - let h = bbox[3] / ratio; - let x = cx - w / 2.; - let y = cy - h / 2.; - let x = x.max(0.0).min(image_width); - let y = y.max(0.0).min(image_height); - let y_bbox = Bbox::default() - .with_xywh(x, y, w, h) - .with_confidence(confidence) - .with_id(id as isize) - .with_id_born(i as isize) - .with_name(self.names.as_ref().map(|names| names[id].to_owned())); - y_bboxes.push(y_bbox); } // nms @@ -425,59 +470,6 @@ impl YOLO { Ok(ys) } - pub fn postprocess_v10( - &self, - xs: Vec>, - xs0: &[DynamicImage], - ) -> Result> { - let mut ys = Vec::new(); - for (idx, preds) in xs[0].axis_iter(Axis(0)).enumerate() { - let image_width = xs0[idx].width() as f32; - let image_height = xs0[idx].height() as f32; - match self.task { - YOLOTask::Detect => { - let ratio = (self.width() as f32 / image_width) - .min(self.height() as f32 / image_height); - - let mut y_bboxes = vec![]; - for (i, pred) in preds.axis_iter(Axis(0)).enumerate() { - let confidence = pred[CXYWH_OFFSET]; - if confidence < self.confs[0] { - continue; - } - let class_id = pred[CXYWH_OFFSET + 1] as isize; - - let bbox = pred.slice(s![0..CXYWH_OFFSET]); - // re-scale - let x = bbox[0] / ratio; - let y = bbox[1] / ratio; - let x2 = bbox[2] / ratio; - let y2 = bbox[3] / ratio; - let w = x2 - x; - let h = y2 - y; - - let y_bbox = Bbox::default() - .with_xywh(x, y, w, h) - .with_confidence(confidence) - .with_id(class_id) - .with_id_born(i as isize) - .with_name( - self.names - .as_ref() - .map(|names| names[class_id as usize].to_owned()), - ); - y_bboxes.push(y_bbox); - } - let y = Y::default().with_bboxes(&y_bboxes); - ys.push(y); - } - _ => todo!("YOLO_V10 Not supported: {:?}", self.task), - } - } - - Ok(ys) - } - pub fn batch(&self) -> isize { self.batch.opt }