mii/usls
1
0
Fork 0
mirror of https://github.com/mii443/usls.git synced 2025-12-03 11:08:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Minor adjustments to YOLO (#17)

Browse Source
This commit is contained in:
Jamjamjon
2024-06-02 12:49:42 +08:00
committed by GitHub
parent 20577e4242
commit 311d49f5b7
5 changed files with 116 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
examples/yolov5/main.rs
View File

@@ -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_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())

5
examples/yolov8/main.rs
View File

@@ -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)?;

8
examples/yolov9/main.rs
View File

@@ -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())

2
src/core/options.rs
View File

@@ -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>,

146
src/models/yolo.rs
View File

@@ -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()
{
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!("Not supported: {x:?} "),
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<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,11 +258,37 @@ impl YOLO {
.axis_iter(if self.anchors_first { Axis(0) } else { Axis(1) })
.enumerate()
{
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]),
pred.slice(
s![CXYWH_OFFSET + 1..CXYWH_OFFSET + self.nc + 1],
),
)
} else {
(1.0, pred.slice(s![CXYWH_OFFSET..CXYWH_OFFSET + self.nc]))
@@ -260,8 +302,7 @@ impl YOLO {
if confidence < self.confs[id] {
continue;
}
// re-scale
let bbox = pred.slice(s![0..CXYWH_OFFSET]);
let cx = bbox[0] / ratio;
let cy = bbox[1] / ratio;
let w = bbox[2] / ratio;
@@ -275,9 +316,13 @@ impl YOLO {
.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()));
.with_name(
self.names.as_ref().map(|names| names[id].to_owned()),
);
y_bboxes.push(y_bbox);
}
}
}
// nms
let mut y = Y::default().with_bboxes(&y_bboxes);
@@ -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
}