CNWei/ddddocr-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(slide,det): 重构目标检测引擎并统一图像输入类型为 DynamicImage以及滑块匹配与比较引擎为 Rust 实现

Browse Source 
- 统一 `predict` 和 `get_bbox` 接口参数为 `&DynamicImage`，消除多步处理时的重复图像解码开销。
- 引入轻量级 `DetectionResult` 结构体和固定大小数组 `[f32; 6]` 替代旧的嵌套 `Vec`，彻底消除后处理中的内存碎片。
- 优化 `preproc` 预处理逻辑，使用连续内存切片批量操作替代原有的逐像素迭代遍历。
- 移除多余的 `multiclass_nms_class_agnostic` 转发层，合并并精简 NMS 聚合函数。
- 优化 `calculate_center` 几何中心点计算函数，提高泛型语义并复用于两种匹配模式
- 在执行核心算法前增加尺寸与通道边界守卫（Guard Clauses），提升库的防防御性编程能力与崩溃安全性
- 移除多余的错误二次包装（map_err），改由 Rust 原生 Result 错误传播机制直接向上层抛出
This commit is contained in:
CNWei
2026-07-03 17:51:28 +08:00
parent 22cc9709ad
commit 7f1ce04f50
7 changed files with 192 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/lib.rs
View File

@@ -11,6 +11,7 @@ use std::fmt::{Display, Formatter};
// 关键点：直接使用 tract 重导出的 ndarray // 关键点：直接使用 tract 重导出的 ndarray
use crate::charset::CharRestrict; use crate::charset::CharRestrict;
use crate::model_metadata::ModelMetadata; use crate::model_metadata::ModelMetadata;
use crate::models::det::DetectionResult;
use crate::utils::color_filter::{ColorPreset, HsvRange}; use crate::utils::color_filter::{ColorPreset, HsvRange};
use models::det::Det; use models::det::Det;
use models::loader::ModelSession; use models::loader::ModelSession;
@@ -114,7 +115,7 @@ impl DdddOcr {
Runtime::Ocr(s) => { Runtime::Ocr(s) => {
let res = s.predictor().probability(true).predict(img)?; let res = s.predictor().probability(true).predict(img)?;
println!("{}", res); println!("{}", res);
Ok("".to_string()) Ok(res.to_string())
} }
// Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img), // Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img),
// Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![ // Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![
@@ -122,17 +123,16 @@ impl DdddOcr {
// // 但上界的 H 通道写成了 240，超过了 180 的法定上限！ // // 但上界的 H 通道写成了 240，超过了 180 的法定上限！
// HsvRange::new((82, 221, 14), (240, 203, 82)), // HsvRange::new((82, 221, 14), (240, 203, 82)),
// ])).predict(img), // ])).predict(img),
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")),
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), }
} }
} pub fn detection(&self, img: &DynamicImage) -> Result<Vec<DetectionResult>> {
pub fn detection(&self, img: &[u8]) -> Result<Vec<Vec<i32>>> { match &self.runtime {
match &self.runtime { Runtime::Det(s) => s.predict(img),
Runtime::Det(s) => s.predict(img), Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")),
Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")), }
} }
} }
}
// struct Classification {} // struct Classification {}
// #[derive(Debug)] // #[derive(Debug)]

90
src/models/det.rs
View File

