From a3c46145740133a98a1367b8fe0a480f3ed593e2 Mon Sep 17 00:00:00 2001 From: CNWei Date: Thu, 30 Jul 2026 16:55:58 +0800 Subject: [PATCH] =?UTF-8?q?refactor(core):=20=E6=8F=90=E7=82=BC=E5=85=AC?= =?UTF-8?q?=E5=85=B1=E7=B1=BB=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将 AxisDim、TensorInfo 等公共类型下沉至 ddddocr_core::types - 项目结构优化 --- Cargo.toml | 2 +- ddddocr-core/src/det/builder.rs | 23 +-- ddddocr-core/src/det/executor.rs | 19 ++- ddddocr-core/src/lib.rs | 27 +--- ddddocr-core/src/ocr/builder.rs | 10 +- ddddocr-core/src/ocr/executor.rs | 25 ++-- ddddocr-core/src/traits.rs | 42 ++++++ ddddocr-core/src/types.rs | 46 ++++++ ddddocr-ort/src/det/session.rs | 16 +-- ddddocr-ort/src/lib.rs | 4 +- ddddocr-ort/src/loader.rs | 4 +- ddddocr-ort/src/loader/metadata.rs | 4 +- ddddocr-ort/src/loader/model.rs | 30 ++-- ddddocr-ort/src/ocr/session.rs | 201 +++++++++------------------ ddddocr-ort/tests/ocr_test.rs | 41 +++--- ddddocr-tract/src/det/session.rs | 12 +- ddddocr-tract/src/lib.rs | 10 +- ddddocr-tract/src/loader.rs | 4 +- ddddocr-tract/src/loader/metadata.rs | 4 +- ddddocr-tract/src/loader/model.rs | 20 +-- ddddocr-tract/src/ocr/session.rs | 170 +++++++++------------- ddddocr-tract/tests/ocr_test.rs | 43 +++--- 22 files changed, 349 insertions(+), 408 deletions(-) create mode 100644 ddddocr-core/src/traits.rs create mode 100644 ddddocr-core/src/types.rs diff --git a/Cargo.toml b/Cargo.toml index 587929d..ece210e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ members = [ ] [workspace.package] -version = "0.2.0" +version = "0.2.1" edition = "2024" license = "MIT OR Apache-2.0" diff --git a/ddddocr-core/src/det/builder.rs b/ddddocr-core/src/det/builder.rs index 539dff0..2469738 100644 --- a/ddddocr-core/src/det/builder.rs +++ b/ddddocr-core/src/det/builder.rs @@ -1,26 +1,11 @@ use crate::det::executor::Detector; // use ddddocr_tract::det::session::DetSession; -use crate::DetEngine; - -pub struct DetBuilder { - use_gpu: bool, - device_id: u8, -} +use crate::traits::DetEngine; +#[derive(Default)] +pub struct DetBuilder; impl DetBuilder { - fn use_gpu(mut self) -> Self { - self.use_gpu = true; - self - } - fn device_id(mut self, device_id: u8) -> Self { - self.device_id = device_id; - self - } fn build(self, session: &E) -> Detector<'_> { - Detector { - session, - use_gpu: self.use_gpu, - device_id: self.device_id, - } + Detector { session } } } diff --git a/ddddocr-core/src/det/executor.rs b/ddddocr-core/src/det/executor.rs index 169244e..22aeaaf 100644 --- a/ddddocr-core/src/det/executor.rs +++ b/ddddocr-core/src/det/executor.rs @@ -5,7 +5,8 @@ use std::fmt; // use tract_onnx::prelude::{Tensor}; // use ddddocr_tract::det::session::DetSession; -use crate::{DetEngine, DetOutput}; +use crate::{DetBuilder, DetOutput, OcrBuilder}; +use crate::traits::DetEngine; #[derive(Debug, Clone, Copy)] pub struct DetectionResult { pub x1: i32, @@ -29,21 +30,17 @@ impl fmt::Display for DetectionResult { pub struct Detector<'a> { pub(crate) session: &'a dyn DetEngine, - #[allow(dead_code)] - pub(crate) use_gpu: bool, - #[allow(dead_code)] - pub(crate) device_id: u8, } impl<'a> Detector<'a> { pub fn new(session: &'a dyn DetEngine) -> Self { - Detector { - session, - use_gpu: false, - device_id: 0, - } + Detector { session } } - + pub fn builder() -> DetBuilder { + DetBuilder::default() + } +} +impl<'a> Detector<'a> { pub fn predict(&self, image: &DynamicImage) -> Result> { // Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑 Ok(self.get_bbox(image)?) diff --git a/ddddocr-core/src/lib.rs b/ddddocr-core/src/lib.rs index a103779..09a5526 100644 --- a/ddddocr-core/src/lib.rs +++ b/ddddocr-core/src/lib.rs @@ -3,14 +3,14 @@ pub mod error; mod ocr; mod slide; pub mod utils; - -use error::{Result, TensorError}; -use std::path::Path; +pub mod types; +pub mod traits; pub use crate::det::{DetBuilder, DetectionResult, Detector}; pub use crate::ocr::{Charset, ModelMetadata, Normalization, Ocr, OcrBuilder, OcrResult, Resize}; pub use crate::slide::{SlideResult, Slider}; + // DetSession pub enum OcrOutput { @@ -22,24 +22,3 @@ pub enum DetOutput { Detection(ndarray::Array3), // 拥有完整所有权的 2维矩阵，可任意传递和返回 } -/// 核心层定义的统一推理引擎接口。 -/// 未来的 ddddocr-tract 和 ddddocr-ort 都必须实现这个 Trait -pub trait InferenceEngine { - /// 关联类型：具体的 Session 需要声明自己到底产出什么枚举 - type Output; - fn inference(&self, input_array: ndarray::Array4) -> Result; -} - -pub trait OcrEngine: InferenceEngine { - fn metadata(&self) -> &ModelMetadata; -} - -pub trait DetEngine: InferenceEngine {} - -pub trait ModelBuilder { - type Session; - type Error; - fn model_for_path>(&self,model_path: P) -> Result; - - fn model_from_bytes(&self,model_bytes: &[u8]) -> Result; -} diff --git a/ddddocr-core/src/ocr/builder.rs b/ddddocr-core/src/ocr/builder.rs index 50c1b7c..628448a 100644 --- a/ddddocr-core/src/ocr/builder.rs +++ b/ddddocr-core/src/ocr/builder.rs @@ -1,9 +1,9 @@ use crate::ocr::executor::Ocr; // use ddddocr_tract::session::OcrSession; -use crate::OcrEngine; +use crate::traits::OcrEngine; use crate::ocr::color_filter::ColorFilter; use crate::ocr::token_filter::TokenFilter; - +#[derive(Default)] pub struct OcrBuilder { /// 是否修复PNG格式问题 png_fix: bool, @@ -49,14 +49,14 @@ impl OcrBuilder { self.charset_restrict = Some(Box::new(restrict)); self } - pub fn build(self, session: &E) -> Ocr<'_> { + pub fn runner(self, runtime: &E) -> Ocr<'_> { // 1. 原地解析颜色过滤器 let final_color_ranges = match &self.color_filter { Some(filter) => filter.collect_to_vec(), None => Ok(None), }; // 2. 原地解析字符集过滤 - let tokens = &session.metadata().charset.tokens; + let tokens = &runtime.metadata().charset.tokens; let final_charset_indices = match &self.charset_restrict { Some(restrict) => restrict.apply_to_charset(tokens), None => None, @@ -64,7 +64,7 @@ impl OcrBuilder { // Ocr::new(session, self) Ocr { - session, + runtime, png_fix: self.png_fix, // 原地解构出来 probability: self.probability, final_color_ranges, diff --git a/ddddocr-core/src/ocr/executor.rs b/ddddocr-core/src/ocr/executor.rs index f51a04f..ea8aaa6 100644 --- a/ddddocr-core/src/ocr/executor.rs +++ b/ddddocr-core/src/ocr/executor.rs @@ -16,7 +16,8 @@ use ndarray::ArrayView2; // Logits(ndarray::Array2), // 拥有完整所有权的 2维矩阵，可任意传递和返回 // } use crate::error::{ImagePreprocessError, Result, TensorError}; -use crate::{OcrEngine, OcrOutput}; +use crate::{OcrBuilder, OcrOutput}; +use crate::traits::OcrEngine; use tracing::{ warn}; #[derive(Debug, Clone)] pub enum OcrResult { @@ -103,7 +104,7 @@ impl fmt::Display for OcrResult { } pub struct Ocr<'a> { - pub(crate) session: &'a dyn OcrEngine, + pub(crate) runtime: &'a dyn OcrEngine, pub(crate) png_fix: bool, pub(crate) probability: bool, /// 颜色过滤：保留的颜色列表 @@ -116,15 +117,19 @@ pub struct Ocr<'a> { impl<'a> Ocr<'a> { // 初始化任务，设置默认参数 - pub fn new(session: &'a dyn OcrEngine) -> Self { + pub fn new(runtime: &'a dyn OcrEngine) -> Self { Ocr { - session, + runtime, png_fix: false, // 默认值 probability: false, final_color_ranges: Ok(None), final_charset_indices: None, } } + pub fn builder() -> OcrBuilder { + OcrBuilder::default() + } + } impl<'a> Ocr<'a> { pub fn predict(&self, image: &DynamicImage) -> Result { @@ -157,7 +162,7 @@ impl<'a> Ocr<'a> { }; let tensor = self.preprocess_image(&img_cow)?; - let raw_tensor = self.session.inference(tensor)?; + let raw_tensor = self.runtime.inference(tensor)?; // 3. 后处理分流：直接返回 OcrResult // let ocr_output = match raw_tensor.datum_type() { @@ -178,7 +183,7 @@ impl<'a> Ocr<'a> { /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 fn preprocess_image(&self, img: &DynamicImage) -> Result,ImagePreprocessError> { // 1. 获取模型元数据配置 - let meta = self.session.metadata(); + let meta = self.runtime.metadata(); let norm = &meta.normalization; // 获取归一化器 // A. 修复 PNG 透明背景 (内部逻辑你之前已实现) @@ -326,7 +331,7 @@ impl<'a> Ocr<'a> { } impl<'a> Ocr<'a> { fn is_valid_indices(&self, idx: usize) -> bool { - if idx >= self.session.metadata().charset.size() { + if idx >= self.runtime.metadata().charset.size() { return false; } @@ -338,7 +343,7 @@ impl<'a> Ocr<'a> { /// 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 /// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！ pub fn valid_tokens(&self) -> Vec<&str> { - let charset = &self.session.metadata().charset; + let charset = &self.runtime.metadata().charset; let tokens = &charset.tokens; match &self.final_charset_indices { Some(indices) => indices @@ -352,7 +357,7 @@ impl<'a> Ocr<'a> { pub fn valid_size(&self) -> usize { match &self.final_charset_indices { Some(indices) => indices.len(), - None => self.session.metadata().charset.tokens.len(), + None => self.runtime.metadata().charset.tokens.len(), } } /// 变体 B 核心处理器：单次遍历 2D 视图，融合计算 Softmax、Argmax、置信度并输出概率大包 @@ -500,7 +505,7 @@ impl<'a> Ocr<'a> { /// 获取有效字符索引列表 (用于外部验证或过滤) fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String { println!("indices模型输出原始数据: {:?}", predicted_indices); - let charset = &self.session.metadata().charset; + let charset = &self.runtime.metadata().charset; let tokens = &charset.tokens; // let valid_indices = &charset.valid_indices; diff --git a/ddddocr-core/src/traits.rs b/ddddocr-core/src/traits.rs new file mode 100644 index 0000000..d720ea0 --- /dev/null +++ b/ddddocr-core/src/traits.rs @@ -0,0 +1,42 @@ +use crate::error::TensorError; +use crate::types::{ModelInfo, TensorInfo}; +use crate::{DetOutput, ModelMetadata, OcrOutput}; +use std::path::Path; + + +pub trait Info { + fn input_info(&self) -> crate::error::Result>; + fn output_info(&self) -> crate::error::Result>; + fn model_info(&self) -> crate::error::Result; +} + +/// 核心层定义的统一推理引擎接口。 +/// 未来的 ddddocr-tract 和 ddddocr-ort 都必须实现这个 Trait +pub trait InferenceEngine { + /// 关联类型：具体的 Session 需要声明自己到底产出什么枚举 + type Output; + fn inference( + &self, + input_array: ndarray::Array4, + ) -> crate::error::Result; +} + +pub trait OcrEngine: InferenceEngine + Info { + fn metadata(&self) -> &ModelMetadata; +} + +pub trait DetEngine: InferenceEngine {} + +pub trait Loader { + type Session; + type Error; + fn build_for_path>( + &self, + model_path: P, + ) -> crate::error::Result; + + fn build_from_bytes( + &self, + model_bytes: &[u8], + ) -> crate::error::Result; +} diff --git a/ddddocr-core/src/types.rs b/ddddocr-core/src/types.rs new file mode 100644 index 0000000..3984400 --- /dev/null +++ b/ddddocr-core/src/types.rs @@ -0,0 +1,46 @@ +#[derive(Debug,Clone)] +pub enum TensorType{ + F32, + I64, + Other +} +/// 明确命名为 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 tensor_type: TensorType, // 对应 Python 的 type +} + +/// 最终返回的模型完整信息 +#[derive(Debug, Clone)] +pub struct ModelInfo { + pub inputs: Vec, + pub outputs: Vec, + /// 硬件执行提供者（采用 Option 兼容不同底层的推理引擎） + pub providers: Option>, +} \ No newline at end of file diff --git a/ddddocr-ort/src/det/session.rs b/ddddocr-ort/src/det/session.rs index 575cc4d..93bcf21 100644 --- a/ddddocr-ort/src/det/session.rs +++ b/ddddocr-ort/src/det/session.rs @@ -1,24 +1,24 @@ use crate::types::Session; +use ddddocr_core::DetOutput; use ddddocr_core::error::{Result, TensorError}; -use ddddocr_core::{DetEngine, DetOutput, InferenceEngine}; +use ddddocr_core::traits::{DetEngine, InferenceEngine}; use ndarray::Ix3; use ort::inputs; use ort::value::TensorRef; // use tract_onnx::prelude::{tvec, IntoTensor, Tensor}; #[derive(Debug)] -pub struct DetSession { +pub struct DetRuntime { pub session: Session, } -impl DetSession { +impl DetRuntime { pub fn new(session: Session) -> Self { Self { session } } - } -impl InferenceEngine for DetSession { +impl InferenceEngine for DetRuntime { type Output = DetOutput; // 明确绑定 OCR 小枚举 fn inference(&self, input_array: ndarray::Array4) -> Result { // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 @@ -28,7 +28,6 @@ impl InferenceEngine for DetSession { .lock() .map_err(|_| TensorError::Engine("获取 Session 锁失败 (Poisoned)".to_string()))?; - let result = session_guard .run(inputs![TensorRef::from_array_view(&input_array).map_err( |e| TensorError::Engine(format!("构建输入失败: {e}")) @@ -45,8 +44,7 @@ impl InferenceEngine for DetSession { TensorError::Engine("Tract 实体张量无法转换为 ndarray::ArrayD".to_string()) })?; // 提前利用克隆(Clone)备份好当前未转维度前的真实 shape (Vec) - let shape_vec: Vec = - shape_ref.to_vec().iter().map(|v| *v as usize).collect(); + let shape_vec: Vec = shape_ref.to_vec().iter().map(|v| *v as usize).collect(); let shape_vec_slice = shape_vec.as_slice(); let view = ndarray::ArrayViewD::from_shape(shape_vec_slice, slice) @@ -64,4 +62,4 @@ impl InferenceEngine for DetSession { } } -impl DetEngine for DetSession {} +impl DetEngine for DetRuntime {} diff --git a/ddddocr-ort/src/lib.rs b/ddddocr-ort/src/lib.rs index 4940d3f..3c0d256 100644 --- a/ddddocr-ort/src/lib.rs +++ b/ddddocr-ort/src/lib.rs @@ -5,5 +5,5 @@ mod types; pub use ddddocr_core::{SlideResult, Slider,OcrBuilder}; -pub use det::session::DetSession; -pub use ocr::session::OcrSession; \ No newline at end of file +pub use det::session::DetRuntime; +pub use ocr::session::OcrRuntime; \ No newline at end of file diff --git a/ddddocr-ort/src/loader.rs b/ddddocr-ort/src/loader.rs index ec95941..4bdafe9 100644 --- a/ddddocr-ort/src/loader.rs +++ b/ddddocr-ort/src/loader.rs @@ -3,5 +3,5 @@ mod metadata; mod model; pub use error::{Error, ParseError, Result}; -pub use metadata::{ModelMetadataDto, NormalizationDto, TractModelMetadata}; -pub use model::OrtModelLoader; \ No newline at end of file +pub use metadata::{ModelMetadataDto, NormalizationDto, Metadata}; +pub use model::ModelLoader; \ No newline at end of file diff --git a/ddddocr-ort/src/loader/metadata.rs b/ddddocr-ort/src/loader/metadata.rs index fb067dc..3f2dd1e 100644 --- a/ddddocr-ort/src/loader/metadata.rs +++ b/ddddocr-ort/src/loader/metadata.rs @@ -42,7 +42,7 @@ fn default_normalization() -> NormalizationDto { } /// Tract 专属扩展trait 或 工具函数 -pub trait TractModelMetadata: Sized { +pub trait Metadata: Sized { fn from_json_str(json_str: &str) -> Result; /// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载） fn from_json_bytes(bytes: &[u8]) -> Result { @@ -50,7 +50,7 @@ pub trait TractModelMetadata: Sized { Self::from_json_str(json_str) } } -impl TractModelMetadata for ModelMetadata { +impl Metadata for ModelMetadata { // --- 优雅的工厂模式构造器 --- fn from_json_str(json_str: &str) -> Result { let dto: ModelMetadataDto = serde_json::from_str(json_str)?; diff --git a/ddddocr-ort/src/loader/model.rs b/ddddocr-ort/src/loader/model.rs index 291f4a6..3afc75c 100644 --- a/ddddocr-ort/src/loader/model.rs +++ b/ddddocr-ort/src/loader/model.rs @@ -1,27 +1,27 @@ use crate::loader::Error; use crate::loader::error::{BuildError, ParseError, Result}; use crate::types::Session; -use ddddocr_core::ModelBuilder; +use ddddocr_core::traits::Loader; use ort::session::Session as OrtSession; use ort::session::builder::SessionBuilder; use std::sync::{Arc, Mutex}; -pub struct OrtModelLoader; - -impl OrtModelLoader { - /// 获取针对 ORT 后端的链式构建器 - pub fn builder() -> OrtModelBuilder { - OrtModelBuilder::default() - } -} +// pub struct OrtModelLoader; +// +// impl OrtModelLoader { +// /// 获取针对 ORT 后端的链式构建器 +// pub fn builder() -> OrtModelBuilder { +// OrtModelBuilder::default() +// } +// } /// ORT 专用的链式构建器 #[derive(Debug, Clone)] -pub struct OrtModelBuilder { +pub struct ModelLoader { use_gpu: bool, device_id: i32, intra_threads: Option, } -impl Default for OrtModelBuilder { +impl Default for ModelLoader { fn default() -> Self { Self { use_gpu: false, @@ -30,7 +30,7 @@ impl Default for OrtModelBuilder { } } } -impl OrtModelBuilder { +impl ModelLoader { /// 开启或关闭 GPU 加速 pub fn use_gpu(mut self, enable: bool) -> Self { self.use_gpu = enable; @@ -83,11 +83,11 @@ impl OrtModelBuilder { } } -impl ModelBuilder for OrtModelBuilder { +impl Loader for ModelLoader { type Session = Session; type Error = Error; - fn model_for_path

