refactor(slide,det): 重构目标检测引擎并统一图像输入类型为 DynamicImage以及滑块匹配与比较引擎为 Rust 实现
- 统一 `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:
10
src/lib.rs
10
src/lib.rs
@@ -11,6 +11,7 @@ use std::fmt::{Display, Formatter};
|
||||
// 关键点:直接使用 tract 重导出的 ndarray
|
||||
use crate::charset::CharRestrict;
|
||||
use crate::model_metadata::ModelMetadata;
|
||||
use crate::models::det::DetectionResult;
|
||||
use crate::utils::color_filter::{ColorPreset, HsvRange};
|
||||
use models::det::Det;
|
||||
use models::loader::ModelSession;
|
||||
@@ -114,7 +115,7 @@ impl DdddOcr {
|
||||
Runtime::Ocr(s) => {
|
||||
let res = s.predictor().probability(true).predict(img)?;
|
||||
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().color_filter(&ColorPreset::Custom(vec![
|
||||
@@ -122,16 +123,15 @@ impl DdddOcr {
|
||||
// // 但上界的 H 通道写成了 240,超过了 180 的法定上限!
|
||||
// HsvRange::new((82, 221, 14), (240, 203, 82)),
|
||||
// ])).predict(img),
|
||||
|
||||
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型,无法执行 OCR")),
|
||||
}
|
||||
}
|
||||
pub fn detection(&self, img: &[u8]) -> Result<Vec<Vec<i32>>> {
|
||||
}
|
||||
pub fn detection(&self, img: &DynamicImage) -> Result<Vec<DetectionResult>> {
|
||||
match &self.runtime {
|
||||
Runtime::Det(s) => s.predict(img),
|
||||
Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型,无法执行检测")),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// struct Classification {}
|
||||
|
||||
@@ -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::{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 {
|
||||
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;
|
||||
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转换,这里直接进入核心逻辑
|
||||
self.get_bbox(image_bytes)
|
||||
self.get_bbox(image)
|
||||
}
|
||||
/// 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 (img_w, img_h) = img.dimensions();
|
||||
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);
|
||||
@@ -35,7 +49,7 @@ impl Det {
|
||||
let new_w = (img_w as f32 * r) as u32;
|
||||
|
||||
// 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>)
|
||||
let resized_rgb = resized.to_rgb8();
|
||||
// 创建 114 灰度填充的背景
|
||||
@@ -45,21 +59,24 @@ impl Det {
|
||||
// 将 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));
|
||||
for (x, y, pixel) in base_img.enumerate_pixels() {
|
||||
let x = x as usize;
|
||||
let y = y as usize;
|
||||
// 核心对标 Python 的 BGR 逻辑:
|
||||
// pixel[0] 是 R, pixel[1] 是 G, pixel[2] 是 B
|
||||
// 如果模型需要 BGR:
|
||||
// array[[0, 0, y as usize, x as usize]] = pixel[0] as f32;
|
||||
// array[[0, 1, y as usize, x as usize]] = pixel[1] as f32;
|
||||
// array[[0, 2, y as usize, x as usize]] = pixel[2] as f32;
|
||||
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
|
||||
|
||||
// 用连续的 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))
|
||||
}
|
||||
@@ -161,13 +178,14 @@ impl Det {
|
||||
}
|
||||
|
||||
/// 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<Vec<f32>> {
|
||||
) -> Vec<[f32; 6]> {
|
||||
let mut candidates = Vec::new();
|
||||
|
||||
// 1. 筛选高分框 (单次遍历完成 Argmax 和 Threshold 过滤)
|
||||
@@ -213,17 +231,17 @@ impl Det {
|
||||
.map(|k_idx| {
|
||||
let (orig_idx, score, cls_id) = candidates[k_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()
|
||||
}
|
||||
/// 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 解码
|
||||
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 (input_tensor, ratio) = self.preproc(&dynamic_img, (416, 416))?;
|
||||
let (input_tensor, ratio) = self.preproc(dynamic_img, (416, 416))?;
|
||||
|
||||
// tract 推理
|
||||
let outputs = self.session.run(tvec!(input_tensor.into()))?;
|
||||
@@ -236,7 +254,13 @@ impl Det {
|
||||
let pred = predictions.slice(s![0, .., ..]);
|
||||
|
||||
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());
|
||||
for i in 0..boxes.nrows() {
|
||||
@@ -247,17 +271,19 @@ impl Det {
|
||||
}
|
||||
|
||||
let detections = self.multiclass_nms(&boxes_xyxy, &scores, 0.45, 0.1);
|
||||
|
||||
Ok(detections
|
||||
let final_results = detections
|
||||
.into_iter()
|
||||
.map(|d| {
|
||||
vec![
|
||||
(d[0] as i32).max(0).min(orig_w as i32),
|
||||
(d[1] as i32).max(0).min(orig_h as i32),
|
||||
(d[2] as i32).max(0).min(orig_w as i32),
|
||||
(d[3] as i32).max(0).min(orig_h as i32),
|
||||
]
|
||||
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())
|
||||
.collect();
|
||||
Ok(final_results )
|
||||
}
|
||||
}
|
||||
|
||||
@@ -21,7 +21,7 @@ use crate::utils::cv_ops::rgb_to_opencv_hsv;
|
||||
|
||||
/// 推理最终输出的强类型外壳(完全 Owned,无任何生命周期,可直接转 JSON)
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub enum OcrOutput {
|
||||
pub enum OcrResult {
|
||||
/// 纯文本分支(对应 probability = false)
|
||||
Text(String),
|
||||
/// 包含全量概率的分支(对应 probability = true)
|
||||
@@ -35,27 +35,31 @@ pub enum OcrOutput {
|
||||
/// 不支持的模型或未知输出
|
||||
Unsupported { message: String },
|
||||
}
|
||||
impl OcrOutput {
|
||||
impl OcrResult {
|
||||
/// 消费自身,直接提取最终文本
|
||||
pub fn into_text(self) -> String {
|
||||
match self {
|
||||
OcrOutput::Text(text) => text,
|
||||
OcrOutput::Probability { text, .. } => text,
|
||||
OcrOutput::Unsupported { message } => {
|
||||
OcrResult::Text(text) => text,
|
||||
OcrResult::Probability { text, .. } => text,
|
||||
OcrResult::Unsupported { message } => {
|
||||
// 作为库,这里可以返回空,或者直接携带错误信息,取决于你的设计
|
||||
format!("Error: {}", message)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
impl fmt::Display for OcrOutput {
|
||||
impl fmt::Display for OcrResult {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
OcrOutput::Text(text) => {
|
||||
OcrResult::Text(text) => {
|
||||
// 纯文本分支,直接输出文本内容
|
||||
write!(f, "{}", text)
|
||||
}
|
||||
OcrOutput::Probability { text,probabilities, confidence } => {
|
||||
OcrResult::Probability {
|
||||
text,
|
||||
probabilities,
|
||||
confidence,
|
||||
} => {
|
||||
// 概率分支,友好地展示文本以及百分比形式的置信度
|
||||
// 1. 基本信息
|
||||
write!(f, "{} (置信度: {:.2}%)", text, confidence * 100.0)?;
|
||||
@@ -92,7 +96,7 @@ impl fmt::Display for OcrOutput {
|
||||
}
|
||||
write!(f, "]")
|
||||
}
|
||||
OcrOutput::Unsupported { message } => {
|
||||
OcrResult::Unsupported { message } => {
|
||||
// 错误分支,直观输出异常原因
|
||||
write!(f, "未识别成功: {}", message)
|
||||
}
|
||||
@@ -106,7 +110,7 @@ pub struct Ocr {
|
||||
}
|
||||
impl ModelSession for Ocr {
|
||||
fn get_model_type(&self) -> ModelType {
|
||||
todo!()
|
||||
todo!("使用thiserror作为错误处理的库,thiserror 专门用于开发库(Library)");
|
||||
}
|
||||
fn desc(&self) -> String {
|
||||
"Ocr Model 加载成功".to_string()
|
||||
@@ -114,6 +118,7 @@ impl ModelSession for Ocr {
|
||||
}
|
||||
impl Ocr {
|
||||
pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error> {
|
||||
|
||||
let session = ModelLoader::load_model(&model_path)?.session;
|
||||
Ok(Self {
|
||||
session,
|
||||
@@ -142,7 +147,6 @@ pub struct OcrPredictor<'a> {
|
||||
/// 是否修复PNG格式问题
|
||||
png_fix: bool,
|
||||
/// 是否返回概率信息
|
||||
#[allow(dead_code)]
|
||||
probability: bool,
|
||||
/// 颜色过滤:保留的颜色列表
|
||||
color_filter: Result<Option<Vec<HsvRange>>, String>,
|
||||
@@ -189,7 +193,7 @@ 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);
|
||||
// =====================================================================
|
||||
// 管道节点 1: 颜色过滤流水线
|
||||
@@ -217,11 +221,11 @@ impl<'a> OcrPredictor<'a> {
|
||||
|
||||
let raw_tensor = self.ocr.inference(tensor)?;
|
||||
|
||||
// 3. 后处理分流:直接返回 OcrOutput
|
||||
// 3. 后处理分流:直接返回 OcrResult
|
||||
let ocr_output = match raw_tensor.datum_type() {
|
||||
DatumType::I64 => self.extract_from_i64_tensor(raw_tensor)?,
|
||||
DatumType::F32 => self.process_f32_pipeline(raw_tensor)?,
|
||||
_ => OcrOutput::Unsupported {
|
||||
DatumType::I64 => self.process_i64_tensor(raw_tensor)?,
|
||||
DatumType::F32 => self.process_f32_tensor(raw_tensor)?,
|
||||
_ => OcrResult::Unsupported {
|
||||
message: format!("不支持的模型输出数据类型: {:?}", raw_tensor.datum_type()),
|
||||
},
|
||||
};
|
||||
@@ -325,17 +329,16 @@ impl<'a> OcrPredictor<'a> {
|
||||
}
|
||||
}
|
||||
impl<'a> OcrPredictor<'a> {
|
||||
// pub fn get_valid_indices(&self) -> HashSet<usize> {
|
||||
// match &self.charset_restrict {
|
||||
// Some(indices) => indices.iter().cloned().collect(),
|
||||
// // 如果是 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;
|
||||
fn is_valid_indices(&self, idx: usize) -> bool {
|
||||
if idx >= self.ocr.model_metadata.charset.size() {
|
||||
return false;
|
||||
}
|
||||
|
||||
match &self.charset_restrict {
|
||||
Some(v) => v.binary_search(&idx).is_ok(),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
/// 【按需延迟打印】:当用户真的需要“知道当前有哪些限制字符”时,一秒反查并打印
|
||||
/// 这里的 &str 完美借用了自 tokens,依然是彻底的零拷贝!
|
||||
pub fn valid_tokens(&self) -> Vec<&str> {
|
||||
@@ -410,7 +413,7 @@ impl<'a> OcrPredictor<'a> {
|
||||
(probabilities_list, confidence, predicted_indices)
|
||||
}
|
||||
/// 变体 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. 拿到底层的动态维度只读视图
|
||||
let view = raw_tensor.to_array_view::<i64>()?;
|
||||
|
||||
@@ -424,17 +427,17 @@ impl<'a> OcrPredictor<'a> {
|
||||
|
||||
// 4. 组装返回
|
||||
if self.probability {
|
||||
Ok(OcrOutput::Probability {
|
||||
Ok(OcrResult::Probability {
|
||||
text: final_text,
|
||||
probabilities: vec![], // I64 模型物理上丢失了全量 Logits 分值网,降级处理
|
||||
confidence: 1.0, // 判定即百分之百置信
|
||||
})
|
||||
} else {
|
||||
Ok(OcrOutput::Text(final_text))
|
||||
Ok(OcrResult::Text(final_text))
|
||||
}
|
||||
}
|
||||
/// 变体二(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();
|
||||
println!("模型输出shape数据: {:?}", shape);
|
||||
let view = raw_tensor.to_array_view::<f32>()?;
|
||||
@@ -476,7 +479,7 @@ impl<'a> OcrPredictor<'a> {
|
||||
// 5. 执行 CTC 解码
|
||||
let final_text = self.ctc_decode_to_string(&predicted_indices);
|
||||
|
||||
Ok(OcrOutput::Probability {
|
||||
Ok(OcrResult::Probability {
|
||||
text: final_text,
|
||||
probabilities: probabilities_list,
|
||||
confidence: confidence as f64,
|
||||
@@ -495,7 +498,7 @@ impl<'a> OcrPredictor<'a> {
|
||||
.collect();
|
||||
|
||||
let final_text = self.ctc_decode_to_string(&predicted_indices);
|
||||
Ok(OcrOutput::Text(final_text))
|
||||
Ok(OcrResult::Text(final_text))
|
||||
}
|
||||
}
|
||||
/// 获取有效字符索引列表 (用于外部验证或过滤)
|
||||
|
||||
@@ -27,7 +27,7 @@ impl Slide {
|
||||
Self
|
||||
}
|
||||
|
||||
/// 对应 Python: slide_match
|
||||
/// 对应 Python: slide_match 滑块匹配接口
|
||||
pub fn slide_match(
|
||||
&self,
|
||||
target_image: &DynamicImage,
|
||||
@@ -38,9 +38,8 @@ impl Slide {
|
||||
let background_array = image_to_ndarray(background_image);
|
||||
|
||||
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(
|
||||
&self,
|
||||
@@ -53,7 +52,6 @@ impl Slide {
|
||||
|
||||
// 2. 执行比较逻辑 (对应 _perform_slide_comparison)
|
||||
self.perform_slide_comparison(target_array.view(), background_array.view())
|
||||
.map_err(|e| anyhow!("滑块比较执行失败: {}", e))
|
||||
}
|
||||
/// 对应 Python: _perform_slide_comparison
|
||||
pub fn perform_slide_comparison(
|
||||
@@ -61,7 +59,7 @@ impl Slide {
|
||||
target: ArrayView3<u8>,
|
||||
background: ArrayView3<u8>,
|
||||
) -> Result<SlideResult> {
|
||||
let (h, w, _) = target.dim();
|
||||
// let (h, w, _) = target.dim();
|
||||
|
||||
// 1. 计算图像差异并灰度化 (对应 cv2.absdiff + cv2.cvtColor)
|
||||
// 使用 OpenCV 标准权重公式:0.299R + 0.587G + 0.114B
|
||||
@@ -77,6 +75,26 @@ impl Slide {
|
||||
// }
|
||||
// }
|
||||
// 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);
|
||||
|
||||
// 2. 转换为灰度数组 (复用你的 cv2.cvtColor)
|
||||
@@ -130,6 +148,30 @@ impl Slide {
|
||||
background: ArrayView3<u8>,
|
||||
simple_target: bool, // 增加这个参数
|
||||
) -> 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. 统一灰度化
|
||||
let target_gray = rgb_to_gray(target);
|
||||
let background_gray = rgb_to_gray(background);
|
||||
|
||||
@@ -95,9 +95,9 @@ pub fn bounding_rect(
|
||||
let h = max_y - min_y;
|
||||
(min_x, min_y, w, h)
|
||||
}
|
||||
pub fn calculate_center(max_loc: (u32, u32), tw: usize, th: usize) -> (i32, i32) {
|
||||
let center_x = max_loc.0 as i32 + (tw as i32 / 2);
|
||||
let center_y = max_loc.1 as i32 + (th as i32 / 2);
|
||||
pub fn calculate_center(top_left: (u32, u32), width: usize, height: usize) -> (i32, i32) {
|
||||
let center_x = top_left.0 as i32 + (width as i32 / 2);
|
||||
let center_y = top_left.1 as i32 + (height as i32 / 2);
|
||||
(center_x, center_y)
|
||||
}
|
||||
pub fn ndarray_to_luma8(array: ArrayView2<u8>) -> ImageBuffer<Luma<u8>, Vec<u8>> {
|
||||
|
||||
@@ -16,28 +16,22 @@ pub fn resize_image(
|
||||
target_height: u32,
|
||||
// resample 参数我们直接使用 FilterType,Lanczos3 是最接近 Python LANCZOS 的
|
||||
) -> DynamicImage {
|
||||
// 使用 resize 算法进行精确缩放
|
||||
// image::imageops::resize(
|
||||
// image,
|
||||
// target_width,
|
||||
// target_height,
|
||||
// FilterType::Lanczos3
|
||||
// )
|
||||
|
||||
// image::imageops::resize 的最高层封装
|
||||
// FilterType::Lanczos3 与 Python Pillow 的 Image.LANCZOS 算法完全对齐,缩放质量最高
|
||||
image.resize_exact(target_width, target_height, FilterType::Lanczos3)
|
||||
}
|
||||
pub fn resize_image1(
|
||||
image: &GrayImage,
|
||||
target_width: u32,
|
||||
target_height: u32,
|
||||
// resample 参数我们直接使用 FilterType,Lanczos3 是最接近 Python LANCZOS 的
|
||||
) -> GrayImage {
|
||||
// 使用 resize 算法进行精确缩放
|
||||
image::imageops::resize(
|
||||
image,
|
||||
target_width,
|
||||
target_height,
|
||||
FilterType::Lanczos3
|
||||
)
|
||||
}
|
||||
// pub fn resize_image(
|
||||
// image: &GrayImage,
|
||||
// target_width: u32,
|
||||
// target_height: u32,
|
||||
// // resample 参数我们直接使用 FilterType,Lanczos3 是最接近 Python LANCZOS 的
|
||||
// ) -> GrayImage {
|
||||
// // 使用 resize 算法进行精确缩放
|
||||
// image::imageops::resize(
|
||||
// image,
|
||||
// target_width,
|
||||
// target_height,
|
||||
// FilterType::Lanczos3
|
||||
// )
|
||||
// }
|
||||
@@ -1,8 +1,10 @@
|
||||
use ddddocr_rs::models::slide::Slide;
|
||||
use ddddocr_rs::{DdddOcr, DdddOcrBuilder}; // 假设你的包名是这个
|
||||
use image::Rgb;
|
||||
use image::{DynamicImage, Rgb};
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use ddddocr_rs::models::det::DetectionResult;
|
||||
|
||||
fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
|
||||
// 1. 先将泛型转为具体的 &Path 引用
|
||||
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(
|
||||
image_bytes: &[u8],
|
||||
bboxes: &Vec<Vec<i32>>,
|
||||
dynamic_img: &DynamicImage, // 【优化点 1】直接传入解码好的引用,拒绝重复解码
|
||||
bboxes: &[DetectionResult], // 【修改点 1】类型改为自定义结构体切片
|
||||
output_path: &str,
|
||||
) -> 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 (width, height) = img.dimensions();
|
||||
let red = Rgb([255u8, 0, 0]);
|
||||
|
||||
for bbox in bboxes {
|
||||
// 基础边界检查
|
||||
let x1 = bbox[0].max(0).min(width as i32 - 1) as u32;
|
||||
let y1 = bbox[1].max(0).min(height as i32 - 1) as u32;
|
||||
let x2 = bbox[2].max(0).min(width as i32 - 1) as u32;
|
||||
let y2 = bbox[3].max(0).min(height as i32 - 1) as u32;
|
||||
// 【修改点 2】将原来的索引 bbox[0].. 改为结构体字段访问 .x1, .y1 ..
|
||||
let x1 = bbox.x1.max(0).min(width as i32 - 1) as u32;
|
||||
let y1 = bbox.y1.max(0).min(height as i32 - 1) as u32;
|
||||
let x2 = bbox.x2.max(0).min(width as i32 - 1) as u32;
|
||||
let y2 = bbox.y2.max(0).min(height as i32 - 1) as u32;
|
||||
|
||||
// 绘制横向线条
|
||||
for x in x1..=x2 {
|
||||
img.put_pixel(x, y1, red);
|
||||
img.put_pixel(x, y2, red);
|
||||
// 如果要加粗,多画一行
|
||||
if y1 + 1 < height {
|
||||
img.put_pixel(x, y1 + 1, red);
|
||||
}
|
||||
@@ -47,7 +48,6 @@ fn save_debug_image(
|
||||
for y in y1..=y2 {
|
||||
img.put_pixel(x1, y, red);
|
||||
img.put_pixel(x2, y, red);
|
||||
// 如果要加粗,多画一列
|
||||
if x1 + 1 < width {
|
||||
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))?;
|
||||
|
||||
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!("检测到的目标数量: {}", bboxes.len());
|
||||
|
||||
if bboxes.is_empty() {
|
||||
println!("未检测到任何目标。");
|
||||
} 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() {
|
||||
// 【修改点 3】将原来的 bbox[0].. 索引访问改为结构体字段访问
|
||||
println!(
|
||||
"目标 [{}]: x1={}, y1={}, x2={}, y2={}",
|
||||
i, bbox[0], bbox[1], bbox[2], bbox[3]
|
||||
"目标 [{}]: x1={}, y1={}, x2={}, y2={}, 分数={:.4}, 类别ID={}",
|
||||
i, bbox.x1, bbox.y1, bbox.x2, bbox.y2, bbox.score, bbox.class_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user