feat(ocr,det,slide): 重构配置解析流程，移除非必要的生命周期方法

- 优化规范化模型目录 - 重构 Ocr，Detector配置解析流程
2026-07-08 15:48:56 +08:00
parent 31271e80db
commit 0cf3d5fefb
19 changed files with 306 additions and 214 deletions
--- a/src/models/det/executor.rs
+++ b/src/models/det/executor.rs
@@ -0,0 +1,297 @@
+use crate::models::ocr::model_metadata::ModelMetadata;
+use crate::models::loader::{ModelLoader, ModelSession, ModelType};
+use anyhow::{Context, Result, anyhow};
+use image::{DynamicImage, GenericImageView, imageops::FilterType};
+use std::path::Path;
+use tract_onnx::prelude::tract_ndarray::{Array2, Array3, Array4, Axis, prelude::*, s};
+use tract_onnx::prelude::{Graph, RunnableModel, Tensor, TypedFact, TypedOp, tvec};
+
+const DEFAULT_DET_PATH: &'static str = "common_det.onnx";
+
+// 预设的提示信息常量
+use crate::error::MODEL_DOWNLOAD_HELP;
+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,
+}
+
+
+
+
+
+
+#[derive(Debug)]
+pub struct Detector<'a> {
+    pub(crate) session: &'a DetSession,
+    pub(crate) use_gpu: bool,
+    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)
+    }
+}