diff --git a/Cargo.toml b/Cargo.toml index 55e6c95..9b910e5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,4 +11,8 @@ image = "0.25.10" base64 = "0.22.1" imageproc = { version = "0.26.2", default-features = true } serde = { version = "1.0.228", features = ["derive"] } -serde_json = "1.0.150" \ No newline at end of file +serde_json = "1.0.150" + +[features] +default = [] +embed-models = [] # 这是一个留给有特殊需求、且自己下载了模型放入 models/ 目录的人的后门 \ No newline at end of file diff --git a/examples/simple_usage.rs b/examples/simple_usage.rs index decfde6..61f2adf 100644 --- a/examples/simple_usage.rs +++ b/examples/simple_usage.rs @@ -1,5 +1,5 @@ fn main() { - let ocr = ddddocr_rs::DdddOcrBuilder::new().build().unwrap(); - let img = image::open("samples/code3.png").unwrap(); - println!("Result: {}", ocr.classification(&img).unwrap()); + // let ocr = ddddocr_rs::DdddOcrBuilder::new().build().unwrap(); + // let img = image::open("samples/code3.png").unwrap(); + // println!("Result: {}", ocr.classification(&img).unwrap()); } \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index fdc0c90..24f37c0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,138 +1,141 @@ mod charset; - +mod error; mod model_metadata; pub mod models; pub mod utils; -use anyhow::{Result, anyhow}; +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; -use crate::model_metadata::ModelMetadata; -use crate::models::det::DetectionResult; +pub use crate::model_metadata::ModelMetadata; use crate::utils::color_filter::{ColorPreset, HsvRange}; -use models::det::Det; use models::loader::ModelSession; -use models::ocr::Ocr; -pub enum ModelSpec { - /// 默认 OCR (使用内置路径) - OcrModel, - DetModel, - /// 自定义 OCR (路径由用户提供) - CustomOcrModel { - path: String, - model_metadata: ModelMetadata, - }, -} -impl ModelSpec { - // 将默认路径定义为内部关联常量 - const DEFAULT_OCR_PATH: &'static str = "models/common_sml2h3_f32.onnx"; - const DEFAULT_DET_PATH: &'static str = "models/common_det.onnx"; -} -pub enum Runtime { - Ocr(Ocr), - Det(Det), -} -impl Runtime { - // 统一获取描述的方法 - pub fn desc(&self) -> String { - match self { - Runtime::Ocr(s) => s.desc(), // 调用 Ocr 结构体的方法 - Runtime::Det(s) => s.desc(), // 调用 Det 结构体的方法 - } - } -} -pub struct DdddOcrBuilder { - mode: ModelSpec, -} -impl DdddOcrBuilder { - pub fn new() -> Self { - Self { - mode: ModelSpec::OcrModel, - } - } - - /// 切换为检测模式 - pub fn det(mut self) -> Self { - self.mode = ModelSpec::DetModel; - self - } - - /// 设置自定义 OCR 路径 - pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self { - // 直接重写枚举，替换掉之前的 Ocr 或 Det - self.mode = ModelSpec::CustomOcrModel { - path, - model_metadata, - }; - self - } - - /// 核心初始化逻辑 - pub fn build(self) -> Result { - let runtime = match self.mode { - ModelSpec::OcrModel => Runtime::Ocr(Ocr::new( - ModelSpec::DEFAULT_OCR_PATH.into(), - ModelMetadata::from_builtin_beta(), - )?), - ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?), - ModelSpec::CustomOcrModel { - path, - model_metadata, - } => Runtime::Ocr(Ocr::new(path, model_metadata)?), - }; - - Ok(DdddOcr { runtime }) - } -} - -pub struct DdddOcr { - runtime: Runtime, -} - -impl Display for DdddOcr { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "DdddOcr(session: {})", self.runtime.desc()) - } -} - -impl DdddOcr { - pub fn classification(&self, img: &DynamicImage) -> Result { - match &self.runtime { - // Runtime::Ocr(s) => s.predict(img).run(), - // Runtime::Ocr(s) => s.predictor().probability(false).predict(img), - // Runtime::Ocr(s) => { - // let predictor = s.predictor(); - // let restricted = predictor.charset_restrict(&CharRestrict::Lowercase); - // let a = restricted.valid_tokens(); - // println!("{:?}", a); - // Ok("".to_string()) - // } - Runtime::Ocr(s) => { - let res = s.predictor().probability(true).predict(img)?; - println!("{}", res); - Ok(res.to_string()) - } - // Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img), - // Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![ - // // 错误：下界 (82, 221, 14) 没问题 - // // 但上界的 H 通道写成了 240，超过了 180 的法定上限！ - // HsvRange::new((82, 221, 14), (240, 203, 82)), - // ])).predict(img), - Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), - } - } - pub fn detection(&self, img: &DynamicImage) -> Result> { - match &self.runtime { - Runtime::Det(s) => s.predict(img), - Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")), - } - } -} +// pub enum ModelSpec { +// /// 默认 OCR (使用内置路径) +// OcrModel, +// DetModel, +// /// 自定义 OCR (路径由用户提供) +// CustomOcrModel { +// path: String, +// model_metadata: ModelMetadata, +// }, +// } +// impl ModelSpec { +// // 将默认路径定义为内部关联常量 +// const DEFAULT_OCR_PATH: &'static str = "models/common_sml2h3_f32.onnx"; +// const DEFAULT_DET_PATH: &'static str = "models/common_det.onnx"; +// } +// pub enum Runtime { +// Ocr(Ocr), +// Det(Det), +// } +// impl Runtime { +// // 统一获取描述的方法 +// pub fn desc(&self) -> String { +// match self { +// Runtime::Ocr(s) => s.desc(), // 调用 Ocr 结构体的方法 +// Runtime::Det(s) => s.desc(), // 调用 Det 结构体的方法 +// } +// } +// } +// pub struct DdddOcrBuilder { +// mode: ModelSpec, +// } +// +// impl DdddOcrBuilder { +// pub fn new() -> Self { +// Self { +// mode: ModelSpec::OcrModel, +// } +// } +// +// /// 切换为检测模式 +// pub fn det(mut self) -> Self { +// self.mode = ModelSpec::DetModel; +// self +// } +// +// /// 设置自定义 OCR 路径 +// pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self { +// // 直接重写枚举，替换掉之前的 Ocr 或 Det +// self.mode = ModelSpec::CustomOcrModel { +// path, +// model_metadata, +// }; +// self +// } +// +// /// 核心初始化逻辑 +// pub fn build(self) -> Result { +// let runtime = match self.mode { +// ModelSpec::OcrModel => Runtime::Ocr(Ocr::new( +// ModelSpec::DEFAULT_OCR_PATH.into(), +// ModelMetadata::from_builtin_beta(), +// )?), +// ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?), +// ModelSpec::CustomOcrModel { +// path, +// model_metadata, +// } => Runtime::Ocr(Ocr::new(path, model_metadata)?), +// }; +// +// Ok(DdddOcr { runtime }) +// } +// } +// +// pub struct DdddOcr { +// runtime: Runtime, +// } +// +// impl Display for DdddOcr { +// fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { +// write!(f, "DdddOcr(session: {})", self.runtime.desc()) +// } +// } +// +// impl DdddOcr { +// pub fn classification(&self, img: &DynamicImage) -> Result { +// match &self.runtime { +// // Runtime::Ocr(s) => s.predict(img).run(), +// // Runtime::Ocr(s) => s.predictor().probability(false).predict(img), +// // Runtime::Ocr(s) => { +// // let predictor = s.predictor(); +// // let restricted = predictor.charset_restrict(&CharRestrict::Lowercase); +// // let a = restricted.valid_tokens(); +// // println!("{:?}", a); +// // Ok("".to_string()) +// // } +// Runtime::Ocr(s) => { +// let res = s.predictor().probability(true).predict(img)?; +// println!("{}", res); +// Ok(res.to_string()) +// } +// // Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img), +// // Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![ +// // // 错误：下界 (82, 221, 14) 没问题 +// // // 但上界的 H 通道写成了 240，超过了 180 的法定上限！ +// // HsvRange::new((82, 221, 14), (240, 203, 82)), +// // ])).predict(img), +// Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), +// } +// } +// pub fn detection(&self, img: &DynamicImage) -> Result> { +// match &self.runtime { +// Runtime::Det(s) => s.predict(img), +// Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")), +// } +// } +// } // struct Classification {} // #[derive(Debug)] diff --git a/src/model_metadata.rs b/src/model_metadata.rs index 1fd5efa..d62270b 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -113,8 +113,53 @@ 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))?; - /// 从外部外部 JSON 文件动态加载字符集 + // 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) + } + + /// 从外部外部 JSON 文件动态加载字符集(在后续优化中移除) pub fn from_json_file>(path: P) -> Result { let path = path.as_ref(); if !path.exists() { diff --git a/src/models/det.rs b/src/models/det.rs index d4413d0..56b5ede 100644 --- a/src/models/det.rs +++ b/src/models/det.rs @@ -1,10 +1,15 @@ +use crate::model_metadata::ModelMetadata; use crate::models::loader::{ModelLoader, ModelSession, ModelType}; -use anyhow::{Context, Result}; +use anyhow::{Context, Result, anyhow}; use image::{DynamicImage, GenericImageView, imageops::FilterType}; +use std::path::Path; use tract_onnx::prelude::tract_ndarray::{Array2, Array3, Array4, Axis, prelude::*, s}; use tract_onnx::prelude::{Graph, RunnableModel, Tensor, TypedFact, TypedOp, tvec}; +const DEFAULT_DET_PATH: &'static str = "common_det.onnx"; +// 预设的提示信息常量 +use crate::error::MODEL_DOWNLOAD_HELP; #[derive(Debug, Clone, Copy)] pub struct DetectionResult { @@ -16,12 +21,11 @@ pub struct DetectionResult { pub class_id: u32, } - - -pub struct Det { +#[derive(Debug)] +pub struct Detector { session: RunnableModel, Graph>>, } -impl ModelSession for Det { +impl ModelSession for Detector { fn get_model_type(&self) -> ModelType { todo!() } @@ -29,11 +33,20 @@ impl ModelSession for Det { "Detection Model 加载成功".to_string() } } -impl Det { - pub fn new(model_path: String) -> Result { - let session = ModelLoader::load_model(&model_path)?.session; +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 }) + } + pub fn predict(&self, image: &DynamicImage) -> Result> { // Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑 self.get_bbox(image) @@ -73,11 +86,10 @@ impl Det { // BGR 赋值 array[[0, 0, y, x]] = slice[idx + 2] as f32; // B array[[0, 1, y, x]] = slice[idx + 1] as f32; // G - array[[0, 2, y, x]] = slice[idx] as f32; // R + array[[0, 2, y, x]] = slice[idx] as f32; // R } } - Ok((array.into(), r)) } @@ -273,17 +285,15 @@ impl Det { let detections = self.multiclass_nms(&boxes_xyxy, &scores, 0.45, 0.1); let final_results = detections .into_iter() - .map(|d| { - DetectionResult{ - x1: (d[0] as i32).max(0).min(orig_w as i32), - y1: (d[1] as i32).max(0).min(orig_h as i32), - x2: (d[2] as i32).max(0).min(orig_w as i32), - y2: (d[3] as i32).max(0).min(orig_h as i32), - score: d[4], - class_id: d[5] as u32, - } + .map(|d| DetectionResult { + x1: (d[0] as i32).max(0).min(orig_w as i32), + y1: (d[1] as i32).max(0).min(orig_h as i32), + x2: (d[2] as i32).max(0).min(orig_w as i32), + y2: (d[3] as i32).max(0).min(orig_h as i32), + score: d[4], + class_id: d[5] as u32, }) .collect(); - Ok(final_results ) + Ok(final_results) } } diff --git a/src/models/loader.rs b/src/models/loader.rs index ae1f1a8..32eedd9 100644 --- a/src/models/loader.rs +++ b/src/models/loader.rs @@ -1,4 +1,4 @@ -use anyhow::Context; +use anyhow::{anyhow, Context}; use image::DynamicImage; use tract_onnx::onnx; use tract_onnx::prelude::*; @@ -6,10 +6,12 @@ use tract_onnx::prelude::*; use crate::utils::image_io::png_rgba_white_preprocess; use crate::utils::image_processor::{convert_to_grayscale, resize_image}; use std::collections::HashMap; +use std::io::Cursor; use tract_onnx::prelude::tract_ndarray::s; +use crate::ModelMetadata; /// OCR 模型：包含路径和字符集 - +const DEFAULT_OCR_PATH: &'static str = "common_sml2h3_f32.onnx"; pub enum ModelType { Ocr, Det, @@ -26,7 +28,7 @@ pub struct ModelLoader { } impl ModelLoader { - pub fn load_model

