diff --git a/src/algo/mod.rs b/src/algo/mod.rs new file mode 100644 index 0000000..16155cf --- /dev/null +++ b/src/algo/mod.rs @@ -0,0 +1,3 @@ +mod slide; + +pub use slide::{SlideResult, Slider}; diff --git a/src/models/slide.rs b/src/algo/slide.rs similarity index 100% rename from src/models/slide.rs rename to src/algo/slide.rs diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..39cd5de --- /dev/null +++ b/src/error.rs @@ -0,0 +1,18 @@ +pub(crate) const MODEL_DOWNLOAD_HELP: &str = "\ +================================================================================ +[ddddocr-rust] 错误：未找到默认的模型文件！ +-------------------------------------------------------------------------------- +由于打包体积限制，本库未内置 ONNX 模型。请按照以下步骤操作： + +1. 前往官方 GitHub 下载对应的模型权重： + - OCR 模型: https://github.com/sml2h3/ddddocr/raw/master/ddddocr/common_sml2h3_f32.onnx + - DET 模型: https://github.com/sml2h3/ddddocr/raw/master/ddddocr/common_det.onnx + +2. 配置加载方式（二选一）： + A. 【推荐】设置环境变量指向您下载的文件： + Linux/macOS: export DDDD_OCR_MODEL=\"/path/to/common_sml2h3_f32.onnx\" + Windows (CMD): set DDDD_OCR_MODEL=C:\\path\\to\\common_sml2h3_f32.onnx + Windows (PowerShell): $env:DDDD_OCR_MODEL=\"C:\\path\\to\\common_sml2h3_f32.onnx\" + + B. 或者直接将模型文件重命名并放置在您运行程序的“当前工作目录”或“可执行文件同级目录”下。 +================================================================================"; \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 24f37c0..f5ffe9c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,24 +1,14 @@ -mod charset; +mod algo; mod error; -mod model_metadata; pub mod models; pub mod utils; -use anyhow::{Context, Result, anyhow}; -use image::DynamicImage; -use std::fmt::{Display, Formatter}; -pub use crate::models::det::{Detector,DetectionResult}; -pub use crate::models::ocr::{Ocr, OcrPredictor, OcrResult}; -pub use crate::models::slide::{Slider, SlideResult}; -use std::path::{Path, PathBuf}; - -// 关键点：直接使用 tract 重导出的 ndarray -use crate::charset::CharRestrict; -pub use crate::model_metadata::ModelMetadata; -use crate::utils::color_filter::{ColorPreset, HsvRange}; +pub use crate::algo::{SlideResult, Slider}; +pub use crate::models::det::{DetBuilder, DetSession, DetectionResult, Detector}; +pub use crate::models::ocr::{Ocr, OcrBuilder, OcrResult, OcrSession}; use models::loader::ModelSession; - +pub use models::ocr::model_metadata::ModelMetadata; // pub enum ModelSpec { // /// 默认 OCR (使用内置路径) @@ -99,7 +89,7 @@ use models::loader::ModelSession; // // impl Display for DdddOcr { // fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { -// write!(f, "DdddOcr(session: {})", self.runtime.desc()) +// write!(f, "DdddOcr(ocr: {})", self.runtime.desc()) // } // } // @@ -175,7 +165,7 @@ use models::loader::ModelSession; mod tests { #[test] fn test_ctc_decode_indices() { - // 模拟一个 DdddOcr 实例（如果 decode 不依赖 session，可以设为相关函数） + // 模拟一个 DdddOcr 实例（如果 decode 不依赖 ocr，可以设为相关函数） // 这里假设你的 decode_ctc 是公开或内部可访问的 let input = vec![1, 1, 0, 1, 2, 2, 0, 2]; // 逻辑：[1, 1] -> 1, [0] -> 跳过, [1] -> 1, [2, 2] -> 2, [0] -> 跳过, [2] -> 2 diff --git a/src/models/base.rs b/src/models/base.rs deleted file mode 100644 index e9c8b2c..0000000 --- a/src/models/base.rs +++ /dev/null @@ -1,40 +0,0 @@ -pub trait ModelArgs { - // 获取模型路径 - fn model_path(&self) -> &str; - - // 获取字符集（由于 Det 没有，所以返回 Option） - fn charset(&self) -> Option<&str>; -} - -pub struct HasCharset { - pub charset: String, -} // 给 Ocr 和 Custom 用 -pub struct NoCharset; // 给 Det 用 - -pub struct Model { - pub path: String, - pub metadata: T, -} -// 针对有字符集的模型 (Ocr / Custom) -impl ModelArgs for Model { - fn model_path(&self) -> &str { - &self.path - } - fn charset(&self) -> Option<&str> { - Some(&self.metadata.charset) - } -} - -// 针对没有字符集的模型 (Det) -impl ModelArgs for Model { - fn model_path(&self) -> &str { - &self.path - } - fn charset(&self) -> Option<&str> { - None - } -} - -pub type OcrModel = Model; -pub type DetModel = Model; -pub type CustomModel = Model; // Ocr 和 Custom 逻辑一致，可以复用 \ No newline at end of file diff --git a/src/models/det/builder.rs b/src/models/det/builder.rs new file mode 100644 index 0000000..7c4713f --- /dev/null +++ b/src/models/det/builder.rs @@ -0,0 +1,25 @@ +use crate::models::det::executor::Detector; +use crate::models::det::session::DetSession; + +pub struct DetBuilder { + use_gpu: bool, + device_id: u8, +} + +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: &DetSession) -> Detector<'_> { + Detector { + session, + use_gpu: self.use_gpu, + device_id: self.device_id, + } + } +} diff --git a/src/models/det.rs b/src/models/det/executor.rs similarity index 92% rename from src/models/det.rs rename to src/models/det/executor.rs index 56b5ede..6cbb8b9 100644 --- a/src/models/det.rs +++ b/src/models/det/executor.rs @@ -1,4 +1,4 @@ -use crate::model_metadata::ModelMetadata; +use crate::models::ocr::model_metadata::ModelMetadata; use crate::models::loader::{ModelLoader, ModelSession, ModelType}; use anyhow::{Context, Result, anyhow}; use image::{DynamicImage, GenericImageView, imageops::FilterType}; @@ -10,6 +10,7 @@ const DEFAULT_DET_PATH: &'static str = "common_det.onnx"; // 预设的提示信息常量 use crate::error::MODEL_DOWNLOAD_HELP; +use crate::models::det::session::DetSession; #[derive(Debug, Clone, Copy)] pub struct DetectionResult { @@ -21,30 +22,25 @@ pub struct DetectionResult { pub class_id: u32, } -#[derive(Debug)] -pub struct Detector { - session: RunnableModel, Graph>>, -} -impl ModelSession for Detector { - fn get_model_type(&self) -> ModelType { - todo!() - } - fn desc(&self) -> String { - "Detection Model 加载成功".to_string() - } -} -impl Detector { - pub fn new

