301 lines
11 KiB
Rust
301 lines
11 KiB
Rust
use anyhow::{Context, Result};
|
||
use image::{imageops::FilterType, DynamicImage, GenericImageView};
|
||
use std::fmt;
|
||
use tract_onnx::prelude::tract_ndarray::{prelude::*, s, Array2, Array3, Array4, Axis};
|
||
use tract_onnx::prelude::{Tensor};
|
||
|
||
|
||
use crate::models::det::session::DetSession;
|
||
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct DetectionResult {
|
||
pub x1: i32,
|
||
pub y1: i32,
|
||
pub x2: i32,
|
||
pub y2: i32,
|
||
pub score: f32,
|
||
pub class_id: u32,
|
||
}
|
||
|
||
impl fmt::Display for DetectionResult {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
// 结构体只管自己这一行怎么显示，不用管外部的索引 [i]
|
||
write!(
|
||
f,
|
||
"x1={}, y1={}, x2={}, y2={}, 分数={:.4}, 类别ID={}",
|
||
self.x1, self.y1, self.x2, self.y2, self.score, self.class_id
|
||
)
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub struct Detector<'a> {
|
||
pub(crate) session: &'a DetSession,
|
||
#[allow(dead_code)]
|
||
pub(crate) use_gpu: bool,
|
||
#[allow(dead_code)]
|
||
pub(crate) device_id: u8,
|
||
}
|
||
|
||
impl<'a> Detector<'a> {
|
||
pub fn new(session: &'a DetSession) -> Self {
|
||
Detector {
|
||
session,
|
||
use_gpu: false,
|
||
device_id: 0,
|
||
}
|
||
}
|
||
|
||
pub fn predict(&self, image: &DynamicImage) -> Result<Vec<DetectionResult>> {
|
||
// Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑
|
||
self.get_bbox(image)
|
||
}
|
||
/// 2. preproc: 纯 Rust 实现 (替代 OpenCV)
|
||
fn preproc(&self, image: &DynamicImage, input_size: (u32, u32)) -> Result<(Tensor, f32)> {
|
||
let (target_h, target_w) = input_size;
|
||
let (img_w, img_h) = image.dimensions();
|
||
|
||
// 计算缩放比例 (Letterbox)
|
||
let r = (target_h as f32 / img_h as f32).min(target_w as f32 / img_w as f32);
|
||
let new_h = (img_h as f32 * r) as u32;
|
||
let new_w = (img_w as f32 * r) as u32;
|
||
|
||
// Resize 图像
|
||
let resized = image.resize_exact(new_w, new_h, FilterType::Triangle);
|
||
// 2. 关键：将 DynamicImage 显式转换为 RgbImage (Rgb<u8>)
|
||
let resized_rgb = resized.to_rgb8();
|
||
// 创建 114 灰度填充的背景
|
||
let mut base_img =
|
||
image::ImageBuffer::from_pixel(target_w, target_h, image::Rgb([114u8, 114, 114]));
|
||
|
||
// 将 resize 后的图像覆盖到左上角 (类似于原始代码中的 padded_img[:h, :w])
|
||
image::imageops::overlay(&mut base_img, &resized_rgb, 0, 0);
|
||
|
||
// 优化：直接获取底层的扁平 raw buffer，比 enumerate_pixels() 快得多
|
||
let raw_samples = base_img.as_flat_samples();
|
||
let slice = raw_samples.as_slice();
|
||
|
||
// 构造 NCHW Tensor
|
||
let mut array = Array4::<f32>::zeros((1, 3, target_h as usize, target_w as usize));
|
||
|
||
// 用连续的 stride 步长进行写入，提高 CPU 缓存利用率
|
||
for y in 0..target_h as usize {
|
||
for x in 0..target_w as usize {
|
||
let idx = (y * target_w as usize + x) * 3;
|
||
// BGR 赋值
|
||
array[[0, 0, y, x]] = slice[idx + 2] as f32; // B
|
||
array[[0, 1, y, x]] = slice[idx + 1] as f32; // G
|
||
array[[0, 2, y, x]] = slice[idx] as f32; // R
|
||
}
|
||
}
|
||
|
||
Ok((array.into(), r))
|
||
}
|
||
|
||
/// 3. demo_postprocess (逻辑与 Python 一致)
|
||
fn demo_postprocess(&self, mut outputs: Array3<f32>, img_size: (i32, i32)) -> Array3<f32> {
|
||
let strides = [8, 16, 32];
|
||
|
||
// 遍历每一个 Batch（支持动态 Batch 推理）
|
||
for mut batch in outputs.axis_iter_mut(Axis(0)) {
|
||
let mut offset = 0;
|
||
|
||
for &stride in &strides {
|
||
// 计算当前特征图的尺寸
|
||
let h = img_size.0 / stride;
|
||
let w = img_size.1 / stride;
|
||
let f_stride = stride as f32;
|
||
|
||
for y in 0..h {
|
||
for x in 0..w {
|
||
// 计算当前格子在 25200 个锚点中的线性索引
|
||
let idx = offset + (y * w + x) as usize;
|
||
// 1. 还原中心点坐标 (cx, cy)
|
||
// 公式: (output + grid_offset) * stride
|
||
batch[[idx, 0]] = (batch[[idx, 0]] + x as f32) * f_stride;
|
||
batch[[idx, 1]] = (batch[[idx, 1]] + y as f32) * f_stride;
|
||
|
||
// 2. 还原宽高 (w, h)
|
||
// 公式: exp(output) * stride
|
||
batch[[idx, 2]] = batch[[idx, 2]].exp() * f_stride;
|
||
batch[[idx, 3]] = batch[[idx, 3]].exp() * f_stride;
|
||
}
|
||
}
|
||
// 移动到下一个步长的起始位置
|
||
offset += (h * w) as usize;
|
||
}
|
||
}
|
||
outputs
|
||
}
|
||
|
||
/// 4. nms
|
||
fn nms(&self, boxes: &Array2<f32>, scores: &Array1<f32>, nms_thr: f32) -> Vec<usize> {
|
||
let mut keep = Vec::new();
|
||
let x1 = boxes.column(0);
|
||
let y1 = boxes.column(1);
|
||
let x2 = boxes.column(2);
|
||
let y2 = boxes.column(3);
|
||
// 在每一项前加上 &，并确保括号内的计算顺序
|
||
// 注意：ndarray 的 View 运算需要 &view1 - &view2
|
||
let areas = (&x2 - &x1 + 1.0) * (&y2 - &y1 + 1.0);
|
||
|
||
// 初始排序索引
|
||
let mut v: Vec<usize> = (0..scores.len()).collect();
|
||
v.sort_unstable_by(|&i, &j| {
|
||
scores[j]
|
||
.partial_cmp(&scores[i])
|
||
.unwrap_or(std::cmp::Ordering::Equal)
|
||
});
|
||
// 我们不使用 v.remove(0)，而是直接通过索引池操作
|
||
let mut active_indices = v;
|
||
|
||
while !active_indices.is_empty() {
|
||
// 取出当前池子中得分最高的框（即第一个元素）
|
||
let i = active_indices[0];
|
||
keep.push(i);
|
||
|
||
// 如果池子里只剩一个了，直接结束
|
||
if active_indices.len() == 1 {
|
||
break;
|
||
}
|
||
|
||
// 5. 核心逻辑：使用 retain 一次性过滤掉：
|
||
// (a) 当前框自己 (idx == i)
|
||
// (b) 与当前框重叠度过高的框 (iou > nms_thr)
|
||
active_indices.retain(|&idx| {
|
||
// 如果是当前正在处理的框，不保留（因为它已经进入 keep 了）
|
||
if idx == i {
|
||
return false;
|
||
}
|
||
|
||
// 计算 IoU
|
||
let xx1 = x1[i].max(x1[idx]);
|
||
let yy1 = y1[i].max(y1[idx]);
|
||
let xx2 = x2[i].min(x2[idx]);
|
||
let yy2 = y2[i].min(y2[idx]);
|
||
|
||
let w = (xx2 - xx1 + 1.0).max(0.0);
|
||
let h = (yy2 - yy1 + 1.0).max(0.0);
|
||
let inter = w * h;
|
||
|
||
let iou = inter / (areas[i] + areas[idx] - inter);
|
||
|
||
// 只保留 IoU 小于阈值的框
|
||
iou <= nms_thr
|
||
});
|
||
}
|
||
|
||
keep
|
||
}
|
||
|
||
/// 5. multiclass_nms
|
||
//multiclass_nms_class_agnostic
|
||
pub fn multiclass_nms(
|
||
&self,
|
||
boxes: &Array2<f32>, // [25200, 4] -> xyxy 格式
|
||
scores: &Array2<f32>, // [25200, 80] -> 已经乘以 objectness 的得分
|
||
nms_thr: f32,
|
||
score_thr: f32,
|
||
) -> Vec<[f32; 6]> {
|
||
let mut candidates = Vec::new();
|
||
|
||
// 1. 筛选高分框 (单次遍历完成 Argmax 和 Threshold 过滤)
|
||
for i in 0..scores.nrows() {
|
||
let row = scores.row(i);
|
||
|
||
// 找到当前行（即当前锚点）得分最高的类别
|
||
let mut max_score = 0.0;
|
||
let mut cls_id = 0;
|
||
for (j, &s) in row.iter().enumerate() {
|
||
if s > max_score {
|
||
max_score = s;
|
||
cls_id = j;
|
||
}
|
||
}
|
||
|
||
// 仅保留超过阈值的候选框
|
||
if max_score > score_thr {
|
||
// 暂时存储索引和元数据，避免频繁创建大数组
|
||
candidates.push((i, max_score, cls_id));
|
||
}
|
||
}
|
||
|
||
if candidates.is_empty() {
|
||
return vec![];
|
||
}
|
||
|
||
// 2. 准备 NMS 输入
|
||
// 构造 NMS 需要的子集数组
|
||
let mut b_subset = Array2::<f32>::zeros((candidates.len(), 4));
|
||
let mut s_subset = Array1::<f32>::zeros(candidates.len());
|
||
|
||
for (new_idx, &(orig_idx, score, _)) in candidates.iter().enumerate() {
|
||
b_subset.row_mut(new_idx).assign(&boxes.row(orig_idx));
|
||
s_subset[new_idx] = score;
|
||
}
|
||
|
||
// 3. 执行 NMS (返回保留下来的子集索引)
|
||
let keep = self.nms(&b_subset, &s_subset, nms_thr);
|
||
|
||
// 4. 组装最终结果 [x1, y1, x2, y2, score, class_id]
|
||
keep.into_iter()
|
||
.map(|k_idx| {
|
||
let (orig_idx, score, cls_id) = candidates[k_idx];
|
||
let b = boxes.row(orig_idx);
|
||
[b[0], b[1], b[2], b[3], score, cls_id as f32]
|
||
})
|
||
.collect()
|
||
}
|
||
/// 6. get_bbox (完全解耦 OpenCV)
|
||
pub fn get_bbox(&self, dynamic_img: &DynamicImage) -> Result<Vec<DetectionResult>> {
|
||
// 使用 utils crate 解码
|
||
// let dynamic_img = image::load_from_memory(image_bytes).context("Failed to decode utils")?;
|
||
let (orig_w, orig_h) = dynamic_img.dimensions();
|
||
|
||
let (input_tensor, ratio) = self.preproc(dynamic_img, (416, 416))?;
|
||
|
||
// tract 推理
|
||
// let outputs = self.session.session.run(tvec!(input_tensor.into()))?;
|
||
let outputs = self.session.inference(input_tensor)?;
|
||
// let output_array = outputs[0]
|
||
let output_array = outputs
|
||
.to_array_view::<f32>()?
|
||
.to_owned()
|
||
.into_dimensionality::<Ix3>()?;
|
||
|
||
let predictions = self.demo_postprocess(output_array, (416, 416));
|
||
let pred = predictions.slice(s![0, .., ..]);
|
||
|
||
let boxes = pred.slice(s![.., 0..4]);
|
||
let obj_conf = pred.slice(s![.., 4..5]);
|
||
let cls_conf = pred.slice(s![.., 5..]);
|
||
let obj_broadcast = obj_conf
|
||
.broadcast(cls_conf.dim())
|
||
.context("ndarray broadcasting failed for scores calculation")?;
|
||
let scores = &obj_broadcast * &cls_conf;
|
||
// let scores = &pred.slice(s![.., 4..5]) * &pred.slice(s![.., 5..]);
|
||
|
||
let mut boxes_xyxy = Array2::<f32>::zeros(boxes.raw_dim());
|
||
for i in 0..boxes.nrows() {
|
||
boxes_xyxy[[i, 0]] = (boxes[[i, 0]] - boxes[[i, 2]] / 2.0) / ratio;
|
||
boxes_xyxy[[i, 1]] = (boxes[[i, 1]] - boxes[[i, 3]] / 2.0) / ratio;
|
||
boxes_xyxy[[i, 2]] = (boxes[[i, 0]] + boxes[[i, 2]] / 2.0) / ratio;
|
||
boxes_xyxy[[i, 3]] = (boxes[[i, 1]] + boxes[[i, 3]] / 2.0) / ratio;
|
||
}
|
||
|
||
let detections = self.multiclass_nms(&boxes_xyxy, &scores, 0.45, 0.1);
|
||
let final_results = detections
|
||
.into_iter()
|
||
.map(|d| DetectionResult {
|
||
x1: (d[0] as i32).max(0).min(orig_w as i32),
|
||
y1: (d[1] as i32).max(0).min(orig_h as i32),
|
||
x2: (d[2] as i32).max(0).min(orig_w as i32),
|
||
y2: (d[3] as i32).max(0).min(orig_h as i32),
|
||
score: d[4],
|
||
class_id: d[5] as u32,
|
||
})
|
||
.collect();
|
||
Ok(final_results)
|
||
}
|
||
}
|