HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

[Example] YOLO-Series(v5-11) ONNXRuntime Rust (#17311)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Jamjamjon 2024-11-02 20:06:07 +08:00 • committed by GitHub
parent d28caa9a58
commit f95dc37311
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 362 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

13
examples/YOLOv8-ONNXRuntime-Rust/src/lib.rs
View file

@ -118,16 +118,15 @@ pub fn check_font(font: &str) -> rusttype::Font<'static> {
rusttype::Font::try_from_vec(buffer).unwrap()
}
use ab_glyph::FontArc;
pub fn load_font() -> FontArc{
pub fn load_font() -> FontArc {
use std::path::Path;
let font_path = Path::new("./font/Arial.ttf");
match font_path.try_exists() {
Ok(true) => {
let buffer = std::fs::read(font_path).unwrap();
FontArc::try_from_vec(buffer).unwrap()
},
}
Ok(false) => {
std::fs::create_dir_all("./font").unwrap();
println!("Downloading font...");
@ -136,7 +135,7 @@ pub fn load_font() -> FontArc{
.timeout(std::time::Duration::from_secs(500))
.call()
.unwrap_or_else(|err| panic!("> Failed to download font: {source_url}: {err:?}"));
// read to buffer
let mut buffer = vec![];
let total_size = resp
@ -153,9 +152,9 @@ pub fn load_font() -> FontArc{
fd.write_all(&buffer).unwrap();
println!("Font saved at: {:?}", font_path.display());
FontArc::try_from_vec(buffer).unwrap()
},
}
Err(e) => {
panic!("Failed to load font {}", e);
},
}
}
}
}

6
examples/YOLOv8-ONNXRuntime-Rust/src/model.rs
View file

@ -8,7 +8,7 @@ use rand::{thread_rng, Rng};
use std::path::PathBuf;
use crate::{
load_font, gen_time_string, non_max_suppression, Args, Batch, Bbox, Embedding, OrtBackend,
gen_time_string, load_font, non_max_suppression, Args, Batch, Bbox, Embedding, OrtBackend,
OrtConfig, OrtEP, Point2, YOLOResult, YOLOTask, SKELETON,
};
@ -40,7 +40,7 @@ impl YOLOv8 {
OrtEP::CUDA(config.device_id)
} else {
OrtEP::CPU
};
};
// batch
let batch = Batch {
@ -463,7 +463,7 @@ impl YOLOv8 {
image::Rgb(self.color_palette[bbox.id()].into()),
bbox.xmin() as i32,
(bbox.ymin() - legend_size as f32) as i32,
legend_size as f32,
legend_size as f32,
&font,
&legend,
);