1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
| use ort::{session::Session, GraphOptimizationLevel, value::Value, Tensor};
use image::{imageops::FilterType, GenericImageView, Pixel};
use ndarray::{s, Array, Array4, Axis};
use std::path::Path;
use std::time::Instant;
// ========== 配置 ==========
const MODEL_PATH: &str = "yolo26n.onnx";
const INPUT_SIZE: usize = 640;
const CONF_THRESH: f32 = 0.25;
const IOU_THRESH: f32 = 0.45;
#[derive(Debug, Clone)]
struct Detection {
x1: f32,
y1: f32,
x2: f32,
y2: f32,
confidence: f32,
class_id: usize,
class_name: &'static str,
}
// COCO 80类名称
const CLASS_NAMES: [&str; 80] = [
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
"umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
"kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
"bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
"chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop",
"mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
"refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
"toothbrush",
];
fn main() -> anyhow::Result<()> {
// 1. 创建会话
let session = Session::builder()?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.with_intra_threads(8)?
.commit_from_file(MODEL_PATH)?;
println!("✅ YOLO26 Rust 推理已就绪");
// 2. 加载并预处理图片
let img_path = "test.jpg";
let img = image::open(img_path)?;
let (orig_w, orig_h) = (img.width() as f32, img.height() as f32);
let start = Instant::now();
// 3. 预处理
let input = preprocess(&img);
// 4. 推理
let input_value = Tensor::from_array(input)?;
let outputs = session.run(ort::inputs!["images" => input_value]?)?;
// 5. 后处理
let output = outputs["output0"].try_extract_tensor::<f32>()?;
let detections = postprocess(output.view(), orig_w, orig_h);
let elapsed = start.elapsed();
// 6. 输出结果
println!("\n📊 检测完成，耗时: {:?}", elapsed);
println!("共检测到 {} 个目标:\n", detections.len());
for (i, det) in detections.iter().enumerate() {
println!(
"{:2}. {:<15} 置信度: {:.3} 位置: [{:.0}, {:.0}, {:.0}, {:.0}]",
i + 1, det.class_name, det.confidence, det.x1, det.y1, det.x2, det.y2
);
}
Ok(())
}
/// 图片预处理：Resize + 归一化 + NCHW格式
fn preprocess(img: &image::DynamicImage) -> Array4<f32> {
// Resize到640x640
let resized = img.resize_exact(
INPUT_SIZE as u32,
INPUT_SIZE as u32,
FilterType::CatmullRom,
);
// 创建NCHW格式数组 [1, 3, H, W]
let mut input = Array::zeros((1, 3, INPUT_SIZE, INPUT_SIZE));
for y in 0..INPUT_SIZE {
for x in 0..INPUT_SIZE {
let pixel = resized.get_pixel(x as u32, y as u32).to_rgb();
input[[0, 0, y, x]] = pixel[0] as f32 / 255.0;
input[[0, 1, y, x]] = pixel[1] as f32 / 255.0;
input[[0, 2, y, x]] = pixel[2] as f32 / 255.0;
}
}
input
}
/// 后处理：解析输出 + NMS
fn postprocess(
output: ndarray::ArrayView3<'_, f32>,
orig_w: f32,
orig_h: f32,
) -> Vec<Detection> {
let scale_x = orig_w / INPUT_SIZE as f32;
let scale_y = orig_h / INPUT_SIZE as f32;
let mut detections = Vec::new();
// 输出形状 [1, 84, 8400] -> 转置为 [8400, 84]
let output = output.permuted_axes((1, 2, 0)).remove_axis(Axis(2));
for i in 0..8400 {
let row = output.slice(s![i, ..]);
// 找最大置信度
let (class_id, confidence) = (4..84)
.map(|c| (c - 4, row[c]))
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
.unwrap();
if confidence < CONF_THRESH {
continue;
}
// 解析坐标 cx, cy, w, h
let cx = row[0] * scale_x;
let cy = row[1] * scale_y;
let w = row[2] * scale_x;
let h = row[3] * scale_y;
detections.push(Detection {
x1: cx - w / 2.0,
y1: cy - h / 2.0,
x2: cx + w / 2.0,
y2: cy + h / 2.0,
confidence,
class_id,
class_name: CLASS_NAMES[class_id],
});
}
// NMS非极大值抑制
nms(&mut detections, IOU_THRESH)
}
/// 非极大值抑制（Rust高效实现）
fn nms(detections: &mut Vec<Detection>, iou_thresh: f32) -> Vec<Detection> {
// 按置信度降序
detections.sort_by(|a, b| b.confidence.partial_cmp(&a.confidence).unwrap());
let mut keep = Vec::new();
let mut suppressed = vec![false; detections.len()];
for i in 0..detections.len() {
if suppressed[i] {
continue;
}
keep.push(detections[i].clone());
for j in (i + 1)..detections.len() {
if suppressed[j] {
continue;
}
if calculate_iou(&detections[i], &detections[j]) > iou_thresh {
suppressed[j] = true;
}
}
}
keep
}
fn calculate_iou(a: &Detection, b: &Detection) -> f32 {
let x1 = a.x1.max(b.x1);
let y1 = a.y1.max(b.y1);
let x2 = a.x2.min(b.x2);
let y2 = a.y2.min(b.y2);
if x2 <= x1 || y2 <= y1 {
return 0.0;
}
let intersection = (x2 - x1) * (y2 - y1);
let area_a = (a.x2 - a.x1) * (a.y2 - a.y1);
let area_b = (b.x2 - b.x1) * (b.y2 - b.y1);
intersection / (area_a + area_b - intersection)
}
|