(model_path: P) -> anyhow::Result + pub fn model_for_path

(model_path: P) -> anyhow::Result where P: AsRef, { @@ -37,4 +39,74 @@ impl ModelLoader { .into_runnable()?; Ok(Self { session }) } + /// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用） + pub fn model_from_bytes(model_bytes: &[u8]) -> anyhow::Result { + // 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read） + let mut cursor = Cursor::new(model_bytes); + + let session = onnx() + .model_for_read(&mut cursor) + .with_context(|| "从内存字节流解析 ONNX 模型失败")? + .into_optimized()? + .into_runnable()?; + + Ok(Self { session }) + } } +// impl ModelLoader { +// pub fn find_model_path(env_var: &str, default_filename: &str) -> Option { +// // 1. 策略一：优先尝试读取环境变量 +// if let Ok(env_path) = std::env::var(env_var) { +// let path = std::path::PathBuf::from(env_path); +// if path.exists() { +// return Some(path); +// } +// } +// // 2. 策略二：尝试在当前工作目录寻找 +// if let Ok(mut path) = std::env::current_dir() { +// path.push(default_filename); +// if path.exists() { +// return Some(path); +// } +// } +// +// // 3. 策略三：尝试在当前可执行文件同级目录寻找 +// if let Ok(mut exe_path) = std::env::current_exe() { +// exe_path.pop(); // 弹出可执行文件名，拿到所在的父目录 +// exe_path.push(default_filename); +// if exe_path.exists() { +// return Some(exe_path); +// } +// } +// // 4. 所有本地探测策略均落空 +// None +// } +// } + + + +// pub fn new_default() -> anyhow::Result { +// let metadata = ModelMetadata::from_builtin_beta(); // 绑定自带的 BETA 字符集 +// +// // 1. 策略一：优先尝试读取环境变量 +// if let Some(path) = ModelLoader::find_model_path("DDDD_OCR_MODEL", DEFAULT_OCR_PATH) { +// return Self::new(path, metadata); +// } +// +// // 4. 策略四：开启了 embed-models 特征时的编译期死穴保底 +// // 如果开启了 feature 但根目录下没有该模型，编译时会在此处直接中断失败 +// +// #[cfg(feature = "embed-models")] +// { +// let model_bytes = include_bytes!("../models/common_sml2h3_f32.onnx"); +// // 注意：这里需要你的 InternalOcr 扩展一个 from_bytes 的方法 +// return Self::model_from_bytes(model_bytes, metadata); +// +// } +// +// // 5. 所有策略落空，抛出保姆级错误 +// #[cfg(not(feature = "embed-models"))] +// { +// Err(anyhow!(MODEL_DOWNLOAD_HELP)) +// } +// } \ No newline at end of file diff --git a/src/models/ocr.rs b/src/models/ocr.rs index fac0338..9060520 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -12,10 +12,12 @@ 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; @@ -117,9 +119,19 @@ impl ModelSession for Ocr { } } impl Ocr { - pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result { - - let session = ModelLoader::load_model(&model_path)?.session; + 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, diff --git a/src/models/slide.rs b/src/models/slide.rs index ab2b702..3a7e096 100644 --- a/src/models/slide.rs +++ b/src/models/slide.rs @@ -20,13 +20,12 @@ pub struct SlideResult { pub confidence: f64, } -pub struct Slide; +pub struct Slider; -impl Slide { - pub fn new() -> Self { - Self +impl Slider { + pub fn new() -> Result { + Ok(Self) } - /// 对应 Python: slide_match 滑块匹配接口 pub fn slide_match( &self, diff --git a/tests/ocr_test.rs b/tests/ocr_test.rs index 4ae2892..73a5946 100644 --- a/tests/ocr_test.rs +++ b/tests/ocr_test.rs @@ -1,5 +1,4 @@ -use ddddocr_rs::models::slide::Slide; -use ddddocr_rs::{DdddOcr, DdddOcrBuilder}; // 假设你的包名是这个 +use ddddocr_rs::{Ocr, Slider, Detector, ModelMetadata}; // 假设你的包名是这个 use image::{DynamicImage, Rgb}; use std::fs; use std::path::Path; @@ -60,23 +59,25 @@ fn save_debug_image( img.save(output_path)?; Ok(()) } + + #[test] fn test_full_classification() { // 1. 初始化模型 - let ocr = DdddOcrBuilder::new().build().expect("模型加载失败"); + let ocr = Ocr::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.classification(&img).expect("识别过程出错"); + let result = ocr.predictor().predict(&img).expect("识别过程出错").into_text(); println!("识别结果: {}", result); assert!(!result.is_empty()); } #[test] fn test_det_load() -> anyhow::Result<()> { - let det = DdddOcrBuilder::new().det().build()?; + let det = Detector::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))?; @@ -88,8 +89,8 @@ fn test_det_load() -> anyhow::Result<()> { .map_err(|e| anyhow::anyhow!("图片解码失败: {}", e))?; // 【修改点 2】传入统一的 &DynamicImage 引用 - let bboxes = det.detection(&img)?; - println!(":?{}", det); + let bboxes = det.predict(&img)?; + println!("{:?}", det); println!("检测到的目标数量: {}", bboxes.len()); if bboxes.is_empty() { @@ -111,7 +112,7 @@ fn test_det_load() -> anyhow::Result<()> { #[test] fn test_real_slide_match() { - let engine = Slide::new(); + let engine = Slider::new().unwrap(); // 1. 加载你准备好的测试图 // 假设图片放在项目根目录下的 assets 文件夹 @@ -142,7 +143,7 @@ fn test_real_slide_match() { #[test] fn test_real_slide_comparison() { - let engine = Slide::new(); + let engine = Slider::new().unwrap(); // 1. 加载你准备好的测试图 // 假设图片放在项目根目录下的 assets 文件夹