Minor adjustments to YOLO (#17)

examples/yolov5/main.rs

@@ -1,15 +1,14 @@
 use usls::{
-    models::{YOLOTask, YOLO},
+    models::{YOLOTask, YOLOVersion, YOLO},
     Annotator, DataLoader, Options,
 };
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // build model
     let options = Options::default()
-        .with_conf_independent(true)
+        .with_yolo_version(YOLOVersion::V5)
-        .with_anchors_first(true)
+        .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())

examples/yolov8/main.rs

@@ -3,11 +3,11 @@ use usls::{coco, models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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<dyn std::error::Error>> {
         .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)?;
 

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<dyn std::error::Error>> {
     // 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())

src/core/options.rs

@@ -54,7 +54,7 @@ pub struct Options {
     pub tokenizer: Option<String>,
     pub vocab: Option<String>,
     pub names: Option<Vec<String>>,  // names
-    pub names2: Option<Vec<String>>, // names2: could be keypoints names
+    pub names2: Option<Vec<String>>, // names2
     pub names3: Option<Vec<String>>, // names3
     pub min_width: Option<f32>,
     pub min_height: Option<f32>,

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
+            None => match engine.try_fetch("task") {
-                .try_fetch("task")
+                None => {
-                .unwrap_or("detect".to_string())
+                    println!("No clear YOLO task specified, using default: Detect");
-                .as_str()
+                    YOLOTask::Detect
-            {
+                }
-                "classify" => YOLOTask::Classify,
+                Some(x) => match x.as_str() {
-                "detect" => YOLOTask::Detect,
+                    "classify" => YOLOTask::Classify,
-                "pose" => YOLOTask::Pose,
+                    "detect" => YOLOTask::Detect,
-                "segment" => YOLOTask::Segment,
+                    "pose" => YOLOTask::Pose,
-                "obb" => YOLOTask::Obb,
+                    "segment" => YOLOTask::Segment,
-                x => todo!("Not supported: {x:?} "),
+                    "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,
+            anchors_first,
-            conf_independent: options.conf_independent,
+            conf_independent,
-            apply_probs_softmax: options.apply_probs_softmax,
+            apply_nms,
+            apply_probs_softmax,
         })
     }
 
@@ -154,11 +176,7 @@ impl YOLO {
         };
         let xs_ = ops::normalize(xs_, 0., 255.);
         let ys = self.engine.run(&[xs_])?;
-
+        self.postprocess(ys, xs)
-        match self.version {
-            YOLOVersion::V10 => self.postprocess_v10(ys, xs),
-            _ => self.postprocess(ys, xs),
-        }
     }
 
     pub fn postprocess(&self, xs: Vec<Array<f32, IxDyn>>, xs0: &[DynamicImage]) -> Result<Vec<Y>> {
@@ -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]);
+                        match self.version {
-                        let (conf_, clss) = if self.conf_independent {
+                            YOLOVersion::V10 => {
-                            (
+                                let class_id = pred[CXYWH_OFFSET + 1] as usize;
-                                pred[CXYWH_OFFSET],
+                                let confidence = pred[CXYWH_OFFSET];
-                                pred.slice(s![CXYWH_OFFSET + 1..CXYWH_OFFSET + self.nc + 1]),
+                                if confidence < self.confs[class_id] {
-                            )
+                                    continue;
-                        } else {
+                                }
-                            (1.0, pred.slice(s![CXYWH_OFFSET..CXYWH_OFFSET + self.nc]))
+                                let bbox = pred.slice(s![0..CXYWH_OFFSET]);
-                        };
+                                let x = bbox[0] / ratio;
-                        let (id, &confidence) = clss
+                                let y = bbox[1] / ratio;
-                            .into_iter()
+                                let x2 = bbox[2] / ratio;
-                            .enumerate()
+                                let y2 = bbox[3] / ratio;
-                            .max_by(|a, b| a.1.total_cmp(b.1))
+                                let w = x2 - x;
-                            .unwrap();
+                                let h = y2 - y;
-                        let confidence = confidence * conf_;
+                                let y_bbox = Bbox::default()
-                        if confidence < self.confs[id] {
+                                    .with_xywh(x, y, w, h)
-                            continue;
+                                    .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<Array<f32, IxDyn>>,
-        xs0: &[DynamicImage],
-    ) -> Result<Vec<Y>> {
-        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
     }