(&self, model_path: P) -> Result + fn build_for_path

(&self, model_path: P) -> Result where P: AsRef, { @@ -102,7 +102,7 @@ impl ModelBuilder for OrtModelBuilder { Ok(Arc::new(Mutex::new(session))) // 这里的session需要包装下 } /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用） - fn model_from_bytes(&self, model_bytes: &[u8]) -> Result { + fn build_from_bytes(&self, model_bytes: &[u8]) -> Result { let mut builder = self.create_session_builder()?; let session = builder diff --git a/ddddocr-ort/src/ocr/session.rs b/ddddocr-ort/src/ocr/session.rs index 204e580..f2e5e98 100644 --- a/ddddocr-ort/src/ocr/session.rs +++ b/ddddocr-ort/src/ocr/session.rs @@ -1,70 +1,69 @@ use crate::types::Session; use ddddocr_core::ModelMetadata; +use ddddocr_core::OcrOutput; use ddddocr_core::error::{DdddError, Result, TensorError}; +use ddddocr_core::traits::{InferenceEngine, Info, OcrEngine}; +use ddddocr_core::types::{AxisDim, ModelInfo, TensorInfo, TensorType}; use ddddocr_core::utils::normalize_ocr_logits; -use ddddocr_core::{InferenceEngine, OcrEngine, OcrOutput}; use ort::inputs; use ort::value::{TensorElementType, TensorRef}; use std::sync::Mutex; // 引入核心层的统一错误类型 /// 明确命名为 AxisDim，代表模型某一个轴的维度特征 -#[derive(Clone, PartialEq, Eq)] -pub enum AxisDim { - /// 静态固定维度（如通道数固定为 1，高度固定为 64） - Static(usize), - /// 动态符号维度（如宽度是动态的 "image_width"） - Dynamic(String), -} +// #[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), - } - } -} +// 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: TensorElementType, // 对应 Python 的 type -} - -/// 最终返回的模型完整信息 -#[derive(Debug, Clone)] -pub struct ModelInfo { - pub inputs: Vec, - pub outputs: Vec, - /// 硬件执行提供者（采用 Option 兼容不同底层的推理引擎） - pub providers: Option>, -} -pub struct OcrSession { +// #[derive(Debug, Clone)] +// pub struct TensorInfo { +// pub name: String, +// pub shape: Vec, // 既包含 Fixed 静态维度，也包含 Dynamic 动态符号 +// pub data_type: TensorElementType, // 对应 Python 的 type +// } +// +// /// 最终返回的模型完整信息 +// #[derive(Debug, Clone)] +// pub struct ModelInfo { +// pub inputs: Vec, +// pub outputs: Vec, +// /// 硬件执行提供者（采用 Option 兼容不同底层的推理引擎） +// pub providers: Option>, +// } +pub struct OcrRuntime { pub session: Session, - pub model_metadata: ModelMetadata, + pub metadata: ModelMetadata, } -impl OcrSession { - pub fn new(session: Session, model_metadata: ModelMetadata) -> Self { - Self { - session, - model_metadata, - } +impl OcrRuntime { + pub fn new(session: Session, metadata: ModelMetadata) -> Self { + Self { session, metadata } } } -impl OcrEngine for OcrSession { +impl OcrEngine for OcrRuntime { fn metadata(&self) -> &ModelMetadata { - &self.model_metadata + &self.metadata } } -impl InferenceEngine for OcrSession { +impl InferenceEngine for OcrRuntime { type Output = OcrOutput; /// 对应 Python 的 _inference fn inference(&self, input_array: ndarray::Array4) -> Result { @@ -76,18 +75,6 @@ impl InferenceEngine for OcrSession { .lock() .map_err(|_| TensorError::Engine("获取 Session 锁失败 (Poisoned)".to_string()))?; - // // 2. 获取输入节点名称 - // let input_name = session_guard - // .inputs() - // .first() - // .map(|i| i.name()) - // .unwrap_or("input"); - // - // // 3. 在 session_guard (&mut Session) 上调用 run - // let outputs = session_guard - // .run(inputs![TensorRef::from_array_view(&input_array).map_err(|e| TensorError::Engine(format!("构建输入失败: {e}")))? ) - // .map_err(|e| TensorError::Engine(format!("执行模型推理失败: {e}")))?; - let result = session_guard .run(inputs![TensorRef::from_array_view(&input_array).map_err( |e| TensorError::Engine(format!("构建输入失败: {e}")) @@ -97,10 +84,7 @@ impl InferenceEngine for OcrSession { println!("模型输出原始数据: {:?}", result); // Ok(result.swap_remove(0).into_tensor()) let raw_value = &result[0]; - // let dtype = raw_tensor - // .dtype(); - // .map_err(|e| TensorError::Engine(format!("无法读取输出数据类型: {e}")))?; - // 在引擎内部消化掉 DatumType 强耦合 + match raw_value.dtype().tensor_type().unwrap() { TensorElementType::Int64 => { let (array_d, slice) = raw_value @@ -108,7 +92,7 @@ impl InferenceEngine for OcrSession { .map_err(|_| TensorError::Engine("Tract 无法获取 i64 内存视图".to_string()))?; // .context("Tract 无法获取 i64 内存视图")?; - // 🌟 提前提取真实维度 + // 提前提取真实维度 let actual_shape = array_d .to_vec() .iter() @@ -151,75 +135,16 @@ impl InferenceEngine for OcrSession { } } } -// 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 -// } -// } +impl Info for OcrRuntime { + fn input_info(&self) -> Result> { + todo!() + } + + fn output_info(&self) -> Result> { + todo!() + } + + fn model_info(&self) -> Result { + todo!() + } +} diff --git a/ddddocr-ort/tests/ocr_test.rs b/ddddocr-ort/tests/ocr_test.rs index 8eff838..d9411e8 100644 --- a/ddddocr-ort/tests/ocr_test.rs +++ b/ddddocr-ort/tests/ocr_test.rs @@ -1,8 +1,9 @@ use anyhow::Context; -use ddddocr_core::{DetectionResult, ModelBuilder}; +use ddddocr_core::{DetectionResult, Ocr}; +use ddddocr_core::traits::Loader; use ddddocr_core::{Detector, ModelMetadata, Normalization, Slider}; // 假设你的包名是这个 -use ddddocr_ort::{DetSession, OcrBuilder, OcrSession}; +use ddddocr_ort::{DetRuntime, OcrBuilder, OcrRuntime}; use image::{DynamicImage, ImageBuffer, Luma, Rgb}; use std::fs; use std::path::Path; @@ -11,7 +12,7 @@ mod char_slice; use char_slice::CHARSET_BETA; use ddddocr_core::Resize; -use ddddocr_ort::loader::OrtModelLoader; +use ddddocr_ort::loader::ModelLoader as OrtModelLoader; fn load_image>(path: P) -> anyhow::Result { // 1. 先将泛型转为具体的 &Path 引用 @@ -104,8 +105,9 @@ fn save_rust_result(result: &ImageBuffer, Vec>, filename: &str) { } #[test] fn test_full_classification() { - let model = OrtModelLoader::builder().use_gpu(true) - .model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx") + let model = OrtModelLoader::default().use_gpu(false) + .build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx") + // .build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_old.onnx") .expect("模型加载失败"); let metadata = ModelMetadata::from_static_slice( CHARSET_BETA, @@ -115,7 +117,7 @@ fn test_full_classification() { Normalization::MinusOneToOne, ); // 1. 初始化模型 - let ocr = OcrSession::new(model, metadata); + let ocr = OcrRuntime::new(model, metadata); // 2. 加载测试图片 let img = image::open("D:/CNWei/CNW/Rust/ddddocr-rs/samples/code2.png").expect("测试图片不存在"); @@ -125,21 +127,24 @@ fn test_full_classification() { // .predict(&img) // .expect("识别过程出错") // .into_text(); - let result = OcrBuilder::new() - .build(&ocr) - .predict(&img) - .expect("识别过程出错") - .into_text(); + // let result = OcrBuilder::new() + // .build(&ocr) + // .predict(&img) + // .expect("识别过程出错") + // .into_text(); + let res=Ocr::builder().runner(&ocr).predict(&img).expect("s").into_text(); - println!("识别结果: {}", result); - assert!(!result.is_empty()); + // println!("识别结果: {}", result); + println!("识别结果: {}", res); + // assert!(!result.is_empty()); + assert!(!res.is_empty()); } #[test] fn test_det_load() -> anyhow::Result<()> { - let det_model = OrtModelLoader::builder() - .model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx") + let det_model = OrtModelLoader::default() + .build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx") .expect("模型加载失败"); - let det = DetSession::new(det_model); + let det = DetRuntime::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))?; @@ -239,8 +244,8 @@ fn test_real_slide_comparison() { #[test] fn test_resolve_shape_logic_direct() { // 创建一个哑 ModelLoader 实例（session 用不上，因为我们直接测私有方法） - let loader = OrtModelLoader::builder() - .model_for_path( + let loader = OrtModelLoader::default() + .build_for_path( // "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx", "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx", ) diff --git a/ddddocr-tract/src/det/session.rs b/ddddocr-tract/src/det/session.rs index 0c7404d..36072b6 100644 --- a/ddddocr-tract/src/det/session.rs +++ b/ddddocr-tract/src/det/session.rs @@ -1,22 +1,22 @@ use crate::types::Session; use ddddocr_core::error::{Result, TensorError}; -use ddddocr_core::{DetEngine, DetOutput, InferenceEngine}; +use ddddocr_core::{ DetOutput}; +use ddddocr_core::traits::{DetEngine, InferenceEngine}; use ndarray::Ix3; -// use tract_onnx::prelude::{tvec, IntoTensor, Tensor}; use tract_onnx::prelude::*; #[derive(Debug)] -pub struct DetSession { +pub struct DetRuntime { pub session: Session, } -impl DetSession { +impl DetRuntime { pub fn new(session: Session) -> Self { Self { session } } } -impl InferenceEngine for DetSession { +impl InferenceEngine for DetRuntime { type Output = DetOutput; // 明确绑定 OCR 小枚举 fn inference(&self, input_array: ndarray::Array4) -> Result { // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 @@ -49,4 +49,4 @@ impl InferenceEngine for DetSession { } } -impl DetEngine for DetSession {} +impl DetEngine for DetRuntime {} diff --git a/ddddocr-tract/src/lib.rs b/ddddocr-tract/src/lib.rs index c67f953..0b1d679 100644 --- a/ddddocr-tract/src/lib.rs +++ b/ddddocr-tract/src/lib.rs @@ -1,9 +1,11 @@ mod det; +mod error; pub mod loader; mod ocr; mod types; -mod error; -pub use ddddocr_core::{SlideResult, Slider,OcrBuilder}; -pub use det::session::DetSession; -pub use ocr::session::OcrSession; +pub use ddddocr_core::{ + DetectionResult, Detector, ModelMetadata, Normalization, Ocr, OcrBuilder, SlideResult, Slider, +}; +pub use det::session::DetRuntime; +pub use ocr::session::OcrRuntime; diff --git a/ddddocr-tract/src/loader.rs b/ddddocr-tract/src/loader.rs index 8690dea..dcd4dd7 100644 --- a/ddddocr-tract/src/loader.rs +++ b/ddddocr-tract/src/loader.rs @@ -3,5 +3,5 @@ mod metadata; mod model; pub use error::{Error, ParseError, Result}; -pub use metadata::{ModelMetadataDto, NormalizationDto, TractModelMetadata}; -pub use model::TractModelLoader; +pub use metadata::{ModelMetadataDto, NormalizationDto, Metadata}; +pub use model::ModelLoader; diff --git a/ddddocr-tract/src/loader/metadata.rs b/ddddocr-tract/src/loader/metadata.rs index fb067dc..3f2dd1e 100644 --- a/ddddocr-tract/src/loader/metadata.rs +++ b/ddddocr-tract/src/loader/metadata.rs @@ -42,7 +42,7 @@ fn default_normalization() -> NormalizationDto { } /// Tract 专属扩展trait 或 工具函数 -pub trait TractModelMetadata: Sized { +pub trait Metadata: Sized { fn from_json_str(json_str: &str) -> Result; /// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载） fn from_json_bytes(bytes: &[u8]) -> Result { @@ -50,7 +50,7 @@ pub trait TractModelMetadata: Sized { Self::from_json_str(json_str) } } -impl TractModelMetadata for ModelMetadata { +impl Metadata for ModelMetadata { // --- 优雅的工厂模式构造器 --- fn from_json_str(json_str: &str) -> Result { let dto: ModelMetadataDto = serde_json::from_str(json_str)?; diff --git a/ddddocr-tract/src/loader/model.rs b/ddddocr-tract/src/loader/model.rs index c94f74c..1d5688b 100644 --- a/ddddocr-tract/src/loader/model.rs +++ b/ddddocr-tract/src/loader/model.rs @@ -1,26 +1,18 @@ use crate::loader::error::{Error, ParseError, Result}; use crate::types::Session; -use ddddocr_core::ModelBuilder; +use ddddocr_core::traits::Loader; use std::io::Cursor; use tract_linalg::multithread::{Executor, set_default_executor}; use tract_onnx::onnx; use tract_onnx::prelude::*; -pub struct TractModelLoader; -impl TractModelLoader { - /// 获取针对 Tract 后端的链式构建器 - pub fn builder() -> TractModelBuilder { - TractModelBuilder::default() - } -} - /// Tract 专用的链式构建器 #[derive(Debug, Clone, Default)] -pub struct TractModelBuilder { +pub struct ModelLoader { num_threads: Option, } -impl TractModelBuilder { +impl ModelLoader { /// 可选扩展：设置 CPU 线程数（不提供任何 GPU 相关的 API） pub fn num_threads(mut self, threads: usize) -> Self { self.num_threads = Some(threads); @@ -41,10 +33,10 @@ impl TractModelBuilder { } } -impl ModelBuilder for TractModelBuilder { +impl Loader for ModelLoader { type Session = Session; type Error = Error; - fn model_for_path

(&self, model_path: P) -> Result + fn build_for_path

(&self, model_path: P) -> Result where P: AsRef, { @@ -64,7 +56,7 @@ impl ModelBuilder for TractModelBuilder { Ok(session) } /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用） - fn model_from_bytes(&self, model_bytes: &[u8]) -> Result { + fn build_from_bytes(&self, model_bytes: &[u8]) -> Result { self.setup_tract_threads(); // 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read） let mut cursor = Cursor::new(model_bytes); diff --git a/ddddocr-tract/src/ocr/session.rs b/ddddocr-tract/src/ocr/session.rs index 4155985..6ef947b 100644 --- a/ddddocr-tract/src/ocr/session.rs +++ b/ddddocr-tract/src/ocr/session.rs @@ -1,69 +1,80 @@ use crate::types::Session; use ddddocr_core::ModelMetadata; +use ddddocr_core::OcrOutput; use ddddocr_core::error::{DdddError, Result, TensorError}; +use ddddocr_core::traits::{InferenceEngine, Info, OcrEngine}; +use ddddocr_core::types::{AxisDim, ModelInfo, TensorInfo, TensorType}; use ddddocr_core::utils::normalize_ocr_logits; -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 struct OcrRuntime { pub session: Session, - pub model_metadata: ModelMetadata, + pub metadata: ModelMetadata, } -impl OcrSession { - pub fn new(session: Session, model_metadata: ModelMetadata) -> Self { - Self { - session, - model_metadata, - } +impl OcrRuntime { + pub fn new(session: Session, metadata: ModelMetadata) -> Self { + Self { session, metadata } + } + /// 获取模型输入的节点信息列表 + + /// 提取出来的公共转换逻辑：将一组 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(); + let tensor_type = match fact.datum_type { + DatumType::F32 => TensorType::F32, + DatumType::I64 => TensorType::I64, + _ => TensorType::Other, + }; + + Ok(TensorInfo { + name: node_name, + shape, + tensor_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 } } -impl OcrEngine for OcrSession { +impl OcrEngine for OcrRuntime { fn metadata(&self) -> &ModelMetadata { - &self.model_metadata + &self.metadata } } -impl InferenceEngine for OcrSession { +impl InferenceEngine for OcrRuntime { type Output = OcrOutput; /// 对应 Python 的 _inference fn inference(&self, input_array: ndarray::Array4) -> Result { @@ -115,16 +126,15 @@ impl InferenceEngine for OcrSession { } } } -impl OcrSession { - /// 获取模型输入的节点信息列表 - pub fn input_info(&self) -> Result> { +impl Info for OcrRuntime { + 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> { + fn output_info(&self) -> Result> { let model = self.session.model(); let outlets = model.output_outlets().map_err(DdddError::new)?; self.resolve_tensors(model, outlets) @@ -133,57 +143,11 @@ impl OcrSession { /// 获取模型详细元数据信息（对标 Python ddddocr 的 get_model_info） /// 完美包容 [1, 1, 64, image_width] 这样的变长图像模型 /// 获取模型详细元数据信息（代码更紧凑、优雅） - pub fn model_info(&self) -> Result { + 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/tests/ocr_test.rs b/ddddocr-tract/tests/ocr_test.rs index 2f7eb5c..26a3616 100644 --- a/ddddocr-tract/tests/ocr_test.rs +++ b/ddddocr-tract/tests/ocr_test.rs @@ -1,8 +1,9 @@ use anyhow::Context; -use ddddocr_core::{DetectionResult, ModelBuilder}; -use ddddocr_core::{Detector, ModelMetadata, Normalization, Slider}; +use ddddocr_core::traits::Loader; +use ddddocr_tract::{DetectionResult, Ocr}; +use ddddocr_tract::{Detector, ModelMetadata, Normalization, Slider}; // 假设你的包名是这个 -use ddddocr_tract::{DetSession, OcrBuilder, OcrSession}; +use ddddocr_tract::{DetRuntime, OcrRuntime}; use image::{DynamicImage, ImageBuffer, Luma, Rgb}; use std::fs; use std::path::Path; @@ -11,7 +12,7 @@ mod char_slice; use char_slice::CHARSET_BETA; use ddddocr_core::Resize; -use ddddocr_tract::loader::{ TractModelLoader}; +use ddddocr_tract::loader::ModelLoader as TractModelLoader; fn load_image>(path: P) -> anyhow::Result { // 1. 先将泛型转为具体的 &Path 引用 @@ -104,9 +105,10 @@ fn save_rust_result(result: &ImageBuffer, Vec>, filename: &str) { } #[test] fn test_full_classification() { - let model = TractModelLoader::builder() - .model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx") + let session = TractModelLoader::default() + .build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx") .expect("模型加载失败"); + let metadata = ModelMetadata::from_static_slice( CHARSET_BETA, false, @@ -115,18 +117,18 @@ fn test_full_classification() { Normalization::MinusOneToOne, ); // 1. 初始化模型 - let ocr = OcrSession::new(model, metadata); + let ocr_runtime = OcrRuntime::new(session, 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_runtime) // .predict(&img) // .expect("识别过程出错") // .into_text(); - let result = OcrBuilder::new() - .build(&ocr) + let result = Ocr::builder() + .runner(&ocr_runtime) .predict(&img) .expect("识别过程出错") .into_text(); @@ -136,11 +138,10 @@ fn test_full_classification() { } #[test] fn test_det_load() -> anyhow::Result<()> { - let det_model = - TractModelLoader::builder() - .model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx") - .expect("模型加载失败"); - let det = DetSession::new(det_model); + let det_model = TractModelLoader::default() + .build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx") + .expect("模型加载失败"); + let det = DetRuntime::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))?; @@ -240,11 +241,11 @@ fn test_real_slide_comparison() { #[test] fn test_resolve_shape_logic_direct() { // 创建一个哑 ModelLoader 实例（session 用不上，因为我们直接测私有方法） - let loader = TractModelLoader::builder() - .model_for_path( - // "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx", - "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx", - ) - .expect("建立测试模型图失败"); + let loader = TractModelLoader::default() + .build_for_path( + // "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx", + "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx", + ) + .expect("建立测试模型图失败"); println!("{:?}", loader.model().inputs); }