(model_path: P) -> Result - where - P: AsRef, - { - let session = ModelLoader::model_for_path(&model_path)?.session; - Ok(Self { session }) - } - pub fn model_from_bytes(model_bytes: &[u8]) -> Result { - let session = ModelLoader::model_from_bytes(model_bytes)?.session; - Ok(Self { session }) + + + + +#[derive(Debug)] +pub struct Detector<'a> { + pub(crate) session: &'a DetSession, + pub(crate) use_gpu: bool, + pub(crate) device_id: u8, +} + +impl<'a> Detector<'a> { + pub fn new(session: &'a DetSession) -> Self { + Detector { + session, + use_gpu: false, + device_id: 0, + } } pub fn predict(&self, image: &DynamicImage) -> Result> { @@ -256,8 +252,10 @@ impl Detector { let (input_tensor, ratio) = self.preproc(dynamic_img, (416, 416))?; // tract 推理 - let outputs = self.session.run(tvec!(input_tensor.into()))?; - let output_array = outputs[0] + // let outputs = self.session.session.run(tvec!(input_tensor.into()))?; + let outputs = self.session.inference(input_tensor)?; + // let output_array = outputs[0] + let output_array = outputs .to_array_view::()? .to_owned() .into_dimensionality::()?; diff --git a/src/models/det/mod.rs b/src/models/det/mod.rs new file mode 100644 index 0000000..eb6a733 --- /dev/null +++ b/src/models/det/mod.rs @@ -0,0 +1,7 @@ +mod builder; +mod executor; +mod session; + +pub use builder::DetBuilder; +pub use executor::{DetectionResult, Detector}; +pub use session::DetSession; diff --git a/src/models/det/session.rs b/src/models/det/session.rs new file mode 100644 index 0000000..92a36ff --- /dev/null +++ b/src/models/det/session.rs @@ -0,0 +1,43 @@ +use crate::models::loader::{ModelLoader, ModelSession, ModelType}; +use anyhow::{Context, Result}; +use std::path::Path; +use tract_onnx::prelude::{tvec, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp}; + +#[derive(Debug)] +pub struct DetSession { + pub(crate) session: RunnableModel, Graph>>, +} + +impl ModelSession for DetSession { + fn get_model_type(&self) -> ModelType { + todo!() + } + fn desc(&self) -> String { + "Detection Model 加载成功".to_string() + } +} + +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 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()) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 981358c..c62f505 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,5 +1,3 @@ -pub mod base; pub mod loader; pub mod ocr; -pub mod det; -pub mod slide; \ No newline at end of file +pub mod det; \ No newline at end of file diff --git a/src/models/ocr/builder.rs b/src/models/ocr/builder.rs new file mode 100644 index 0000000..20b07ec --- /dev/null +++ b/src/models/ocr/builder.rs @@ -0,0 +1,72 @@ +use crate::models::ocr::charset::TokenFilter; +use crate::models::ocr::executor::Ocr; +use crate::models::ocr::session::OcrSession; +use crate::models::ocr::color_filter::ColorFilter; +pub struct OcrBuilder { + /// 是否修复PNG格式问题 + png_fix: bool, + /// 是否返回概率信息 + probability: bool, + /// 颜色过滤：保留的颜色列表 + color_filter: Option>, + /// 字符集范围 + charset_restrict: Option>, +} + +impl OcrBuilder { + // 初始化任务，设置默认参数 + pub fn new() -> Self { + Self { + png_fix: false, // 默认值 + probability: false, + color_filter: None, + charset_restrict: None, + } + } + pub fn png_fix(mut self, value: bool) -> Self { + self.png_fix = value; + self + } + pub fn probability(mut self, value: bool) -> Self { + self.probability = value; + self + } + + pub fn color_filter(mut self, filter: T) -> Self + where + T: ColorFilter + Send + Sync + 'static, + { + self.color_filter = Some(Box::new(filter)); + self + } + + pub fn charset_restrict(mut self, restrict: T) -> Self + where + T: TokenFilter + Send + Sync + 'static, + { + self.charset_restrict = Some(Box::new(restrict)); + self + } + pub fn build(self, session: &OcrSession) -> Ocr<'_> { + // 1. 原地解析颜色过滤器 + let final_color_ranges = match &self.color_filter { + Some(filter) => filter.collect_to_vec(), + None => Ok(None), + }; + // 2. 原地解析字符集过滤 + let tokens = &session.model_metadata.charset.tokens; + let final_charset_indices = match &self.charset_restrict { + Some(restrict) => restrict.apply_to_charset(tokens), + None => None, + }; + + // Ocr::new(session, self) + Ocr { + session, + png_fix: self.png_fix, // 原地解构出来 + probability: self.probability, + final_color_ranges, + final_charset_indices, + } + } +} diff --git a/src/charset.rs b/src/models/ocr/charset.rs similarity index 100% rename from src/charset.rs rename to src/models/ocr/charset.rs diff --git a/src/utils/color_filter.rs b/src/models/ocr/color_filter.rs similarity index 100% rename from src/utils/color_filter.rs rename to src/models/ocr/color_filter.rs diff --git a/src/models/ocr.rs b/src/models/ocr/executor.rs similarity index 80% rename from src/models/ocr.rs rename to src/models/ocr/executor.rs index 9060520..8ce4d4b 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr/executor.rs @@ -1,27 +1,16 @@ -use crate::charset::{TokenFilter, ValidationCtx}; -use crate::model_metadata::{ModelMetadata, Resize}; -use crate::models::base::ModelArgs; -use crate::models::loader::{ModelLoader, ModelSession, ModelType}; -use crate::utils::color_filter::{ColorFilter, HsvRange, filter_image}; +use crate::models::ocr::model_metadata::Resize; + +use crate::models::ocr::session::OcrSession; +use crate::models::ocr::color_filter::{HsvRange, filter_image}; use crate::utils::image_io::png_rgba_white_preprocess; use crate::utils::image_processor::{convert_to_grayscale, resize_image}; -use anyhow::Context; -use anyhow::{Result, anyhow}; -use image::{DynamicImage, ImageBuffer, Rgb}; +use anyhow::Result; +use image::DynamicImage; use serde::Serialize; use std::borrow::Cow; -use std::collections::HashSet; use std::fmt; -use std::path::Path; use tract_onnx::prelude::tract_ndarray::{ArrayView2, Ix2, s}; -use tract_onnx::prelude::{ - DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec, -}; - -// 引入 cv_ops 模块中的 OpenCV HSV 转换算子 -use crate::utils::cv_ops::rgb_to_opencv_hsv; - -/// 推理最终输出的强类型外壳（完全 Owned，无任何生命周期，可直接转 JSON） +use tract_onnx::prelude::{DatumType, Tensor, tract_ndarray}; #[derive(Debug, Clone, Serialize)] pub enum OcrResult { /// 纯文本分支（对应 probability = false） @@ -106,113 +95,40 @@ impl fmt::Display for OcrResult { } } -pub struct Ocr { - pub session: RunnableModel, Graph>>, - pub model_metadata: ModelMetadata, -} -impl ModelSession for Ocr { - fn get_model_type(&self) -> ModelType { - todo!("使用thiserror作为错误处理的库,thiserror 专门用于开发库（Library）"); - } - fn desc(&self) -> String { - "Ocr Model 加载成功".to_string() - } -} -impl Ocr { - pub fn new

(model_path: P, model_metadata: ModelMetadata) -> Result - where - P: AsRef, - { - let session = ModelLoader::model_for_path(model_path)?.session; - Ok(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, - }) - } - /// 对应 Python 的 _inference - fn inference(&self, tensor: Tensor) -> anyhow::Result { - // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 - // let result = self.session.run(tvec!(tensor.into()))?; - let mut result = self - .session - .run(tvec!(tensor.into())) - .context("执行模型推理失败")?; - println!("模型输出原始数据: {:?}", result); - Ok(result.swap_remove(0).into_tensor()) - } - - pub fn predictor(&'_ self) -> OcrPredictor<'_> { - OcrPredictor::new(self) - } -} - -pub struct OcrPredictor<'a> { - ocr: &'a Ocr, - /// 是否修复PNG格式问题 - png_fix: bool, - /// 是否返回概率信息 - probability: bool, +pub struct Ocr<'a> { + pub(crate) session: &'a OcrSession, + pub(crate) png_fix: bool, + pub(crate) probability: bool, /// 颜色过滤：保留的颜色列表 - color_filter: Result>, String>, + pub(crate) final_color_ranges: Result>, String>, /// 字符集范围 - charset_restrict: Option>, + pub(crate) final_charset_indices: Option>, } -impl<'a> OcrPredictor<'a> { +impl<'a> Ocr<'a> { // 初始化任务，设置默认参数 - pub fn new(ocr: &'a Ocr) -> Self { - Self { - ocr, + + pub fn new(session: &'a OcrSession) -> Self { + Ocr { + session, png_fix: false, // 默认值 probability: false, - color_filter: Ok(None), - charset_restrict: None, + final_color_ranges: Ok(None), + final_charset_indices: None, } } - pub fn png_fix(mut self, value: bool) -> Self { - self.png_fix = value; - self - } - pub fn probability(mut self, value: bool) -> Self { - self.probability = value; - self - } - - pub fn color_filter(mut self, filter: &dyn ColorFilter) -> Self { - // 一句话把活全包了！错误信息无缝传递，完美熔断 - match filter.collect_to_vec() { - Ok(new_ranges) => self.color_filter = Ok(new_ranges), - Err(err_msg) => self.color_filter = Err(err_msg), // 校验失败，Builder 正式中毒 - } - - self - } - - pub fn charset_restrict(mut self, restrict: &dyn TokenFilter) -> Self { - let charset = &self.ocr.model_metadata.charset; - let tokens = &charset.tokens; - self.charset_restrict = restrict.apply_to_charset(tokens); - self - } } -impl<'a> OcrPredictor<'a> { - pub fn predict(self, image: &DynamicImage) -> anyhow::Result { - println!("当前颜色过滤器状态: {:?}", self.color_filter); +impl<'a> Ocr<'a> { + pub fn predict(&self, image: &DynamicImage) -> anyhow::Result { + println!("当前颜色过滤器状态: {:?}", self.final_color_ranges); + // ===================================================================== // 管道节点 1: 颜色过滤流水线 // 使用 Cow (Copy-On-Write) 智能指针。 // 如果未开启过滤，img_cow 内部只是持有原图的【只读借用】，发生【零内存分配】！ // ===================================================================== - let img_cow = match &self.color_filter { + let img_cow = match &self.final_color_ranges { Err(err_msg) => { return Err(anyhow::anyhow!( "颜色过滤器初始化失败，全链路短路: {}", @@ -231,7 +147,7 @@ impl<'a> OcrPredictor<'a> { }; let tensor = self.preprocess_image(&img_cow)?; - let raw_tensor = self.ocr.inference(tensor)?; + let raw_tensor = self.session.inference(tensor)?; // 3. 后处理分流：直接返回 OcrResult let ocr_output = match raw_tensor.datum_type() { @@ -252,7 +168,7 @@ impl<'a> OcrPredictor<'a> { /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result { // 1. 获取模型元数据配置 - let meta = &self.ocr.model_metadata; + let meta = &self.session.model_metadata; let norm = &meta.normalization; // 获取归一化器 // A. 修复 PNG 透明背景 (内部逻辑你之前已实现) @@ -340,13 +256,13 @@ impl<'a> OcrPredictor<'a> { // Ok(tensor) } } -impl<'a> OcrPredictor<'a> { +impl<'a> Ocr<'a> { fn is_valid_indices(&self, idx: usize) -> bool { - if idx >= self.ocr.model_metadata.charset.size() { + if idx >= self.session.model_metadata.charset.size() { return false; } - match &self.charset_restrict { + match &self.final_charset_indices { Some(v) => v.binary_search(&idx).is_ok(), None => true, } @@ -354,9 +270,9 @@ impl<'a> OcrPredictor<'a> { /// 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 /// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！ pub fn valid_tokens(&self) -> Vec<&str> { - let charset = &self.ocr.model_metadata.charset; + let charset = &self.session.model_metadata.charset; let tokens = &charset.tokens; - match &self.charset_restrict { + match &self.final_charset_indices { Some(indices) => indices .iter() .filter_map(|&idx| tokens.get(idx).map(|cow| cow.as_ref())) @@ -366,9 +282,9 @@ impl<'a> OcrPredictor<'a> { } } pub fn valid_size(&self) -> usize { - match &self.charset_restrict { + match &self.final_charset_indices { Some(indices) => indices.len(), - None => self.ocr.model_metadata.charset.tokens.len(), + None => self.session.model_metadata.charset.tokens.len(), } } /// 变体 B 核心处理器：单次遍历 2D 视图，融合计算 Softmax、Argmax、置信度并输出概率大包 @@ -516,7 +432,7 @@ impl<'a> OcrPredictor<'a> { /// 获取有效字符索引列表 (用于外部验证或过滤) fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String { println!("indices模型输出原始数据: {:?}", predicted_indices); - let charset = &self.ocr.model_metadata.charset; + let charset = &self.session.model_metadata.charset; let tokens = &charset.tokens; // let valid_indices = &charset.valid_indices; @@ -544,7 +460,7 @@ impl<'a> OcrPredictor<'a> { // 史诗级加速点：如果是 None，说明没限制，根本不进入分支，直接放行！ // 只有当有具体限制（Some）时，才去跑 4-5 次 CPU 寄存器级别的二分查找 - if let Some(ref indices) = self.charset_restrict { + if let Some(ref indices) = self.final_charset_indices { if indices.binary_search(&u_idx).is_err() { continue; } diff --git a/src/models/ocr/mod.rs b/src/models/ocr/mod.rs new file mode 100644 index 0000000..ecf7fac --- /dev/null +++ b/src/models/ocr/mod.rs @@ -0,0 +1,10 @@ +mod builder; +mod executor; +mod session; +pub mod charset; +pub mod model_metadata; +pub mod color_filter; + +pub use builder::OcrBuilder; +pub use executor::{Ocr, OcrResult}; +pub use session::OcrSession; diff --git a/src/model_metadata.rs b/src/models/ocr/model_metadata.rs similarity index 99% rename from src/model_metadata.rs rename to src/models/ocr/model_metadata.rs index d62270b..63aef00 100644 --- a/src/model_metadata.rs +++ b/src/models/ocr/model_metadata.rs @@ -1,4 +1,4 @@ -use crate::charset::{CHARSET_BETA, CHARSET_OLD, Charset}; +use crate::models::ocr::charset::{CHARSET_BETA, CHARSET_OLD, Charset}; use anyhow::{Result, anyhow}; use serde::Deserialize; use std::borrow::Cow; diff --git a/src/models/ocr/session.rs b/src/models/ocr/session.rs new file mode 100644 index 0000000..4bd056d --- /dev/null +++ b/src/models/ocr/session.rs @@ -0,0 +1,53 @@ +use crate::models::ocr::model_metadata::ModelMetadata; +use crate::models::loader::{ModelLoader, ModelSession, ModelType}; +use anyhow::Context; +use anyhow::Result; +use std::path::Path; +use tract_onnx::prelude::{tvec, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp}; + +pub struct OcrSession { + pub session: RunnableModel, Graph>>, + pub model_metadata: ModelMetadata, +} +impl ModelSession for OcrSession { + fn get_model_type(&self) -> ModelType { + todo!("使用thiserror作为错误处理的库,thiserror 专门用于开发库（Library）"); + } + fn desc(&self) -> String { + "Ocr Model 加载成功".to_string() + } +} +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 { + 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, + }) + } + /// 对应 Python 的 _inference + 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()) + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 7e0f806..2f2430c 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,4 +1,3 @@ pub mod image_io; pub mod image_processor; pub mod cv_ops; -pub mod color_filter; \ No newline at end of file diff --git a/tests/ocr_test.rs b/tests/ocr_test.rs index 73a5946..467c0fc 100644 --- a/tests/ocr_test.rs +++ b/tests/ocr_test.rs @@ -1,4 +1,4 @@ -use ddddocr_rs::{Ocr, Slider, Detector, ModelMetadata}; // 假设你的包名是这个 +use ddddocr_rs::{Ocr, Slider, Detector, ModelMetadata, OcrSession, DetBuilder, DetSession}; // 假设你的包名是这个 use image::{DynamicImage, Rgb}; use std::fs; use std::path::Path; @@ -64,20 +64,20 @@ fn save_debug_image( #[test] fn test_full_classification() { // 1. 初始化模型 - let ocr = Ocr::new("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx",ModelMetadata::from_builtin_beta()).expect("模型加载失败"); + let ocr = OcrSession::new("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx",ModelMetadata::from_builtin_beta()).expect("模型加载失败"); // 2. 加载测试图片 let img = image::open("samples/code2.png").expect("测试图片不存在"); // 3. 执行识别 - let result = ocr.predictor().predict(&img).expect("识别过程出错").into_text(); + let result = Ocr::new(&ocr).predict(&img).expect("识别过程出错").into_text(); println!("识别结果: {}", result); assert!(!result.is_empty()); } #[test] fn test_det_load() -> anyhow::Result<()> { - let det = Detector::new("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx")?; + let det = DetSession::new("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx")?; let image_path = "samples/det1.png"; let image_bytes = fs::read(image_path).map_err(|e| anyhow::anyhow!("无法读取图片 {}: {}", image_path, e))?; @@ -89,8 +89,8 @@ fn test_det_load() -> anyhow::Result<()> { .map_err(|e| anyhow::anyhow!("图片解码失败: {}", e))?; // 【修改点 2】传入统一的 &DynamicImage 引用 - let bboxes = det.predict(&img)?; - println!("{:?}", det); + let bboxes = Detector::new(&det).predict(&img)?; + // println!("{:?}", det); println!("检测到的目标数量: {}", bboxes.len()); if bboxes.is_empty() {