Refactor YOLO with outputs format

examples/yolov8/README.md → examples/yolo/README.md

@@ -1,5 +1,83 @@
 ## Quick Start
 
+
+
+## Options for all YOLOs
+
+```Rust
+// YOLOv5-Classify
+let options = options
+    .with_yolo_version(args.version)
+    .with_yolo_task(args.task)
+    // .with_yolo_format(YOLOFormat::NClss)
+    .with_model("../models/yolov5s-cls.onnx")?;
+
+// YOLOv5-Detect
+let options = options
+    .with_yolo_version(YOLOVersion::V5)
+    .with_yolo_task(YOLOTask::Detect)
+    // .with_yolo_format(YOLOFormat::NACxcywhConfClss)
+    .with_model("../models/yolov5s.onnx")?;
+
+// YOLOv5-Segment
+let options = options
+    .with_yolo_version(YOLOVersion::V5)
+    .with_yolo_task(YOLOTask::Segment)
+    // .with_yolo_format(YOLOFormat::NACxcywhConfClssCoefs)
+    .with_model("../models/yolov5s.onnx")?;
+
+// YOLOv8-Detect
+let options = options
+    .with_yolo_version(YOLOVersion::V8)
+    .with_yolo_task(YOLOTask::Detect)
+    // .with_yolo_format(YOLOFormat::NCxcywhClssA)
+    .with_model("yolov8m-dyn.onnx")?;
+
+// YOLOv8-Classify
+let options = options
+    .with_yolo_version(YOLOVersion::V8)
+    .with_yolo_task(YOLOTask::Classify)
+    // .with_yolo_format(YOLOFormat::NClss)
+    .with_model("yolov8m-cls-dyn.onnx")?;
+
+// YOLOv8-Pose
+let options = options
+    .with_yolo_version(YOLOVersion::V8)
+    .with_yolo_task(YOLOTask::Pose)
+    // .with_yolo_format(YOLOFormat::NCxcywhClssXycsA)
+    .with_model("yolov8m-pose-dyn.onnx")?;
+
+// YOLOv8-Segment
+let options = options
+    .with_yolo_version(YOLOVersion::V8)
+    .with_yolo_task(YOLOTask::Segment)
+    // .with_yolo_format(YOLOFormat::NCxcywhClssCoefsA)
+    .with_model("yolov8m-seg-dyn.onnx")?;
+
+// YOLOv8-Obb
+let options = options
+    .with_yolo_version(YOLOVersion::V8)
+    .with_yolo_task(YOLOTask::Obb)
+    // .with_yolo_format(YOLOFormat::NCxcywhClssRA)
+    .with_model("yolov8m-obb-dyn.onnx")?;
+
+// YOLOv9-Detect
+let options = options
+    .with_yolo_version(YOLOVersion::V9)
+    .with_yolo_task(YOLOTask::Detect)
+    // .with_yolo_format(YOLOFormat::NCxcywhClssA)
+    .with_model("yolov9-c-dyn-f16.onnx")?;
+
+// YOLOv10-Detect
+let options = options
+    .with_yolo_version(YOLOVersion::V10)
+    .with_yolo_task(YOLOTask::Detect)
+    // .with_yolo_format(YOLOFormat::NAXyxyConfCls) //
+    .with_model("yolov10n-dyn.onnx")?;
+```
+
+
+
 ```shell
 cargo run -r --example yolov8
 ```

examples/yolo/main.rs

@@ -0,0 +1,145 @@
+use clap::Parser;
+
+use usls::{
+    coco, models::YOLO, Annotator, DataLoader, Options, Vision, YOLOFormat, YOLOTask, YOLOVersion,
+};
+
+#[derive(Parser, Clone)]
+#[command(author, version, about, long_about = None)]
+pub struct Args {
+    #[arg(long, default_value_t = String::from("./assets/bus.jpg"))]
+    pub source: String,
+
+    #[arg(long, value_enum, default_value_t = YOLOTask::Detect)]
+    pub task: YOLOTask,
+
+    #[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, default_value_t = 224)]
+    pub width_min: isize,
+
+    #[arg(long, default_value_t = 640)]
+    pub width: isize,
+
+    #[arg(long, default_value_t = 800)]
+    pub width_max: isize,
+
+    #[arg(long, default_value_t = 224)]
+    pub height_min: isize,
+
+    #[arg(long, default_value_t = 640)]
+    pub height: isize,
+
+    #[arg(long, default_value_t = 800)]
+    pub height_max: isize,
+
+    #[arg(long)]
+    pub trt: bool,
+
+    #[arg(long)]
+    pub cuda: bool,
+
+    #[arg(long)]
+    pub half: bool,
+
+    #[arg(long)]
+    pub coreml: bool,
+
+    #[arg(long, default_value_t = 0)]
+    pub device_id: usize,
+
+    #[arg(long)]
+    pub profile: bool,
+
+    #[arg(long)]
+    pub plot: bool,
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+
+    // build options
+    let options = Options::default();
+
+    // 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),
+        },
+        YOLOVersion::V8 => match args.task {
+            YOLOTask::Classify => options.with_model("yolov8m-cls-dyn-cls.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),
+        },
+        YOLOVersion::V10 => match args.task {
+            YOLOTask::Detect => options.with_model("yolov10n-dyn.onnx")?,
+            t => todo!("{t:?} is unsupported for {:?}", args.version),
+        },
+    }
+    .with_yolo_version(args.version)
+    .with_yolo_task(args.task);
+
+    // device
+
+    let options = if args.cuda {
+        options.with_cuda(args.device_id)
+    } else if args.trt {
+        let options = options.with_trt(args.device_id);
+        if args.half {
+            options.with_fp16(true)
+        } else {
+            options
+        }
+    } else if args.coreml {
+        options.with_coreml(args.device_id)
+    } else {
+        options.with_cpu()
+    };
+    let options = options
+        .with_i00((1, 1, 4).into())
+        .with_i02((args.height_min, args.height, args.height_max).into())
+        .with_i03((args.width_min, args.width, args.width_max).into())
+        .with_confs(&[0.4, 0.15]) // class_0: 0.4, others: 0.15
+        .with_names2(&coco::KEYPOINTS_NAMES_17)
+        .with_profile(args.profile);
+    let mut model = YOLO::new(options)?;
+
+    // build dataloader
+    let dl = DataLoader::default()
+        .with_batch(model.batch() as _)
+        .load(args.source)?;
+
+    // build annotator
+    let annotator = Annotator::default()
+        .with_skeletons(&coco::SKELETONS_16)
+        .with_bboxes_thickness(7)
+        .without_masks(true) // No masks plotting.
+        .with_saveout("YOLO-Series");
+
+    // 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.plot {
+            annotator.annotate(&xs, &ys);
+        } else {
+            println!("{:?}", ys);
+        }
+    }
+
+    Ok(())
+}