@@ -4,6 +4,20 @@ use image::{DynamicImage, GenericImageView, imageops::FilterType};
use tract_onnx::prelude::tract_ndarray::{Array2, Array3, Array4, Axis, prelude::*, s}; use tract_onnx::prelude::tract_ndarray::{Array2, Array3, Array4, Axis, prelude::*, s};
use tract_onnx::prelude::{Graph, RunnableModel, Tensor, TypedFact, TypedOp, tvec}; use tract_onnx::prelude::{Graph, RunnableModel, Tensor, TypedFact, TypedOp, tvec};
#[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,
}
pub struct Det { pub struct Det {
session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>, session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
} }
@@ -20,14 +34,14 @@ impl Det {
let session = ModelLoader::load_model(&model_path)?.session; let session = ModelLoader::load_model(&model_path)?.session;
Ok(Self { session }) Ok(Self { session })
} }
pub fn predict(&self, image_bytes: &[u8]) -> Result<Vec<Vec<i32>>> { pub fn predict(&self, image: &DynamicImage) -> Result<Vec<DetectionResult>> {
// Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑 // Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑
self.get_bbox(image_bytes) self.get_bbox(image)
} }
/// 2. preproc: 纯 Rust 实现 (替代 OpenCV) /// 2. preproc: 纯 Rust 实现 (替代 OpenCV)
fn preproc(&self, img: &DynamicImage, input_size: (u32, u32)) -> Result<(Tensor, f32)> { fn preproc(&self, image: &DynamicImage, input_size: (u32, u32)) -> Result<(Tensor, f32)> {
let (target_h, target_w) = input_size; let (target_h, target_w) = input_size;
let (img_w, img_h) = img.dimensions(); let (img_w, img_h) = image.dimensions();
// 计算缩放比例 (Letterbox) // 计算缩放比例 (Letterbox)
let r = (target_h as f32 / img_h as f32).min(target_w as f32 / img_w as f32); let r = (target_h as f32 / img_h as f32).min(target_w as f32 / img_w as f32);
@@ -35,7 +49,7 @@ impl Det {
let new_w = (img_w as f32 * r) as u32; let new_w = (img_w as f32 * r) as u32;
// Resize 图像 // Resize 图像
let resized = img.resize_exact(new_w, new_h, FilterType::Triangle); let resized = image.resize_exact(new_w, new_h, FilterType::Triangle);
// 2. 关键：将 DynamicImage 显式转换为 RgbImage (Rgb<u8>) // 2. 关键：将 DynamicImage 显式转换为 RgbImage (Rgb<u8>)
let resized_rgb = resized.to_rgb8(); let resized_rgb = resized.to_rgb8();
// 创建 114 灰度填充的背景 // 创建 114 灰度填充的背景
@@ -45,22 +59,25 @@ impl Det {
// 将 resize 后的图像覆盖到左上角 (类似于原始代码中的 padded_img[:h, :w]) // 将 resize 后的图像覆盖到左上角 (类似于原始代码中的 padded_img[:h, :w])
image::imageops::overlay(&mut base_img, &resized_rgb, 0, 0); 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 // 构造 NCHW Tensor
let mut array = Array4::<f32>::zeros((1, 3, target_h as usize, target_w as usize)); let mut array = Array4::<f32>::zeros((1, 3, target_h as usize, target_w as usize));
for (x, y, pixel) in base_img.enumerate_pixels() {
let x = x as usize; // 用连续的 stride 步长进行写入，提高 CPU 缓存利用率
let y = y as usize; for y in 0..target_h as usize {
// 核心对标 Python 的 BGR 逻辑： for x in 0..target_w as usize {
// pixel[0] 是 R, pixel[1] 是 G, pixel[2] 是 B let idx = (y * target_w as usize + x) * 3;
// 如果模型需要 BGR： // BGR 赋值
// array[[0, 0, y as usize, x as usize]] = pixel[0] as f32; array[[0, 0, y, x]] = slice[idx + 2] as f32; // B
// array[[0, 1, y as usize, x as usize]] = pixel[1] as f32; array[[0, 1, y, x]] = slice[idx + 1] as f32; // G
// array[[0, 2, y as usize, x as usize]] = pixel[2] as f32; array[[0, 2, y, x]] = slice[idx] as f32; // R
array[[0, 0, y, x]] = pixel[2] as f32; // B }
array[[0, 1, y, x]] = pixel[1] as f32; // G
array[[0, 2, y, x]] = pixel[0] as f32; // R
} }
Ok((array.into(), r)) Ok((array.into(), r))
} }
@@ -161,13 +178,14 @@ impl Det {
} }
/// 5. multiclass_nms /// 5. multiclass_nms
//multiclass_nms_class_agnostic
pub fn multiclass_nms( pub fn multiclass_nms(
&self, &self,
boxes: &Array2<f32>, // [25200, 4] -> xyxy 格式 boxes: &Array2<f32>, // [25200, 4] -> xyxy 格式
scores: &Array2<f32>, // [25200, 80] -> 已经乘以 objectness 的得分 scores: &Array2<f32>, // [25200, 80] -> 已经乘以 objectness 的得分
nms_thr: f32, nms_thr: f32,
score_thr: f32, score_thr: f32,
) -> Vec<Vec<f32>> { ) -> Vec<[f32; 6]> {
let mut candidates = Vec::new(); let mut candidates = Vec::new();
// 1. 筛选高分框 (单次遍历完成 Argmax 和 Threshold 过滤) // 1. 筛选高分框 (单次遍历完成 Argmax 和 Threshold 过滤)
@@ -213,17 +231,17 @@ impl Det {
.map(|k_idx| { .map(|k_idx| {
let (orig_idx, score, cls_id) = candidates[k_idx]; let (orig_idx, score, cls_id) = candidates[k_idx];
let b = boxes.row(orig_idx); let b = boxes.row(orig_idx);
vec![b[0], b[1], b[2], b[3], score, cls_id as f32] [b[0], b[1], b[2], b[3], score, cls_id as f32]
}) })
.collect() .collect()
} }
/// 6. get_bbox (完全解耦 OpenCV) /// 6. get_bbox (完全解耦 OpenCV)
pub fn get_bbox(&self, image_bytes: &[u8]) -> Result<Vec<Vec<i32>>> { pub fn get_bbox(&self, dynamic_img: &DynamicImage) -> Result<Vec<DetectionResult>> {
// 使用 utils crate 解码 // 使用 utils crate 解码
let dynamic_img = image::load_from_memory(image_bytes).context("Failed to decode utils")?; // let dynamic_img = image::load_from_memory(image_bytes).context("Failed to decode utils")?;
let (orig_w, orig_h) = dynamic_img.dimensions(); let (orig_w, orig_h) = dynamic_img.dimensions();
let (input_tensor, ratio) = self.preproc(&dynamic_img, (416, 416))?; let (input_tensor, ratio) = self.preproc(dynamic_img, (416, 416))?;
// tract 推理 // tract 推理
let outputs = self.session.run(tvec!(input_tensor.into()))?; let outputs = self.session.run(tvec!(input_tensor.into()))?;
@@ -236,7 +254,13 @@ impl Det {
let pred = predictions.slice(s![0, .., ..]); let pred = predictions.slice(s![0, .., ..]);
let boxes = pred.slice(s![.., 0..4]); let boxes = pred.slice(s![.., 0..4]);
let scores = &pred.slice(s![.., 4..5]) * &pred.slice(s![.., 5..]); 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()); let mut boxes_xyxy = Array2::<f32>::zeros(boxes.raw_dim());
for i in 0..boxes.nrows() { for i in 0..boxes.nrows() {
@@ -247,17 +271,19 @@ impl Det {
} }
let detections = self.multiclass_nms(&boxes_xyxy, &scores, 0.45, 0.1); let detections = self.multiclass_nms(&boxes_xyxy, &scores, 0.45, 0.1);
let final_results = detections
Ok(detections
.into_iter() .into_iter()
.map(|d| { .map(|d| {
vec![ DetectionResult{
(d[0] as i32).max(0).min(orig_w as i32), x1: (d[0] as i32).max(0).min(orig_w as i32),
(d[1] as i32).max(0).min(orig_h as i32), y1: (d[1] as i32).max(0).min(orig_h as i32),
(d[2] as i32).max(0).min(orig_w as i32), x2: (d[2] as i32).max(0).min(orig_w as i32),
(d[3] as i32).max(0).min(orig_h as i32), y2: (d[3] as i32).max(0).min(orig_h as i32),
] score: d[4],
class_id: d[5] as u32,
}
}) })
.collect()) .collect();
Ok(final_results )
} }
} }

67
src/models/ocr.rs
View File

@@ -21,7 +21,7 @@ use crate::utils::cv_ops::rgb_to_opencv_hsv;
/// 推理最终输出的强类型外壳（完全 Owned，无任何生命周期，可直接转 JSON） /// 推理最终输出的强类型外壳（完全 Owned，无任何生命周期，可直接转 JSON）
#[derive(Debug, Clone, Serialize)] #[derive(Debug, Clone, Serialize)]
pub enum OcrOutput { pub enum OcrResult {
/// 纯文本分支（对应 probability = false） /// 纯文本分支（对应 probability = false）
Text(String), Text(String),
/// 包含全量概率的分支（对应 probability = true） /// 包含全量概率的分支（对应 probability = true）
@@ -35,27 +35,31 @@ pub enum OcrOutput {
/// 不支持的模型或未知输出 /// 不支持的模型或未知输出
Unsupported { message: String }, Unsupported { message: String },
} }
impl OcrOutput { impl OcrResult {
/// 消费自身，直接提取最终文本 /// 消费自身，直接提取最终文本
pub fn into_text(self) -> String { pub fn into_text(self) -> String {
match self { match self {
OcrOutput::Text(text) => text, OcrResult::Text(text) => text,
OcrOutput::Probability { text, .. } => text, OcrResult::Probability { text, .. } => text,
OcrOutput::Unsupported { message } => { OcrResult::Unsupported { message } => {
// 作为库，这里可以返回空，或者直接携带错误信息，取决于你的设计 // 作为库，这里可以返回空，或者直接携带错误信息，取决于你的设计
format!("Error: {}", message) format!("Error: {}", message)
} }
} }
} }
} }
impl fmt::Display for OcrOutput { impl fmt::Display for OcrResult {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self { match self {
OcrOutput::Text(text) => { OcrResult::Text(text) => {
// 纯文本分支，直接输出文本内容 // 纯文本分支，直接输出文本内容
write!(f, "{}", text) write!(f, "{}", text)
} }
OcrOutput::Probability { text,probabilities, confidence } => { OcrResult::Probability {
text,
probabilities,
confidence,
} => {
// 概率分支，友好地展示文本以及百分比形式的置信度 // 概率分支，友好地展示文本以及百分比形式的置信度
// 1. 基本信息 // 1. 基本信息
write!(f, "{} (置信度: {:.2}%)", text, confidence * 100.0)?; write!(f, "{} (置信度: {:.2}%)", text, confidence * 100.0)?;
@@ -92,7 +96,7 @@ impl fmt::Display for OcrOutput {
} }
write!(f, "]") write!(f, "]")
} }
OcrOutput::Unsupported { message } => { OcrResult::Unsupported { message } => {
// 错误分支，直观输出异常原因 // 错误分支，直观输出异常原因
write!(f, "未识别成功: {}", message) write!(f, "未识别成功: {}", message)
} }
@@ -106,7 +110,7 @@ pub struct Ocr {
} }
impl ModelSession for Ocr { impl ModelSession for Ocr {
fn get_model_type(&self) -> ModelType { fn get_model_type(&self) -> ModelType {
todo!() todo!("使用thiserror作为错误处理的库,thiserror 专门用于开发库（Library）");
} }
fn desc(&self) -> String { fn desc(&self) -> String {
"Ocr Model 加载成功".to_string() "Ocr Model 加载成功".to_string()
@@ -114,6 +118,7 @@ impl ModelSession for Ocr {
} }
impl Ocr { impl Ocr {
pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error> { pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error> {
let session = ModelLoader::load_model(&model_path)?.session; let session = ModelLoader::load_model(&model_path)?.session;
Ok(Self { Ok(Self {
session, session,
@@ -142,7 +147,6 @@ pub struct OcrPredictor<'a> {
/// 是否修复PNG格式问题 /// 是否修复PNG格式问题
png_fix: bool, png_fix: bool,
/// 是否返回概率信息 /// 是否返回概率信息
#[allow(dead_code)]
probability: bool, probability: bool,
/// 颜色过滤：保留的颜色列表 /// 颜色过滤：保留的颜色列表
color_filter: Result<Option<Vec<HsvRange>>, String>, color_filter: Result<Option<Vec<HsvRange>>, String>,
@@ -189,7 +193,7 @@ impl<'a> OcrPredictor<'a> {
} }
} }
impl<'a> OcrPredictor<'a> { impl<'a> OcrPredictor<'a> {
pub fn predict(self, image: &DynamicImage) -> anyhow::Result<OcrOutput> { pub fn predict(self, image: &DynamicImage) -> anyhow::Result<OcrResult> {
println!("当前颜色过滤器状态: {:?}", self.color_filter); println!("当前颜色过滤器状态: {:?}", self.color_filter);
// ===================================================================== // =====================================================================
// 管道节点 1: 颜色过滤流水线 // 管道节点 1: 颜色过滤流水线
@@ -217,11 +221,11 @@ impl<'a> OcrPredictor<'a> {
let raw_tensor = self.ocr.inference(tensor)?; let raw_tensor = self.ocr.inference(tensor)?;
// 3. 后处理分流：直接返回 OcrOutput // 3. 后处理分流：直接返回 OcrResult
let ocr_output = match raw_tensor.datum_type() { let ocr_output = match raw_tensor.datum_type() {
DatumType::I64 => self.extract_from_i64_tensor(raw_tensor)?, DatumType::I64 => self.process_i64_tensor(raw_tensor)?,
DatumType::F32 => self.process_f32_pipeline(raw_tensor)?, DatumType::F32 => self.process_f32_tensor(raw_tensor)?,
_ => OcrOutput::Unsupported { _ => OcrResult::Unsupported {
message: format!("不支持的模型输出数据类型: {:?}", raw_tensor.datum_type()), message: format!("不支持的模型输出数据类型: {:?}", raw_tensor.datum_type()),
}, },
}; };
@@ -325,17 +329,16 @@ impl<'a> OcrPredictor<'a> {
} }
} }
impl<'a> OcrPredictor<'a> { impl<'a> OcrPredictor<'a> {
// pub fn get_valid_indices(&self) -> HashSet<usize> { fn is_valid_indices(&self, idx: usize) -> bool {
// match &self.charset_restrict { if idx >= self.ocr.model_metadata.charset.size() {
// Some(indices) => indices.iter().cloned().collect(), return false;
// // 如果是 None，现场映射出全量索引集给外部 }
// None => (0..self.ocr.model_metadata.charset.tokens.len()).collect(),
// }
// }
// compute_valid_indices
// fn valid_indices(&self) -> (bool, HashSet<usize>) {
// let charset = &self.ocr.model_metadata.charset;
match &self.charset_restrict {
Some(v) => v.binary_search(&idx).is_ok(),
None => true,
}
}
/// 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 /// 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印
/// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！ /// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！
pub fn valid_tokens(&self) -> Vec<&str> { pub fn valid_tokens(&self) -> Vec<&str> {
@@ -410,7 +413,7 @@ impl<'a> OcrPredictor<'a> {
(probabilities_list, confidence, predicted_indices) (probabilities_list, confidence, predicted_indices)
} }
/// 变体 A 专属提取器：直接从 I64 Tensor 零拷贝提取 CTC 文本与初始概率包 /// 变体 A 专属提取器：直接从 I64 Tensor 零拷贝提取 CTC 文本与初始概率包
fn extract_from_i64_tensor(&self, raw_tensor: Tensor) -> anyhow::Result<OcrOutput> { fn process_i64_tensor(&self, raw_tensor: Tensor) -> anyhow::Result<OcrResult> {
// 1. 拿到底层的动态维度只读视图 // 1. 拿到底层的动态维度只读视图
let view = raw_tensor.to_array_view::<i64>()?; let view = raw_tensor.to_array_view::<i64>()?;
@@ -424,17 +427,17 @@ impl<'a> OcrPredictor<'a> {
// 4. 组装返回 // 4. 组装返回
if self.probability { if self.probability {
Ok(OcrOutput::Probability { Ok(OcrResult::Probability {
text: final_text, text: final_text,
probabilities: vec![], // I64 模型物理上丢失了全量 Logits 分值网，降级处理 probabilities: vec![], // I64 模型物理上丢失了全量 Logits 分值网，降级处理
confidence: 1.0, // 判定即百分之百置信 confidence: 1.0, // 判定即百分之百置信
}) })
} else { } else {
Ok(OcrOutput::Text(final_text)) Ok(OcrResult::Text(final_text))
} }
} }
/// 变体二（F32）的总体管线：负责降维，并分流文本和概率 /// 变体二（F32）的总体管线：负责降维，并分流文本和概率
fn process_f32_pipeline(&self, raw_tensor: Tensor) -> anyhow::Result<OcrOutput> { fn process_f32_tensor(&self, raw_tensor: Tensor) -> anyhow::Result<OcrResult> {
let shape = raw_tensor.shape(); let shape = raw_tensor.shape();
println!("模型输出shape数据: {:?}", shape); println!("模型输出shape数据: {:?}", shape);
let view = raw_tensor.to_array_view::<f32>()?; let view = raw_tensor.to_array_view::<f32>()?;
@@ -476,7 +479,7 @@ impl<'a> OcrPredictor<'a> {
// 5. 执行 CTC 解码 // 5. 执行 CTC 解码
let final_text = self.ctc_decode_to_string(&predicted_indices); let final_text = self.ctc_decode_to_string(&predicted_indices);
Ok(OcrOutput::Probability { Ok(OcrResult::Probability {
text: final_text, text: final_text,
probabilities: probabilities_list, probabilities: probabilities_list,
confidence: confidence as f64, confidence: confidence as f64,
@@ -495,7 +498,7 @@ impl<'a> OcrPredictor<'a> {
.collect(); .collect();
let final_text = self.ctc_decode_to_string(&predicted_indices); let final_text = self.ctc_decode_to_string(&predicted_indices);
Ok(OcrOutput::Text(final_text)) Ok(OcrResult::Text(final_text))
} }
} }
/// 获取有效字符索引列表 (用于外部验证或过滤) /// 获取有效字符索引列表 (用于外部验证或过滤)

52
src/models/slide.rs
View File

@@ -27,7 +27,7 @@ impl Slide {
Self Self
} }
/// 对应 Python: slide_match /// 对应 Python: slide_match 滑块匹配接口
pub fn slide_match( pub fn slide_match(
&self, &self,
target_image: &DynamicImage, target_image: &DynamicImage,
@@ -38,9 +38,8 @@ impl Slide {
let background_array = image_to_ndarray(background_image); let background_array = image_to_ndarray(background_image);
self.perform_slide_match(target_array.view(), background_array.view(), simple_target) self.perform_slide_match(target_array.view(), background_array.view(), simple_target)
.map_err(|e| anyhow!("滑块匹配失败: {}", e))
} }
/// 对应 Python: slide_comparison /// 对应 Python: slide_comparison 差异比较接口
/// 用于比较带坑位的图片与原始背景图，定位差异点 /// 用于比较带坑位的图片与原始背景图，定位差异点
pub fn slide_comparison( pub fn slide_comparison(
&self, &self,
@@ -53,7 +52,6 @@ impl Slide {
// 2. 执行比较逻辑 (对应 _perform_slide_comparison) // 2. 执行比较逻辑 (对应 _perform_slide_comparison)
self.perform_slide_comparison(target_array.view(), background_array.view()) self.perform_slide_comparison(target_array.view(), background_array.view())
.map_err(|e| anyhow!("滑块比较执行失败: {}", e))
} }
/// 对应 Python: _perform_slide_comparison /// 对应 Python: _perform_slide_comparison
pub fn perform_slide_comparison( pub fn perform_slide_comparison(
@@ -61,7 +59,7 @@ impl Slide {
target: ArrayView3<u8>, target: ArrayView3<u8>,
background: ArrayView3<u8>, background: ArrayView3<u8>,
) -> Result<SlideResult> { ) -> Result<SlideResult> {
let (h, w, _) = target.dim(); // let (h, w, _) = target.dim();
// 1. 计算图像差异并灰度化 (对应 cv2.absdiff + cv2.cvtColor) // 1. 计算图像差异并灰度化 (对应 cv2.absdiff + cv2.cvtColor)
// 使用 OpenCV 标准权重公式：0.299R + 0.587G + 0.114B // 使用 OpenCV 标准权重公式：0.299R + 0.587G + 0.114B
@@ -77,6 +75,26 @@ impl Slide {
// } // }
// } // }
// 1. 计算差异数组 (复用 cv2::absdiff) // 1. 计算差异数组 (复用 cv2::absdiff)
let (th, tw, tc) = target.dim();
let (bh, bw, bc) = background.dim();
// 1. 比较模式下的严格尺寸校验
if th != bh || tw != bw || tc != bc {
return Err(anyhow!(
"比较模式要求两张图分辨率与通道数完全一致！Target: [{}x{}x{}], Background: [{}x{}x{}]",
tw,
th,
tc,
bw,
bh,
bc
));
}
if th == 0 || tw == 0 {
return Err(anyhow!("输入图像尺寸不能为0"));
}
let diff_array = abs_diff(&target, &background); let diff_array = abs_diff(&target, &background);
// 2. 转换为灰度数组 (复用你的 cv2.cvtColor) // 2. 转换为灰度数组 (复用你的 cv2.cvtColor)
@@ -130,6 +148,30 @@ impl Slide {
background: ArrayView3<u8>, background: ArrayView3<u8>,
simple_target: bool, // 增加这个参数 simple_target: bool, // 增加这个参数
) -> Result<SlideResult> { ) -> Result<SlideResult> {
let (th, tw, tc) = target.dim();
let (bh, bw, bc) = background.dim();
// 1. 严格的鲁棒性校验（防止底层的 imageproc 算子崩溃）
if th == 0 || tw == 0 || bh == 0 || bw == 0 {
return Err(anyhow!("输入图像的宽度或高度不能为0"));
}
if th > bh || tw > bw {
return Err(anyhow!(
"尺寸不匹配：滑块模板(target)尺寸 [{}x{}] 不能大于背景图(background) [{}x{}]",
tw,
th,
bw,
bh
));
}
if tc != bc {
return Err(anyhow!(
"目标图与背景图的通道数不一致 (target: {}, bg: {})",
tc,
bc
));
}
// 1. 统一灰度化 // 1. 统一灰度化
let target_gray = rgb_to_gray(target); let target_gray = rgb_to_gray(target);
let background_gray = rgb_to_gray(background); let background_gray = rgb_to_gray(background);

6
src/utils/cv_ops.rs
View File

@@ -95,9 +95,9 @@ pub fn bounding_rect(
let h = max_y - min_y; let h = max_y - min_y;
(min_x, min_y, w, h) (min_x, min_y, w, h)
} }
pub fn calculate_center(max_loc: (u32, u32), tw: usize, th: usize) -> (i32, i32) { pub fn calculate_center(top_left: (u32, u32), width: usize, height: usize) -> (i32, i32) {
let center_x = max_loc.0 as i32 + (tw as i32 / 2); let center_x = top_left.0 as i32 + (width as i32 / 2);
let center_y = max_loc.1 as i32 + (th as i32 / 2); let center_y = top_left.1 as i32 + (height as i32 / 2);
(center_x, center_y) (center_x, center_y)
} }
pub fn ndarray_to_luma8(array: ArrayView2<u8>) -> ImageBuffer<Luma<u8>, Vec<u8>> { pub fn ndarray_to_luma8(array: ArrayView2<u8>) -> ImageBuffer<Luma<u8>, Vec<u8>> {

36
src/utils/image_processor.rs
View File

@@ -16,28 +16,22 @@ pub fn resize_image(
target_height: u32, target_height: u32,
// resample 参数我们直接使用 FilterType，Lanczos3 是最接近 Python LANCZOS 的 // resample 参数我们直接使用 FilterType，Lanczos3 是最接近 Python LANCZOS 的
) -> DynamicImage { ) -> DynamicImage {
// 使用 resize 算法进行精确缩放
// image::imageops::resize(
// image,
// target_width,
// target_height,
// FilterType::Lanczos3
// )
// image::imageops::resize 的最高层封装 // image::imageops::resize 的最高层封装
// FilterType::Lanczos3 与 Python Pillow 的 Image.LANCZOS 算法完全对齐，缩放质量最高 // FilterType::Lanczos3 与 Python Pillow 的 Image.LANCZOS 算法完全对齐，缩放质量最高
image.resize_exact(target_width, target_height, FilterType::Lanczos3) image.resize_exact(target_width, target_height, FilterType::Lanczos3)
} }
pub fn resize_image1( // pub fn resize_image(
image: &GrayImage, // image: &GrayImage,
target_width: u32, // target_width: u32,
target_height: u32, // target_height: u32,
// resample 参数我们直接使用 FilterType，Lanczos3 是最接近 Python LANCZOS 的 // // resample 参数我们直接使用 FilterType，Lanczos3 是最接近 Python LANCZOS 的
) -> GrayImage { // ) -> GrayImage {
// 使用 resize 算法进行精确缩放 // // 使用 resize 算法进行精确缩放
image::imageops::resize( // image::imageops::resize(
image, // image,
target_width, // target_width,
target_height, // target_height,
FilterType::Lanczos3 // FilterType::Lanczos3
) // )
} // }

40
tests/ocr_test.rs
View File

@@ -1,8 +1,10 @@
use ddddocr_rs::models::slide::Slide; use ddddocr_rs::models::slide::Slide;
use ddddocr_rs::{DdddOcr, DdddOcrBuilder}; // 假设你的包名是这个 use ddddocr_rs::{DdddOcr, DdddOcrBuilder}; // 假设你的包名是这个
use image::Rgb; use image::{DynamicImage, Rgb};
use std::fs; use std::fs;
use std::path::Path; use std::path::Path;
use ddddocr_rs::models::det::DetectionResult;
fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> { fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
// 1. 先将泛型转为具体的 &Path 引用 // 1. 先将泛型转为具体的 &Path 引用
let path_ref = path.as_ref(); let path_ref = path.as_ref();
@@ -15,27 +17,26 @@ fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
} }
/// 将检测结果绘制在图像上并保存 /// 将检测结果绘制在图像上并保存
fn save_debug_image( fn save_debug_image(
image_bytes: &[u8], dynamic_img: &DynamicImage, // 【优化点 1】直接传入解码好的引用，拒绝重复解码
bboxes: &Vec<Vec<i32>>, bboxes: &[DetectionResult], // 【修改点 1】类型改为自定义结构体切片
output_path: &str, output_path: &str,
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
let dynamic_img = image::load_from_memory(image_bytes)?; // 删除了原本的 let dynamic_img = image::load_from_memory(image_bytes)?;
let mut img = dynamic_img.to_rgb8(); let mut img = dynamic_img.to_rgb8();
let (width, height) = img.dimensions(); let (width, height) = img.dimensions();
let red = Rgb([255u8, 0, 0]); let red = Rgb([255u8, 0, 0]);
for bbox in bboxes { for bbox in bboxes {
// 基础边界检查 // 【修改点 2】将原来的索引 bbox[0].. 改为结构体字段访问 .x1, .y1 ..
let x1 = bbox[0].max(0).min(width as i32 - 1) as u32; let x1 = bbox.x1.max(0).min(width as i32 - 1) as u32;
let y1 = bbox[1].max(0).min(height as i32 - 1) as u32; let y1 = bbox.y1.max(0).min(height as i32 - 1) as u32;
let x2 = bbox[2].max(0).min(width as i32 - 1) as u32; let x2 = bbox.x2.max(0).min(width as i32 - 1) as u32;
let y2 = bbox[3].max(0).min(height as i32 - 1) as u32; let y2 = bbox.y2.max(0).min(height as i32 - 1) as u32;
// 绘制横向线条 // 绘制横向线条
for x in x1..=x2 { for x in x1..=x2 {
img.put_pixel(x, y1, red); img.put_pixel(x, y1, red);
img.put_pixel(x, y2, red); img.put_pixel(x, y2, red);
// 如果要加粗，多画一行
if y1 + 1 < height { if y1 + 1 < height {
img.put_pixel(x, y1 + 1, red); img.put_pixel(x, y1 + 1, red);
} }
@@ -47,7 +48,6 @@ fn save_debug_image(
for y in y1..=y2 { for y in y1..=y2 {
img.put_pixel(x1, y, red); img.put_pixel(x1, y, red);
img.put_pixel(x2, y, red); img.put_pixel(x2, y, red);
// 如果要加粗，多画一列
if x1 + 1 < width { if x1 + 1 < width {
img.put_pixel(x1 + 1, y, red); img.put_pixel(x1 + 1, y, red);
} }
@@ -82,17 +82,27 @@ fn test_det_load() -> anyhow::Result<()> {
fs::read(image_path).map_err(|e| anyhow::anyhow!("无法读取图片 {}: {}", image_path, e))?; fs::read(image_path).map_err(|e| anyhow::anyhow!("无法读取图片 {}: {}", image_path, e))?;
println!("图片读取成功，字节大小: {}", image_bytes.len()); println!("图片读取成功，字节大小: {}", image_bytes.len());
let bboxes = det.detection(&image_bytes)?;
// 【修改点 1】将字节流解码为统一的 DynamicImage
let img = image::load_from_memory(&image_bytes)
.map_err(|e| anyhow::anyhow!("图片解码失败: {}", e))?;
// 【修改点 2】传入统一的 &DynamicImage 引用
let bboxes = det.detection(&img)?;
println!(":?{}", det); println!(":?{}", det);
println!("检测到的目标数量: {}", bboxes.len()); println!("检测到的目标数量: {}", bboxes.len());
if bboxes.is_empty() { if bboxes.is_empty() {
println!("未检测到任何目标。"); println!("未检测到任何目标。");
} else { } else {
save_debug_image(&image_bytes, &bboxes, "samples/result.jpg")?; // 如果 save_debug_image 报错，记得去把它的入参类型和内部访问也改为 DetectionResult
save_debug_image(&img, &bboxes, "samples/result.jpg")?;
for (i, bbox) in bboxes.iter().enumerate() { for (i, bbox) in bboxes.iter().enumerate() {
// 【修改点 3】将原来的 bbox[0].. 索引访问改为结构体字段访问
println!( println!(
"目标 [{}]: x1={}, y1={}, x2={}, y2={}", "目标 [{}]: x1={}, y1={}, x2={}, y2={}, 分数={:.4}, 类别ID={}",
i, bbox[0], bbox[1], bbox[2], bbox[3] i, bbox.x1, bbox.y1, bbox.x2, bbox.y2, bbox.score, bbox.class_id
); );
} }
} }