diff --git a/ddddocr-core/Cargo.toml b/ddddocr-core/Cargo.toml index 5b5fcb7..318b703 100644 --- a/ddddocr-core/Cargo.toml +++ b/ddddocr-core/Cargo.toml @@ -5,7 +5,6 @@ edition = { workspace = true } license = { workspace = true } [dependencies] -anyhow = "1.0.102" image = "0.25.10" base64 = "0.22.1" imageproc = { version = "0.26.2", default-features = true } diff --git a/ddddocr-core/src/det/builder.rs b/ddddocr-core/src/det/builder.rs index 27d4520..539dff0 100644 --- a/ddddocr-core/src/det/builder.rs +++ b/ddddocr-core/src/det/builder.rs @@ -16,7 +16,7 @@ impl DetBuilder { self.device_id = device_id; self } - fn build(self, session: &dyn DetEngine) -> Detector<'_> { + fn build(self, session: &E) -> Detector<'_> { Detector { session, use_gpu: self.use_gpu, diff --git a/ddddocr-core/src/det/executor.rs b/ddddocr-core/src/det/executor.rs index eea7ad8..169244e 100644 --- a/ddddocr-core/src/det/executor.rs +++ b/ddddocr-core/src/det/executor.rs @@ -1,10 +1,9 @@ use crate::error::{Result, TensorError}; -use image::{imageops::FilterType, DynamicImage, GenericImageView}; -use ndarray::{prelude::*, s, Array2, Array3, Array4, Axis}; +use image::{DynamicImage, GenericImageView, imageops::FilterType}; +use ndarray::{Array2, Array3, Array4, Axis, prelude::*, s}; use std::fmt; // use tract_onnx::prelude::{Tensor}; - // use ddddocr_tract::det::session::DetSession; use crate::{DetEngine, DetOutput}; #[derive(Debug, Clone, Copy)] @@ -28,7 +27,6 @@ impl fmt::Display for DetectionResult { } } - pub struct Detector<'a> { pub(crate) session: &'a dyn DetEngine, #[allow(dead_code)] @@ -48,7 +46,7 @@ impl<'a> Detector<'a> { pub fn predict(&self, image: &DynamicImage) -> Result> { // Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑 - self.get_bbox(image) + Ok(self.get_bbox(image)?) } /// 2. preproc: 纯 Rust 实现 (替代 OpenCV) fn preproc(&self, image: &DynamicImage, input_size: (u32, u32)) -> (Array4, f32) { @@ -247,7 +245,10 @@ impl<'a> Detector<'a> { .collect() } /// 6. get_bbox (完全解耦 OpenCV) - pub fn get_bbox(&self, dynamic_img: &DynamicImage) -> Result> { + pub fn get_bbox( + &self, + dynamic_img: &DynamicImage, + ) -> Result, TensorError> { // 使用 utils crate 解码 // let dynamic_img = image::load_from_memory(image_bytes).context("Failed to decode utils")?; let (orig_w, orig_h) = dynamic_img.dimensions(); @@ -267,14 +268,14 @@ impl<'a> Detector<'a> { 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()).ok_or_else(|| { - TensorError::DimensionMismatch { - expected: format!("可广播至 cls_conf 形状 {:?}", cls_conf.shape()), - actual: obj_conf.shape().to_vec(), - } - })?; - // .context("ndarray broadcasting failed for scores calculation")?; + let obj_broadcast = + obj_conf + .broadcast(cls_conf.dim()) + .ok_or_else(|| TensorError::DimensionMismatch { + expected: format!("可广播至 cls_conf 形状 {:?}", cls_conf.shape()), + actual: obj_conf.shape().to_vec(), + })?; + // .context("ndarray broadcasting failed for scores calculation")?; let scores = &obj_broadcast * &cls_conf; // let scores = &pred.slice(s![.., 4..5]) * &pred.slice(s![.., 5..]); diff --git a/ddddocr-core/src/error.rs b/ddddocr-core/src/error.rs index d3be367..b0e2cb3 100644 --- a/ddddocr-core/src/error.rs +++ b/ddddocr-core/src/error.rs @@ -99,11 +99,10 @@ pub enum DdddError { /// 框架内部不可恢复的逻辑断言错误（例如解析节点 Fact 失败） #[error("内部严重逻辑错误: {0}")] Internal(String), - /// 【流派核心】接替 anyhow::Error 的用户自定义扩展错误 /// 承载任何第三方扩展、解密、特定预处理插件在执行时产生的自定义错误 #[error("用户自定义扩展错误: {0}")] - Other(#[source]Box), + Other(#[source] Box), } // ===================================================================== @@ -142,7 +141,7 @@ pub enum ImagePreprocessError { // ================= 新增：针对 HSV 和 Preset 的强类型错误 ================= /// HSV 颜色区间非法 (例如 H > 180 或 lower > upper) #[error("HSV 颜色区间参数非法: {0}")] - InvalidHsvRange (String ), + InvalidHsvRange(String), /// 不支持或未知的颜色预设名称 #[error("不支持的颜色预设名称: {0}")] @@ -151,6 +150,17 @@ pub enum ImagePreprocessError { /// 颜色过滤器/预处理规则配置非法导致失败 #[error("颜色过滤器配置无效或初始化失败: {0}")] FilterConfigInvalid(String), + + #[error("图像维度不匹配！{0}")] + MismatchDimensions (String), + + #[error("滑块模板尺寸 [{target_w}x{target_h}] 大于背景图 [{bg_w}x{bg_h}]")] + TargetExceedsBackground { + target_w: usize, + target_h: usize, + bg_w: usize, + bg_h: usize, + }, // #[error("Base64 解码失败: {0}")] // Base64(#[from] base64::DecodeError), // diff --git a/ddddocr-core/src/lib.rs b/ddddocr-core/src/lib.rs index bdb1f08..1ae5ea6 100644 --- a/ddddocr-core/src/lib.rs +++ b/ddddocr-core/src/lib.rs @@ -4,12 +4,13 @@ pub mod ocr; mod slide; pub mod utils; -use crate::error::Result; +use crate::error::{Result, TensorError}; pub use crate::slide::{SlideResult, Slider}; pub use crate::det::{DetBuilder, DetectionResult, Detector}; pub use crate::ocr::{Ocr, OcrBuilder, OcrResult}; -pub use ocr::metadata::ModelMetadata; +pub use crate::ocr::{ModelMetadata,Normalization}; +pub use ocr::Charset; // DetSession pub enum OcrOutput { @@ -23,11 +24,10 @@ pub enum DetOutput { /// 核心层定义的统一推理引擎接口。 /// 未来的 ddddocr-tract 和 ddddocr-ort 都必须实现这个 Trait - pub trait InferenceEngine { /// 关联类型：具体的 Session 需要声明自己到底产出什么枚举 type Output; - fn inference(&self, input_array: ndarray::Array4) -> Result; + fn inference(&self, input_array: ndarray::Array4) -> Result; } pub trait OcrEngine: InferenceEngine { diff --git a/ddddocr-core/src/ocr.rs b/ddddocr-core/src/ocr.rs index b25e6e8..3eac2f1 100644 --- a/ddddocr-core/src/ocr.rs +++ b/ddddocr-core/src/ocr.rs @@ -1,9 +1,13 @@ mod builder; +mod charset; +mod color_filter; mod executor; -pub mod metadata; -pub mod color_filter; +mod metadata; mod token_filter; pub use builder::OcrBuilder; +pub use charset::Charset; pub use executor::{Ocr, OcrResult}; +pub use metadata::{ModelMetadata, Normalization, Resize}; +pub use token_filter::TokenFilter; // pub use ddddocr_tract::session::OcrSession; diff --git a/ddddocr-core/src/ocr/builder.rs b/ddddocr-core/src/ocr/builder.rs index cfc0329..50c1b7c 100644 --- a/ddddocr-core/src/ocr/builder.rs +++ b/ddddocr-core/src/ocr/builder.rs @@ -1,8 +1,8 @@ use crate::ocr::executor::Ocr; // use ddddocr_tract::session::OcrSession; +use crate::OcrEngine; use crate::ocr::color_filter::ColorFilter; use crate::ocr::token_filter::TokenFilter; -use crate::OcrEngine; pub struct OcrBuilder { /// 是否修复PNG格式问题 @@ -49,7 +49,7 @@ impl OcrBuilder { self.charset_restrict = Some(Box::new(restrict)); self } - pub fn build(self, session: &dyn OcrEngine) -> Ocr<'_> { + pub fn build(self, session: &E) -> Ocr<'_> { // 1. 原地解析颜色过滤器 let final_color_ranges = match &self.color_filter { Some(filter) => filter.collect_to_vec(), diff --git a/ddddocr-core/src/ocr/charset.rs b/ddddocr-core/src/ocr/charset.rs new file mode 100644 index 0000000..c738611 --- /dev/null +++ b/ddddocr-core/src/ocr/charset.rs @@ -0,0 +1,66 @@ +use std::borrow::Cow; +use std::collections::HashMap; + +// ========================================== +// 3. 字符集核心结构体 (重命名为 Charset) +// ========================================== +#[derive(Debug, Clone)] +pub struct Charset { + // 使用 Cow 统一静态切片和动态读取的 Vec，内部实现真正的零拷贝 + pub tokens: Vec>, + // 反向查找表，保证字符转索引为 O(1) + pub char_to_idx: HashMap, usize>, + // 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速) + // pub valid_indices: HashSet, +} + +impl Charset { + // 内部底层统一收拢构造 + pub fn new(tokens: Vec>) -> Self { + let mut char_to_idx = HashMap::with_capacity(tokens.len()); + for (idx, token) in tokens.iter().enumerate() { + char_to_idx.entry(token.clone()).or_insert(idx); + // 如果字符集有重复，保留第一个遇到的索引 (符合 Python .index 逻辑) + // char_to_idx.entry(token.to_string()).or_insert(idx); + } + + Self { + tokens, + char_to_idx, + } + } + + // --- 业务策略方法 --- + + /// 将字符转为索引，不存在返回 -1 (保持与原 Python 库行为一致) + pub fn char_to_index(&self, char_str: &str) -> i32 { + if let Some(&idx) = self.char_to_idx.get(char_str) { + idx as i32 + } else { + -1 + } + } + + /// 将索引转为字符引用，零拷贝。若越界返回 None + pub fn index_to_char_ref(&self, index: usize) -> Option<&str> { + self.tokens.get(index).map(|cow| cow.as_ref()) + } + + pub fn is_valid_char(&self, char_str: &str) -> bool { + self.char_to_idx.get(char_str).is_some() + } + pub fn size(&self) -> usize { + self.tokens.len() + } +} + +// ========================================== +// 4. 标准 Display 接口实现 (对应 __str__) +// ========================================== +impl std::fmt::Display for Charset { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Charset [Total Size: {}", self.size(),) + } +} + + diff --git a/ddddocr-core/src/ocr/metadata.rs b/ddddocr-core/src/ocr/metadata.rs index 285399c..8d9a0eb 100644 --- a/ddddocr-core/src/ocr/metadata.rs +++ b/ddddocr-core/src/ocr/metadata.rs @@ -1,78 +1,11 @@ -use crate::error::{ Result}; -use serde::Deserialize; -use std::borrow::Cow; -use std::collections::HashMap; - -// ========================================== -// 3. 字符集核心结构体 (重命名为 Charset) -// ========================================== -#[derive(Debug, Clone)] -pub struct Charset { - // 使用 Cow 统一静态切片和动态读取的 Vec，内部实现真正的零拷贝 - pub tokens: Vec>, - // 反向查找表，保证字符转索引为 O(1) - pub char_to_idx: HashMap, usize>, - // 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速) - // pub valid_indices: HashSet, -} - -impl Charset { - // 内部底层统一收拢构造 - pub fn new(tokens: Vec>) -> Self { - let mut char_to_idx = HashMap::with_capacity(tokens.len()); - for (idx, token) in tokens.iter().enumerate() { - char_to_idx.entry(token.clone()).or_insert(idx); - // 如果字符集有重复，保留第一个遇到的索引 (符合 Python .index 逻辑) - // char_to_idx.entry(token.to_string()).or_insert(idx); - } - - Self { - tokens, - char_to_idx, - } - } - - // --- 业务策略方法 --- - - /// 将字符转为索引，不存在返回 -1 (保持与原 Python 库行为一致) - pub fn char_to_index(&self, char_str: &str) -> i32 { - if let Some(&idx) = self.char_to_idx.get(char_str) { - idx as i32 - } else { - -1 - } - } - - /// 将索引转为字符引用，零拷贝。若越界返回 None - pub fn index_to_char_ref(&self, index: usize) -> Option<&str> { - self.tokens.get(index).map(|cow| cow.as_ref()) - } - - pub fn is_valid_char(&self, char_str: &str) -> bool { - self.char_to_idx.get(char_str).is_some() - } - pub fn size(&self) -> usize { - self.tokens.len() - } -} - -// ========================================== -// 4. 标准 Display 接口实现 (对应 __str__) -// ========================================== -impl std::fmt::Display for Charset { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Charset [Total Size: {}", self.size(),) - } -} - - - // ===================================================================== // 1. 辅助定义的枚举与结构体 // ===================================================================== -#[derive(Debug, Clone, Copy, Deserialize)] -#[serde(rename_all = "snake_case")] // 支持 json 中写 "zero_to_one" 或 "minus_one_to_one" +use crate::ocr::Charset; +use std::borrow::Cow; + +#[derive(Debug, Clone, Copy)] pub enum Normalization { /// 映射到 [0.0, 1.0] -> pixel / 255.0 ZeroToOne, @@ -102,23 +35,6 @@ pub enum Resize { Square(u32), } -/// 仅用于反序列化 JSON 的中间临时结构体（DTO） -#[derive(Deserialize)] -struct ModelMetadataDto { - charset: Vec, - word: bool, - #[serde(alias = "image")] - resize: Vec, - channel: u8, - /// 新增：允许在配置文件中指定归一化策略。 - /// 使用 serde(default) 可以在不配置时提供一个默认值（比如默认 ZeroToOne） - #[serde(default = "default_normalization")] - normalization: Normalization, -} -fn default_normalization() -> Normalization { - Normalization::ZeroToOne -} - #[derive(Debug, Clone)] pub struct ModelMetadata { /// 字符集管理器 @@ -134,6 +50,21 @@ pub struct ModelMetadata { } impl ModelMetadata { + pub fn new( + charset: Charset, + word: bool, + resize: Resize, + channel: u8, + normalization: Normalization, + ) -> Self { + Self { + charset, + word, + resize, + channel, + normalization, + } + } // --- 优雅的工厂模式构造器 --- /// 通用的静态切片转换构造器 pub fn from_static_slice( @@ -152,49 +83,4 @@ impl ModelMetadata { normalization, } } - pub fn from_json_str(json_str: &str) -> Result { - let dto: ModelMetadataDto = serde_json::from_str(json_str) - .map_err(|e| anyhow!("JSON 反序列化失败，请检查字段是否完整: {}", e))?; - - // 1. 将 DTO 的字符串数组转化为强类型的 Charset - let tokens: Vec> = - dto.charset.into_iter().map(|s| Cow::Owned(s)).collect(); - let charset = Charset::new(tokens); - - // 2. 解析 resize 策略（重现 Python 的复杂条件判断） - if dto.resize.len() != 2 { - return Err(anyhow!( - "'resize (or image)' 字段必须是包含两个元素的数组，例如 [-1, 64]" - )); - } - let r0 = dto.resize[0]; - let r1 = dto.resize[1]; - - let resize = if r0 == -1 { - if dto.word { - // 如果 word 为 true，且包含 -1，Python 里是 resize 为 (r1, r1) 的正方形 - Resize::Square(r1 as u32) - } else { - // 如果 word 为 false，且包含 -1，Python 里是高度固定为 r1，宽度按原图比例缩放 - Resize::DynamicWidth(r1 as u32) - } - } else { - // 正常的固定宽高 - Resize::Fixed(r0 as u32, r1 as u32) - }; - - Ok(Self { - charset, - word: dto.word, - resize, - channel: dto.channel, - normalization: dto.normalization, - }) - } - /// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载） - pub fn from_json_bytes(bytes: &[u8]) -> Result { - let json_str = std::str::from_utf8(bytes) - .map_err(|e| anyhow!("JSON 字节流不是合法的 UTF-8 编码: {}", e))?; - Self::from_json_str(json_str) - } } diff --git a/ddddocr-core/src/slide.rs b/ddddocr-core/src/slide.rs index 3c87b32..7ba5067 100644 --- a/ddddocr-core/src/slide.rs +++ b/ddddocr-core/src/slide.rs @@ -1,7 +1,7 @@ -use crate::utils::image_convert::{image_to_ndarray, ColorMode}; +use crate::error::{ImagePreprocessError, Result}; +use crate::utils::image_convert::{ColorMode, image_to_ndarray}; use crate::utils::image_processor; use crate::utils::image_processor::{abs_diff, min_max_loc, ndarray_to_luma8, rgb_to_gray}; -use anyhow::{Result, anyhow}; use image::DynamicImage; use image::Luma; use imageproc::contrast::{ThresholdType, threshold}; @@ -32,8 +32,8 @@ impl fmt::Display for SlideResult { pub struct Slider; impl Slider { - pub fn new() -> Result { - Ok(Self) + pub fn new() -> Self { + Self } /// 对应 Python: slide_match 滑块匹配接口 pub fn slide_match( @@ -42,10 +42,11 @@ impl Slider { background_image: &DynamicImage, simple_target: bool, ) -> Result { - let target_array = image_to_ndarray(target_image,ColorMode::RGB)?; - let background_array = image_to_ndarray(background_image,ColorMode::RGB)?; + let target_array = image_to_ndarray(target_image, ColorMode::RGB)?; + let background_array = image_to_ndarray(background_image, ColorMode::RGB)?; self.perform_slide_match(target_array.view(), background_array.view(), simple_target) + .map_err(Into::into) } /// 对应 Python: slide_comparison 差异比较接口 /// 用于比较带坑位的图片与原始背景图，定位差异点 @@ -55,36 +56,35 @@ impl Slider { background_image: &DynamicImage, ) -> Result { // 1. 转换为 ndarray (HWC RGB) - let target_array = image_to_ndarray(target_image,ColorMode::RGB)?; - let background_array = image_to_ndarray(background_image,ColorMode::RGB)?; + let target_array = image_to_ndarray(target_image, ColorMode::RGB)?; + let background_array = image_to_ndarray(background_image, ColorMode::RGB)?; // 2. 执行比较逻辑 (对应 _perform_slide_comparison) self.perform_slide_comparison(target_array.view(), background_array.view()) + .map_err(Into::into) } /// 对应 Python: _perform_slide_comparison pub fn perform_slide_comparison( &self, target: ArrayView3, background: ArrayView3, - ) -> Result { + ) -> Result { // 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!( + return Err(ImagePreprocessError::MismatchDimensions(format!( "比较模式要求两张图分辨率与通道数完全一致！Target: [{}x{}x{}], Background: [{}x{}x{}]", - tw, - th, - tc, - bw, - bh, - bc - )); + tw, th, tc, bw, bh, bc + ))); } if th == 0 || tw == 0 { - return Err(anyhow!("输入图像尺寸不能为0")); + return Err(ImagePreprocessError::InvalidDimensions { + expected: "输入图像尺寸不能为0".to_string(), + actual: vec![th, tw], + }); } let diff_array = abs_diff(&target, &background); @@ -140,29 +140,31 @@ impl Slider { target: ArrayView3, background: ArrayView3, simple_target: bool, // 增加这个参数 - ) -> Result { + ) -> Result { 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")); + return Err(ImagePreprocessError::InvalidDimensions { + expected: "输入图像的宽度或高度不能为0".to_string(), + actual: vec![th, tw, tc], + }); } if th > bh || tw > bw { - return Err(anyhow!( - "尺寸不匹配：滑块模板(target)尺寸 [{}x{}] 不能大于背景图(background) [{}x{}]", - tw, - th, - bw, - bh - )); + return Err(ImagePreprocessError::TargetExceedsBackground { + // "尺寸不匹配：滑块模板(target)尺寸 [{}x{}] 不能大于背景图(background) [{}x{}]", + target_w: tw, + target_h: th, + bg_w: bw, + bg_h: bh, + }); } if tc != bc { - return Err(anyhow!( + return Err(ImagePreprocessError::MismatchDimensions(format!( "目标图与背景图的通道数不一致 (target: {}, bg: {})", - tc, - bc - )); + tc, bc + ))); } // 1. 统一灰度化 @@ -171,11 +173,11 @@ impl Slider { if simple_target { // 2a. 简单模式：直接在灰度图上匹配 - self.simple_template_match(target_gray.view(), background_gray.view()) + Ok(self.simple_template_match(target_gray.view(), background_gray.view())) } else { // 2b. 复杂模式：先提取边缘，再匹配 - self.edge_based_match(target_gray.view(), background_gray.view()) + Ok(self.edge_based_match(target_gray.view(), background_gray.view())) } } /// 对应 Python: _simple_template_match @@ -185,7 +187,7 @@ impl Slider { &self, target: ArrayView2, background: ArrayView2, - ) -> Result { + ) -> SlideResult { // 1. 将 ndarray 转换为 imageproc 需要的 ImageBuffer (无拷贝或轻量转换) // 转换逻辑 (假设你已经有方法转回 ImageBuffer) let t_buf = ndarray_to_luma8(target); @@ -211,12 +213,12 @@ impl Slider { // println!("Rust Target Width (tw): {}", tw); // println!("Rust Best Max Loc X: {}", max_loc.0); // println!("Rust Final Center X: {}", center_x); - Ok(SlideResult { + SlideResult { target: [center_x, center_y], target_x: center_x, target_y: center_y, confidence: max_val as f64, - }) + } } /// 对应 Python: _edge_based_match @@ -225,7 +227,7 @@ impl Slider { &self, target: ArrayView2, background: ArrayView2, - ) -> Result { + ) -> SlideResult { // 1. 将 ndarray 转换为 ImageBuffer // 注意：Canny 和 match_template 需要 ImageBuffer 格式 let t_buf = ndarray_to_luma8(target); @@ -260,11 +262,11 @@ impl Slider { println!("-Rust Target Width (tw): {}", tw); println!("-Rust Best Max Loc X: {}", max_loc.0); println!("-Rust Final Center X: {}", center_x); - Ok(SlideResult { + SlideResult { target: [center_x, center_y], target_x: center_x, target_y: center_y, confidence: max_val as f64, - }) + } } } diff --git a/ddddocr-core/src/utils/image_convert.rs b/ddddocr-core/src/utils/image_convert.rs index 30808ce..5900116 100644 --- a/ddddocr-core/src/utils/image_convert.rs +++ b/ddddocr-core/src/utils/image_convert.rs @@ -1,4 +1,4 @@ -use crate::error::{DdddError, ImagePreprocessError, Result}; +use crate::error::{ImagePreprocessError, Result}; use image::{DynamicImage, GenericImageView, ImageBuffer, Luma, Rgb, Rgba}; use ndarray::{Array3, ArrayViewD}; @@ -11,7 +11,7 @@ pub enum ColorMode { /// 封装数组转图像的逻辑， // 对应 Python 版 _numpy_to_pil_image -pub fn ndarray_to_hwc_image(array: ArrayViewD) -> Result { +pub fn ndarray_to_hwc_image(array: ArrayViewD) -> Result { let shape = array.shape(); let dim = shape.len(); @@ -35,7 +35,7 @@ pub fn ndarray_to_hwc_image(array: ArrayViewD) -> Result { // 对应 Python: array.shape[2] == 4 (RGBA H, W, 4) 4 => ColorMode::RGBA, _ => { - return Err(ImagePreprocessError::UnsupportedChannels(c))?; + return Err(ImagePreprocessError::UnsupportedChannels(c)); } } } @@ -43,8 +43,7 @@ pub fn ndarray_to_hwc_image(array: ArrayViewD) -> Result { return Err(ImagePreprocessError::InvalidDimensions { expected: "2D (H,W) 或 3D (H,W,C)".to_string(), actual: shape.to_vec(), - } - .into()); + }); } }; from_ndarray(array, color_mode) @@ -98,7 +97,7 @@ pub fn png_rgba_white_preprocess(img: &DynamicImage) -> DynamicImage { DynamicImage::ImageRgb8(background) } /// 将 DynamicImage 转换为 array 数组 -pub fn image_to_ndarray(image: &DynamicImage, mode: ColorMode) -> Result> { +pub fn image_to_ndarray(image: &DynamicImage, mode: ColorMode) -> Result,ImagePreprocessError> { // 1. 模式转换 (对应 utils.convert(target_mode)),此函数在时保留看后续优化是否需要替代image_to_ndarray // Rust utils 库通过 to_rgb8, to_luma8 等方法实现转换 let (width, height) = image.dimensions(); @@ -114,7 +113,7 @@ pub fn image_to_ndarray(image: &DynamicImage, mode: ColorMode) -> Result, mode: ColorMode) -> Result { +pub fn ndarray_to_image(array: ArrayViewD, mode: ColorMode) -> Result { let shape = array.shape(); // 基础边界检查：至少要有 H 和 W 两个维度 @@ -127,7 +126,7 @@ pub fn ndarray_to_image(array: ArrayViewD, mode: ColorMode) -> Result, mode: ColorMode) -> Result { +fn from_ndarray(array: ArrayViewD, mode: ColorMode) -> Result { let shape = array.shape(); // 映射：ndarray 的 shape 默认是 [Height, Width, (Channels)] @@ -158,7 +157,6 @@ fn from_ndarray(array: ArrayViewD, mode: ColorMode) -> Result height, channels, } - .into() }; // 2. 重新解释内存并构建 ImageBuffer diff --git a/ddddocr-core/src/utils/tensor_transform.rs b/ddddocr-core/src/utils/tensor_transform.rs index deec7c6..8eddc7e 100644 --- a/ddddocr-core/src/utils/tensor_transform.rs +++ b/ddddocr-core/src/utils/tensor_transform.rs @@ -1,5 +1,5 @@ use crate::OcrOutput; -use crate::error::{DdddError, Result, TensorError}; +use crate::error::{Result, TensorError}; use ndarray::s; /// 核心层复用资产：将异构的动态维度矩阵转化为标准 OCR 2D Logits 矩阵 pub fn normalize_ocr_logits(array: ndarray::ArrayViewD, shape: &[usize]) -> Result { diff --git a/ddddocr-tract/src/det/session.rs b/ddddocr-tract/src/det/session.rs index bd75cd7..cbfecc2 100644 --- a/ddddocr-tract/src/det/session.rs +++ b/ddddocr-tract/src/det/session.rs @@ -1,68 +1,45 @@ -use crate::loader::ModelLoader; -use anyhow::Context; -use ddddocr_core::error::{DdddError, Result, TensorErrorReason}; +use crate::types::Session; +use ddddocr_core::error::{Result, TensorError}; use ddddocr_core::{DetEngine, DetOutput, InferenceEngine}; use ndarray::Ix3; -use std::path::Path; -use tract_onnx::prelude::{Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tvec}; +use tract_onnx::prelude::{tvec, IntoTensor, Tensor}; #[derive(Debug)] pub struct DetSession { - pub session: RunnableModel, Graph>>, + pub session: Session, } impl DetSession { - pub fn new

(model_path: P) -> Result - where - P: AsRef, - { - let session = ModelLoader::model_for_path(&model_path)?.session; - Ok(Self { session }) + pub fn new(session: Session) -> Self { + Self { session } } - pub fn model_from_bytes(model_bytes: &[u8]) -> Result { - let session = ModelLoader::model_from_bytes(model_bytes)?.session; - Ok(Self { session }) - } - // pub fn inference(&self, tensor: Tensor) -> anyhow::Result { - // // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 - // // let result = self.ocr.run(tvec!(tensor.into()))?; - // let mut result = self - // .session - // .run(tvec!(tensor.into())) - // .context("执行模型推理失败")?; - // println!("模型输出原始数据: {:?}", result); - // Ok(result.swap_remove(0).into_tensor()) - // } } impl InferenceEngine for DetSession { type Output = DetOutput; // 明确绑定 OCR 小枚举 - fn inference(&self, input_array: ndarray::Array4) -> Result { + fn inference(&self, input_array: ndarray::Array4) -> Result { // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 // let result = self.ocr.run(tvec!(tensor.into()))?; let tensor = Tensor::from(input_array); - let mut result = self.session.run(tvec!(tensor.into())).map_err(|_| { - DdddError::Inference(TensorErrorReason::EngineError( - "执行模型推理失败".to_string(), - )) - })?; + let mut result = self + .session + .run(tvec!(tensor.into())) + .map_err(|_| TensorError::Engine("执行模型推理失败".to_string()))?; println!("模型输出原始数据: {:?}", result); // Ok(result.swap_remove(0).into_tensor()) let raw_tensor = result.swap_remove(0).into_tensor(); let array_d = raw_tensor.into_array::().map_err(|_| { - DdddError::Inference(TensorErrorReason::EngineError( - "Tract 实体张量无法转换为 ndarray::ArrayD".to_string(), - )) + TensorError::Engine("Tract 实体张量无法转换为 ndarray::ArrayD".to_string()) })?; // 提前利用克隆(Clone)备份好当前未转维度前的真实 shape (Vec) let actual_shape = array_d.shape().to_vec(); let array3 = array_d.into_dimensionality::().map_err(|_| { - DdddError::Inference(TensorErrorReason::TensorDimensionMismatch { + TensorError::DimensionMismatch { expected: "3D 检测矩阵 [Batch, Box_Count, Box_Attributes]".to_string(), actual: actual_shape, // 优雅降维失败时动态捕获 - }) + } })?; Ok(DetOutput::Detection(array3)) diff --git a/ddddocr-tract/src/error.rs b/ddddocr-tract/src/error.rs new file mode 100644 index 0000000..91daa51 --- /dev/null +++ b/ddddocr-tract/src/error.rs @@ -0,0 +1,8 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum TensorError { + /// 替换原有的 anyhow::Error，明确将 Tract/ONNX 引擎底层报错序列化为干净的 String + #[error("推理引擎内部发生异常: {0}")] + Engine(String), +} diff --git a/ddddocr-tract/src/lib.rs b/ddddocr-tract/src/lib.rs index 9725b5d..27fcca1 100644 --- a/ddddocr-tract/src/lib.rs +++ b/ddddocr-tract/src/lib.rs @@ -1,6 +1,10 @@ mod det; pub mod loader; mod ocr; +mod types; +mod error; +pub use ddddocr_core::ocr::OcrBuilder; +pub use ddddocr_core::{SlideResult, Slider}; pub use det::session::DetSession; -pub use ocr::session::OcrSession; \ No newline at end of file +pub use ocr::session::OcrSession; diff --git a/ddddocr-tract/src/loader.rs b/ddddocr-tract/src/loader.rs index 2cb0554..1e4294b 100644 --- a/ddddocr-tract/src/loader.rs +++ b/ddddocr-tract/src/loader.rs @@ -1,214 +1,7 @@ -use anyhow::Context; -use ddddocr_core::error::{DdddError, Result}; -use std::fmt; -use std::io::Cursor; -use tract_onnx::onnx; -use tract_onnx::prelude::*; -// 引入核心层的统一错误类型 -/// 明确命名为 AxisDim，代表模型某一个轴的维度特征 -#[derive(Clone, PartialEq, Eq)] -pub enum AxisDim { - /// 静态固定维度（如通道数固定为 1，高度固定为 64） - Static(usize), - /// 动态符号维度（如宽度是动态的 "image_width"） - Dynamic(String), -} +mod error; +mod metadata; +mod model; -impl AxisDim { - /// 便捷方法：判断是否为动态维度 - pub fn is_dynamic(&self) -> bool { - matches!(self, AxisDim::Dynamic(_)) - } -} -/// 自定义 Debug 格式化输出，彻底融化套娃外壳，保证日志干净漂亮 -impl fmt::Debug for AxisDim { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - AxisDim::Static(size) => write!(f, "{}", size), - AxisDim::Dynamic(expr) => write!(f, "Dynamic(\"{}\")", expr), - } - } -} -/// 模拟 Python 的 input_info 和 output_info 结构 -#[derive(Debug, Clone)] -pub struct TensorInfo { - pub name: String, - pub shape: Vec, // 既包含 Fixed 静态维度，也包含 Dynamic 动态符号 - pub data_type: DatumType, // 对应 Python 的 type -} - -/// 最终返回的模型完整信息 -#[derive(Debug, Clone)] -pub struct ModelInfo { - pub inputs: Vec, - pub outputs: Vec, - /// 硬件执行提供者（采用 Option 兼容不同底层的推理引擎） - pub providers: Option>, -} - -pub struct ModelLoader { - pub session: RunnableModel, Graph>>, -} - -impl ModelLoader { - pub fn model_for_path

(model_path: P) -> Result - where - P: AsRef, - { - let session = onnx() - .model_for_path(model_path).map_err(DdddError::new)? - // .with_context(|| "加载 ONNX 模型失败，请检查路径是否正确")? - .into_optimized().map_err(DdddError::new)? - // .with_context(|| "优化 Tract 模型图失败")? - .into_runnable().map_err(DdddError::new)?; - // .with_context(|| "构建可运行 Tract 实例失败")?; - Ok(Self { session }) - } - /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用） - pub fn model_from_bytes(model_bytes: &[u8]) -> Result { - // 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read） - let mut cursor = Cursor::new(model_bytes); - - let session = onnx() - .model_for_read(&mut cursor) - .map_err(DdddError::new)? - // .with_context(|| "从内存字节流解析 ONNX 模型失败")? - .into_optimized() - .map_err(DdddError::new)? - // .with_context(|| "优化 Tract 模型图失败")? - .into_runnable().map_err(DdddError::new)?; - // .with_context(|| "构建可运行 Tract 实例失败")?; - - Ok(Self { session }) - } -} -impl ModelLoader { - /// 获取模型详细元数据信息（对标 Python ddddocr 的 get_model_info） - /// 完美包容 [1, 1, 64, image_width] 这样的变长图像模型 - /// 获取模型详细元数据信息（代码更紧凑、优雅） - pub fn model_info(&self) -> Result { - let model = self.session.model(); - - // 使用私有辅助函数统一处理，消除重复代码 - let inputs = self.resolve_tensors( - model - .input_outlets().map_err(DdddError::new)? - // .map_err(|e| DdddError::InternalError(format!("获取输入节点失败: {:?}", e)))?, - )?; - let outputs = self.resolve_tensors( - model - .output_outlets().map_err(DdddError::new)? - // .map_err(|e| DdddError::InternalError(format!("获取输出节点失败: {:?}", e)))?, - )?; - - Ok(ModelInfo { - inputs, - outputs, - providers: None, - }) - } - - /// 提取出来的公共转换逻辑：将一组 OutletId 解析为 TensorInfo 列表 - fn resolve_tensors(&self, outlets: &[OutletId]) -> Result> { - let model = self.session.model(); - - outlets - .iter() - .map(|&outlet_id| { - let fact = model.outlet_fact(outlet_id).map_err(DdddError::new)?; - // .map_err(|e| { - // DdddError::InternalError(format!("解析节点 Fact 失败: {:?}", e)) - // })?; - - let shape = self.resolve_shape(&fact.shape)?; - let node_name = model.node(outlet_id.node).name.clone(); - - Ok(TensorInfo { - name: node_name, - shape, - data_type: fact.datum_type, - }) - }) - .collect() // 函数式声明：自动传播第一处发生的错误 - } - - /// 安全还原 Tract 维度至 Vec - fn resolve_shape(&self, shape_fact: &ShapeFact) -> Result> { - let tract_shape = shape_fact.to_tvec(); - - let resolved = tract_shape - .iter() - .map(|dim| { - // 防御性编程：必须同时满足能够转换为 i64 且 大于等于 0 - if let Ok(size) = dim.to_i64() { - if size >= 0 { - AxisDim::Static(size as usize) - } else { - // 如果 ONNX 导出时某些动态维度被标记为了 -1，安全地作为动态符号捕获 - AxisDim::Dynamic(dim.to_string()) - } - } else { - AxisDim::Dynamic(dim.to_string()) - } - }) - .collect(); - - Ok(resolved) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// 辅助函数：动态构建一个简单的 ONNX/Tract 内存模型图用于测试 - fn create_test_model() -> std::result::Result< - RunnableModel, Graph>>, - anyhow::Error, - > { - let mut rect = tract_onnx::prelude::Graph::default(); - - // 0.21.10 最稳妥的静态 Fact 构建 - let input_fact = TypedFact::dt_shape(DatumType::F32, &[1, 3, 224, 224]); - - let input_node = rect - .add_source("input", input_fact) - .map_err(|e| anyhow::anyhow!("{:?}", e))?; - - rect.set_input_outlets(&[input_node.into()]) - .map_err(|e| anyhow::anyhow!("{:?}", e))?; - rect.set_output_outlets(&[input_node.into()]) - .map_err(|e| anyhow::anyhow!("{:?}", e))?; - - let typed = rect - .into_optimized() - .map_err(|e| anyhow::anyhow!("{:?}", e))?; - let runnable = typed - .into_runnable() - .map_err(|e| anyhow::anyhow!("{:?}", e))?; - Ok(runnable) - } - - #[test] - fn test_input_output_shapes_and_type() { - let session = create_test_model().expect("建立测试模型图失败"); - let loader = ModelLoader { session }; - println!("{:?}", loader.model_info().unwrap()); - // 1. 测试输入维度解析 - } - - #[test] - fn test_resolve_shape_logic_direct() { - // 创建一个哑 ModelLoader 实例（session 用不上，因为我们直接测私有方法） - let session = create_test_model().expect("建立测试模型图失败"); - let loader = ModelLoader { session }; - let dims: Vec = vec![TDim::from(1), TDim::from(3), TDim::from(224)]; - // 方案二的精髓：我们直接利用已导出的 ShapeFact 来纯手工验证边界逻辑！ - // 1. 验证纯静态维度是否被正确还原 - let static_shape = ShapeFact::from_dims(dims); - - let res = loader - .resolve_shape(&static_shape) - .expect("解析静态 shape 失败"); - } -} +pub use error::{Error, ParseError, Result}; +pub use metadata::ModelMetadata; +pub use model::ModelLoader; diff --git a/ddddocr-tract/src/loader/error.rs b/ddddocr-tract/src/loader/error.rs new file mode 100644 index 0000000..fd39b5d --- /dev/null +++ b/ddddocr-tract/src/loader/error.rs @@ -0,0 +1,50 @@ +use tract_onnx::prelude::TractError; +pub type Result = std::result::Result; +#[derive(thiserror::Error, Debug)] +pub enum Error { + /// 解析 ONNX 模型/路径失败（如文件损坏、算子不支持、路径非法） + #[error("解析 ONNX 模型结构失败: {0}")] + ModelParse(#[from] ParseError), + + /// 模型计算图优化失败（如常量折叠、形状推导失败） + #[error("优化 Tract 模型图失败: {0}")] + OptimizationFailed(#[source] TractError), + + /// 构建可执行 Session 失败（如输入输出 Tensor 类型/形状未确定） + #[error("构建可运行 Tract 实例失败: {0}")] + RunnableBuildFailed(#[source] TractError), + + /// JSON 反序列化失败（自动透传 serde_json 报错） + #[error("模型 Metadata JSON 解析失败: {0}")] + JsonParse(#[from] serde_json::Error), + + /// 字节流非合法 UTF-8 编码（自动透传 Utf8Error） + #[error("Metadata 字节流不是合法的 UTF-8 编码: {0}")] + InvalidUtf8(#[from] std::str::Utf8Error), + + #[error("模型元数据解析失败: {0}")] + MetadataParse(String), + + /// 承载任何第三方扩展、解密、特定预处理插件在执行时产生的自定义错误 + #[error("{0}: {1}")] + Other(String, #[source] Box), +} +impl Error { + /// 方便将任何第三方 Error 包装为 Error::Other + pub fn new(msg: impl Into, err: E) -> Self + where + E: Into>, + { + Self::Other(msg.into(), err.into()) + } +} +#[derive(thiserror::Error,Debug)] +pub enum ParseError{ + /// 策略 A：从文件路径加载失败（附带路径上下文信息，方便排查是找不到文件还是格式不对） + #[error("从路径 '{0}' 加载 ONNX 模型失败: {1}")] + Path(String, #[source] TractError), + + /// 策略 B：从内存字节流加载失败（如 include_bytes! 传入的字节流损坏） + #[error("从内存字节流解析 ONNX 模型失败: {0}")] + Bytes(#[source] TractError), +} \ No newline at end of file diff --git a/ddddocr-tract/src/loader/metadata.rs b/ddddocr-tract/src/loader/metadata.rs new file mode 100644 index 0000000..7388fa3 --- /dev/null +++ b/ddddocr-tract/src/loader/metadata.rs @@ -0,0 +1,93 @@ +use crate::loader::error::{Error, Result}; + +pub use ddddocr_core::ModelMetadata; +use ddddocr_core::ocr::Resize; +use ddddocr_core::{Charset, Normalization}; +use serde::Deserialize; +use std::borrow::Cow; + +#[derive(Deserialize)] +#[serde(rename_all = "snake_case")] // 支持 json 中写 "zero_to_one" 或 "minus_one_to_one" +enum NormalizationDto { + /// 映射到 [0.0, 1.0] -> pixel / 255.0 + ZeroToOne, + /// 映射到 [-1.0, 1.0] -> (pixel / 255.0 - 0.5) / 0.5 + MinusOneToOne, +} + +impl From for Normalization { + fn from(dto: NormalizationDto) -> Self { + match dto { + NormalizationDto::ZeroToOne => Normalization::ZeroToOne, + NormalizationDto::MinusOneToOne => Normalization::MinusOneToOne, + } + } +} + +/// 仅用于反序列化 JSON 的中间临时结构体（DTO） +#[derive(Deserialize)] +struct ModelMetadataDto { + charset: Vec, + word: bool, + #[serde(alias = "image")] + resize: Vec, + channel: u8, + /// 新增：允许在配置文件中指定归一化策略。 + /// 使用 serde(default) 可以在不配置时提供一个默认值（比如默认 ZeroToOne） + #[serde(default = "default_normalization")] + normalization: NormalizationDto, +} +fn default_normalization() -> NormalizationDto { + NormalizationDto::ZeroToOne +} + +/// Tract 专属扩展trait 或 工具函数 +pub trait TractModelMetadata: Sized { + fn from_json_str(json_str: &str) -> Result; + /// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载） + fn from_json_bytes(bytes: &[u8]) -> Result { + let json_str = std::str::from_utf8(bytes)?; + Self::from_json_str(json_str) + } +} +impl TractModelMetadata for ModelMetadata { + // --- 优雅的工厂模式构造器 --- + fn from_json_str(json_str: &str) -> Result { + let dto: ModelMetadataDto = serde_json::from_str(json_str)?; + + // 1. 将 DTO 的字符串数组转化为强类型的 Charset + let tokens: Vec> = + dto.charset.into_iter().map(|s| Cow::Owned(s)).collect(); + let charset = Charset::new(tokens); + + // 2. 解析 resize 策略（重现 Python 的复杂条件判断） + if dto.resize.len() != 2 { + return Err(Error::MetadataParse( + "'resize (or image)' 字段必须是包含两个元素的数组，例如 [-1, 64]".to_string(), + )); + } + let r0 = dto.resize[0]; + let r1 = dto.resize[1]; + + let resize = if r0 == -1 { + if dto.word { + // 如果 word 为 true，且包含 -1，Python 里是 resize 为 (r1, r1) 的正方形 + Resize::Square(r1 as u32) + } else { + // 如果 word 为 false，且包含 -1，Python 里是高度固定为 r1，宽度按原图比例缩放 + Resize::DynamicWidth(r1 as u32) + } + } else { + // 正常的固定宽高 + Resize::Fixed(r0 as u32, r1 as u32) + }; + + Ok(ModelMetadata::new( + charset, + dto.word, + resize, + dto.channel, + dto.normalization.into(), + )) + } +} diff --git a/ddddocr-tract/src/loader/model.rs b/ddddocr-tract/src/loader/model.rs new file mode 100644 index 0000000..fcfc6ac --- /dev/null +++ b/ddddocr-tract/src/loader/model.rs @@ -0,0 +1,99 @@ +use crate::loader::error; +use crate::loader::error::{Error, ParseError, Result}; +use crate::types::Session; +use std::io::Cursor; +use tract_onnx::onnx; +use tract_onnx::prelude::*; + +pub struct ModelLoader; + +impl ModelLoader { + pub fn model_for_path

(model_path: P) -> Result + where + P: AsRef, + { + let path_ref = model_path.as_ref(); + + let session = onnx() + .model_for_path(path_ref) + .map_err(|e| ParseError::Path(path_ref.display().to_string(), e))? + // .with_context(|| "加载 ONNX 模型失败，请检查路径是否正确")? + .into_optimized() + .map_err(Error::OptimizationFailed)? + // .with_context(|| "优化 Tract 模型图失败")? + .into_runnable() + .map_err(Error::RunnableBuildFailed)?; + // .with_context(|| "构建可运行 Tract 实例失败")?; + Ok(session) + } + /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用） + pub fn model_from_bytes(model_bytes: &[u8]) -> Result { + // 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read） + let mut cursor = Cursor::new(model_bytes); + + let session = onnx() + .model_for_read(&mut cursor) + .map_err(ParseError::Bytes)? + // .with_context(|| "从内存字节流解析 ONNX 模型失败")? + .into_optimized() + .map_err(Error::OptimizationFailed)? + // .with_context(|| "优化 Tract 模型图失败")? + .into_runnable() + .map_err(Error::RunnableBuildFailed)?; + // .with_context(|| "构建可运行 Tract 实例失败")?; + + Ok(session) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 辅助函数：动态构建一个简单的 ONNX/Tract 内存模型图用于测试 + fn create_test_model() -> std::result::Result { + let mut rect = tract_onnx::prelude::Graph::default(); + + // 0.21.10 最稳妥的静态 Fact 构建 + let input_fact = TypedFact::dt_shape(DatumType::F32, &[1, 3, 224, 224]); + + let input_node = rect + .add_source("input", input_fact) + .map_err(|e| anyhow::anyhow!("{:?}", e))?; + + rect.set_input_outlets(&[input_node.into()]) + .map_err(|e| anyhow::anyhow!("{:?}", e))?; + rect.set_output_outlets(&[input_node.into()]) + .map_err(|e| anyhow::anyhow!("{:?}", e))?; + + let typed = rect + .into_optimized() + .map_err(|e| anyhow::anyhow!("{:?}", e))?; + let runnable = typed + .into_runnable() + .map_err(|e| anyhow::anyhow!("{:?}", e))?; + Ok(runnable) + } + + // #[test] + // fn test_input_output_shapes_and_type() { + // let session = create_test_model().expect("建立测试模型图失败"); + // + // println!("{:?}", ModelLoader::model_info(&session).unwrap()); + // // 1. 测试输入维度解析 + // } + // + // #[test] + // fn test_resolve_shape_logic_direct() { + // // 创建一个哑 ModelLoader 实例（session 用不上，因为我们直接测私有方法） + // let session = create_test_model().expect("建立测试模型图失败"); + // + // let dims: Vec = vec![TDim::from(1), TDim::from(3), TDim::from(224)]; + // // 方案二的精髓：我们直接利用已导出的 ShapeFact 来纯手工验证边界逻辑！ + // // 1. 验证纯静态维度是否被正确还原 + // let static_shape = ShapeFact::from_dims(dims); + // + // let res = ModelLoader + // ::resolve_shape(&static_shape); + // } +} diff --git a/ddddocr-tract/src/ocr.rs b/ddddocr-tract/src/ocr.rs index 3e7b3f8..f52f1c4 100644 --- a/ddddocr-tract/src/ocr.rs +++ b/ddddocr-tract/src/ocr.rs @@ -1 +1 @@ -pub mod session; \ No newline at end of file +pub mod session; diff --git a/ddddocr-tract/src/ocr/session.rs b/ddddocr-tract/src/ocr/session.rs index bb55dfc..247b5bb 100644 --- a/ddddocr-tract/src/ocr/session.rs +++ b/ddddocr-tract/src/ocr/session.rs @@ -1,35 +1,61 @@ -use crate::loader::ModelLoader; -use anyhow::Context; +use crate::loader::ModelMetadata; +use crate::types::Session; use ddddocr_core::error::{DdddError, Result, TensorError}; use ddddocr_core::utils::normalize_ocr_logits; -use ddddocr_core::{InferenceEngine, ModelMetadata, OcrEngine, OcrOutput}; -use ndarray::s; -use std::path::Path; -use tract_onnx::prelude::DatumType; -use tract_onnx::prelude::{Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tvec}; +use ddddocr_core::{InferenceEngine, OcrEngine, OcrOutput}; +use tract_onnx::prelude::{DatumType, OutletId, ShapeFact, TypedModel}; +use tract_onnx::prelude::{IntoTensor, Tensor, tvec}; +// 引入核心层的统一错误类型 +/// 明确命名为 AxisDim，代表模型某一个轴的维度特征 +#[derive(Clone, PartialEq, Eq)] +pub enum AxisDim { + /// 静态固定维度（如通道数固定为 1，高度固定为 64） + Static(usize), + /// 动态符号维度（如宽度是动态的 "image_width"） + Dynamic(String), +} +impl AxisDim { + /// 便捷方法：判断是否为动态维度 + pub fn is_dynamic(&self) -> bool { + matches!(self, AxisDim::Dynamic(_)) + } +} +/// 自定义 Debug 格式化输出，彻底融化套娃外壳，保证日志干净漂亮 +impl std::fmt::Debug for AxisDim { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AxisDim::Static(size) => write!(f, "{}", size), + AxisDim::Dynamic(expr) => write!(f, "Dynamic(\"{}\")", expr), + } + } +} +/// 模拟 Python 的 input_info 和 output_info 结构 +#[derive(Debug, Clone)] +pub struct TensorInfo { + pub name: String, + pub shape: Vec, // 既包含 Fixed 静态维度，也包含 Dynamic 动态符号 + pub data_type: DatumType, // 对应 Python 的 type +} + +/// 最终返回的模型完整信息 +#[derive(Debug, Clone)] +pub struct ModelInfo { + pub inputs: Vec, + pub outputs: Vec, + /// 硬件执行提供者（采用 Option 兼容不同底层的推理引擎） + pub providers: Option>, +} pub struct OcrSession { - pub session: RunnableModel, Graph>>, + pub session: Session, pub model_metadata: ModelMetadata, } impl OcrSession { - pub fn new

(model_path: P, model_metadata: ModelMetadata) -> Result - where - P: AsRef, - { - let session = ModelLoader::model_for_path(model_path)?.session; - Ok(Self { + pub fn new(session: Session, model_metadata: ModelMetadata) -> Self { + Self { session, model_metadata, - }) - } - - pub fn model_from_bytes(model_bytes: &[u8], model_metadata: ModelMetadata) -> Result { - let session = ModelLoader::model_from_bytes(model_bytes)?.session; - Ok(Self { - session, - model_metadata, - }) + } } } impl OcrEngine for OcrSession { @@ -40,7 +66,7 @@ impl OcrEngine for OcrSession { impl InferenceEngine for OcrSession { type Output = OcrOutput; /// 对应 Python 的 _inference - fn inference(&self, input_array: ndarray::Array4) -> Result { + fn inference(&self, input_array: ndarray::Array4) -> Result { // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 // let result = self.ocr.run(tvec!(tensor.into()))?; let tensor = Tensor::from(input_array); @@ -48,12 +74,8 @@ impl InferenceEngine for OcrSession { let mut result = self .session .run(tvec!(tensor.into())) - .map_err(|_| { - DdddError::Inference(TensorError::EngineError( - "执行模型推理失败".to_string(), - )) - })?; - // .context("执行模型推理失败")?; + .map_err(|_| TensorError::Engine("执行模型推理失败".to_string()))?; + // .context("执行模型推理失败")?; println!("模型输出原始数据: {:?}", result); // Ok(result.swap_remove(0).into_tensor()) let raw_tensor = result.swap_remove(0).into_tensor(); @@ -62,23 +84,17 @@ impl InferenceEngine for OcrSession { DatumType::I64 => { let array_d = raw_tensor .into_array::() - .map_err(|_| { - DdddError::Inference(TensorErrorReason::EngineError( - "Tract 无法获取 i64 内存视图".to_string(), - )) - })?; - // .context("Tract 无法获取 i64 内存视图")?; + .map_err(|_| TensorError::Engine("Tract 无法获取 i64 内存视图".to_string()))?; + // .context("Tract 无法获取 i64 内存视图")?; // 🌟 提前提取真实维度 let actual_shape = array_d.shape().to_vec(); // 转成标准的 Array1 传给 core let array1 = array_d .to_owned() .into_dimensionality::() - .map_err(|_| { - DdddError::Inference(TensorErrorReason::TensorDimensionMismatch { - expected: "1D 字符索引静态矩阵".to_string(), - actual: actual_shape, - }) + .map_err(|_| TensorError::DimensionMismatch { + expected: "1D 字符索引静态矩阵".to_string(), + actual: actual_shape, })?; Ok(OcrOutput::Indices(array1)) } @@ -87,18 +103,86 @@ impl InferenceEngine for OcrSession { println!("模型输出shape数据: {:?}", shape); let view = raw_tensor .to_array_view::() - .map_err(|_| { - DdddError::Inference(TensorErrorReason::EngineError( - "Tract 无法获取 f32 内存视图".to_string(), - )) - })?; + .map_err(|_| TensorError::Engine("Tract 无法获取 f32 内存视图".to_string()))?; // 1. 极其纯粹的、无拷贝的多维 Shape 压扁清洗 normalize_ocr_logits(view, shape) } _ => Err( // anyhow::anyhow!("不支持的模型输出数据类型: {:?}",raw_tensor.datum_type()) - DdddError::Inference(TensorErrorReason::UnknownOutputFormat) + TensorError::UnknownOutputFormat, ), } } } +impl OcrSession { + /// 获取模型输入的节点信息列表 + pub fn input_info(&self) -> Result> { + let model = self.session.model(); + let outlets = model.input_outlets().map_err(DdddError::new)?; + self.resolve_tensors(model, outlets) + } + + /// 获取模型输出的节点信息列表 + pub fn output_info(&self) -> Result> { + let model = self.session.model(); + let outlets = model.output_outlets().map_err(DdddError::new)?; + self.resolve_tensors(model, outlets) + } + + /// 获取模型详细元数据信息（对标 Python ddddocr 的 get_model_info） + /// 完美包容 [1, 1, 64, image_width] 这样的变长图像模型 + /// 获取模型详细元数据信息（代码更紧凑、优雅） + pub fn model_info(&self) -> Result { + Ok(ModelInfo { + inputs: self.input_info()?, + outputs: self.output_info()?, + providers: None, + }) + } + + /// 提取出来的公共转换逻辑：将一组 OutletId 解析为 TensorInfo 列表 + fn resolve_tensors(&self, model: &TypedModel, outlets: &[OutletId]) -> Result> { + outlets + .iter() + .map(|&outlet_id| { + let fact = model.outlet_fact(outlet_id).map_err(DdddError::new)?; + // .map_err(|e| { + // DdddError::InternalError(format!("解析节点 Fact 失败: {:?}", e)) + // })?; + + let shape = self.resolve_shape(&fact.shape); + let node_name = model.node(outlet_id.node).name.clone(); + + Ok(TensorInfo { + name: node_name, + shape, + data_type: fact.datum_type, + }) + }) + .collect() // 函数式声明：自动传播第一处发生的错误 + } + + /// 安全还原 Tract 维度至 Vec + fn resolve_shape(&self, shape_fact: &ShapeFact) -> Vec { + let tract_shape = shape_fact.to_tvec(); + + let resolved = tract_shape + .iter() + .map(|dim| { + // 防御性编程：必须同时满足能够转换为 i64 且 大于等于 0 + if let Ok(size) = dim.to_i64() { + if size >= 0 { + AxisDim::Static(size as usize) + } else { + // 如果 ONNX 导出时某些动态维度被标记为了 -1，安全地作为动态符号捕获 + AxisDim::Dynamic(dim.to_string()) + } + } else { + AxisDim::Dynamic(dim.to_string()) + } + }) + .collect(); + + resolved + } +} diff --git a/ddddocr-tract/src/types.rs b/ddddocr-tract/src/types.rs new file mode 100644 index 0000000..9c1fbb6 --- /dev/null +++ b/ddddocr-tract/src/types.rs @@ -0,0 +1,3 @@ +use tract_onnx::prelude::{Graph, RunnableModel, TypedFact, TypedOp}; + +pub type Session = RunnableModel, Graph>>; \ No newline at end of file diff --git a/ddddocr-tract/tests/char_slice.rs b/ddddocr-tract/tests/char_slice.rs index d858300..6c41dd3 100644 --- a/ddddocr-tract/tests/char_slice.rs +++ b/ddddocr-tract/tests/char_slice.rs @@ -2,8 +2,8 @@ use std::borrow::Cow; use std::fs::File; use std::path::Path; use anyhow::anyhow; -use ddddocr_core::ocr::metadata::Charset; -use ddddocr_core::ocr::metadata::{Normalization, Resize}; +use ddddocr_core::ocr::Charset; +use ddddocr_core::ocr::{Normalization, Resize}; pub const CHARSET_BETA: &[&str] = &[ "", "笤", "谴", "膀", "荔", "佰", "电", "臁", "矍", "同", "奇", "芄", "吠", "6", "曛", "荇", diff --git a/ddddocr-tract/tests/ocr_test.rs b/ddddocr-tract/tests/ocr_test.rs index ef40f6e..61354d5 100644 --- a/ddddocr-tract/tests/ocr_test.rs +++ b/ddddocr-tract/tests/ocr_test.rs @@ -1,15 +1,17 @@ use anyhow::Context; use ddddocr_core::det::DetectionResult; -use ddddocr_core::{DetBuilder, Detector, ModelMetadata, Ocr, Slider}; // 假设你的包名是这个 -use ddddocr_tract::{DetSession, OcrSession}; +use ddddocr_core::{Detector, ModelMetadata, Normalization, Ocr, Slider}; +// 假设你的包名是这个 +use ddddocr_tract::{DetSession, OcrSession,OcrBuilder}; use image::{DynamicImage, ImageBuffer, Luma, Rgb}; use std::fs; use std::path::Path; -use tract_onnx::prelude::{ShapeFact, TDim}; +use tract_onnx::model; mod char_slice; use char_slice::CHARSET_BETA; -use ddddocr_core::ocr::metadata::{Normalization, Resize}; +use ddddocr_core::ocr::Resize; + use ddddocr_tract::loader::ModelLoader; fn load_image>(path: P) -> anyhow::Result { @@ -103,25 +105,29 @@ fn save_rust_result(result: &ImageBuffer, Vec>, filename: &str) { } #[test] fn test_full_classification() { - // 1. 初始化模型 - let ocr = OcrSession::new( + let model = ModelLoader::model_for_path( "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx", - ModelMetadata::from_static_slice( - CHARSET_BETA, - false, - Resize::DynamicWidth(64), - 1, - Normalization::MinusOneToOne, - ), ) .expect("模型加载失败"); - + let metadata = ModelMetadata::from_static_slice( + CHARSET_BETA, + false, + Resize::DynamicWidth(64), + 1, + Normalization::MinusOneToOne, + ); + // 1. 初始化模型 + let ocr = OcrSession::new(model, metadata); // 2. 加载测试图片 let img = image::open("D:/CNWei/CNW/Rust/ddddocr-rs/samples/code2.png").expect("测试图片不存在"); // 3. 执行识别 - let result = Ocr::new(&ocr) + // let result = Ocr::new(&ocr) + // .predict(&img) + // .expect("识别过程出错") + // .into_text(); + let result = OcrBuilder::new().build(&ocr) .predict(&img) .expect("识别过程出错") .into_text(); @@ -131,7 +137,10 @@ fn test_full_classification() { } #[test] fn test_det_load() -> anyhow::Result<()> { - let det = DetSession::new("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx")?; + let det_model = + ModelLoader::model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx") + .expect("模型加载失败"); + let det = DetSession::new(det_model); let image_path = "D:/CNWei/CNW/Rust/ddddocr-rs/samples/det1.png"; let image_bytes = fs::read(image_path).map_err(|e| anyhow::anyhow!("无法读取图片 {}: {}", image_path, e))?; @@ -167,7 +176,7 @@ fn test_det_load() -> anyhow::Result<()> { #[test] fn test_real_slide_match() { - let engine = Slider::new().unwrap(); + let engine = Slider::new(); // 1. 加载你准备好的测试图 // 假设图片放在项目根目录下的 assets 文件夹 @@ -198,7 +207,7 @@ fn test_real_slide_match() { #[test] fn test_real_slide_comparison() { - let engine = Slider::new().unwrap(); + let engine = Slider::new(); // 1. 加载你准备好的测试图 // 假设图片放在项目根目录下的 assets 文件夹 @@ -236,6 +245,4 @@ fn test_resolve_shape_logic_direct() { "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx", ) .expect("建立测试模型图失败"); - let md_info = &loader.model_info().context("信息"); - println!("md_info: {:?}", md_info); }