CNWei/ddddocr-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(slide,det): 优化项目结构，移除不必要的逻辑

Browse Source 
- 优化 项目结构，移除不必要的逻辑
This commit is contained in:
CNWei
2026-07-07 09:55:00 +08:00
parent 7f1ce04f50
commit 31271e80db
9 changed files with 312 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
Cargo.toml
View File

@@ -12,3 +12,7 @@ base64 = "0.22.1"
imageproc = { version = "0.26.2", default-features = true }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150"
[features]
default = []
embed-models = [] # 这是一个留给有特殊需求、且自己下载了模型放入 models/ 目录的人的后门

6
examples/simple_usage.rs
View File

@@ -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());
}

245
src/lib.rs
View File

@@ -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<DdddOcr> {
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<String> {
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<Vec<DetectionResult>> {
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<DdddOcr> {
// 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<String> {
// 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<Vec<DetectionResult>> {
// match &self.runtime {
// Runtime::Det(s) => s.predict(img),
// Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")),
// }
// }
// }
// struct Classification {}
// #[derive(Debug)]

47
src/model_metadata.rs
View File

@@ -113,8 +113,53 @@ impl ModelMetadata {
normalization,
}
}
pub fn from_json_str(json_str: &str) -> Result<Self> {
let dto: ModelMetadataDto = serde_json::from_str(json_str)
.map_err(|e| anyhow!("JSON 反序列化失败，请检查字段是否完整: {}", e))?;
/// 从外部外部 JSON 文件动态加载字符集
// 1. 将 DTO 的字符串数组转化为强类型的 Charset
let tokens: Vec<Cow<'static, str>> =
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<Self> {
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<P: AsRef<Path>>(path: P) -> Result<Self> {
let path = path.as_ref();
if !path.exists() {

36
src/models/det.rs
View File

@@ -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<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
}
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<Self, anyhow::Error> {
let session = ModelLoader::load_model(&model_path)?.session;
impl Detector {
pub fn new<P>(model_path: P) -> Result<Self, anyhow::Error>
where
P: AsRef<Path>,
{
let session = ModelLoader::model_for_path(&model_path)?.session;
Ok(Self { session })
}
pub fn model_from_bytes(model_bytes: &[u8]) -> Result<Self, anyhow::Error> {
let session = ModelLoader::model_from_bytes(model_bytes)?.session;
Ok(Self { session })
}
pub fn predict(&self, image: &DynamicImage) -> Result<Vec<DetectionResult>> {
// Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑
self.get_bbox(image)
@@ -77,7 +90,6 @@ impl Det {
}
}
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{
.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)
}
}

78
src/models/loader.rs
View File

@@ -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<P>(model_path: P) -> anyhow::Result<Self>
pub fn model_for_path<P>(model_path: P) -> anyhow::Result<Self>
where
P: AsRef<std::path::Path>,
{
@@ -37,4 +39,74 @@ impl ModelLoader {
.into_runnable()?;
Ok(Self { session })
}
/// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用）
pub fn model_from_bytes(model_bytes: &[u8]) -> anyhow::Result<Self> {
// 使用 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<std::path::PathBuf> {
// // 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<Self> {
// 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))
// }
// }

16
src/models/ocr.rs
View File

@@ -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<Self, anyhow::Error> {
pub fn new<P>(model_path: P, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error>
where
P: AsRef<Path>,
{
let session = ModelLoader::model_for_path(model_path)?.session;
Ok(Self {
session,
model_metadata,
})
}
let session = ModelLoader::load_model(&model_path)?.session;
pub fn model_from_bytes(model_bytes: &[u8], model_metadata: ModelMetadata)->Result<Self, anyhow::Error>{
let session = ModelLoader::model_from_bytes(model_bytes)?.session;
Ok(Self {
session,
model_metadata,

9
src/models/slide.rs
View File

@@ -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<Self, anyhow::Error> {
Ok(Self)
}
/// 对应 Python: slide_match 滑块匹配接口
pub fn slide_match(
&self,

19
tests/ocr_test.rs
View File

@@ -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 文件夹